jps bazel compiler - simplify and improve "mark dirty" logic in kotlin jps builder in the light of what Baze and new storage provide

GitOrigin-RevId: 9d2fd17a1775e1c7ebf2bf88e2a5a8ee60889960

build/jvm-rules/src/jps-builder/JpsBuilder.kt

@@ -10,20 +10,48 @@ import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.trace.Span
 import io.opentelemetry.api.trace.Tracer
 import it.unimi.dsi.fastutil.objects.Object2ObjectArrayMap
-import kotlinx.coroutines.*
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.ensureActive
+import kotlinx.coroutines.job
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import org.apache.arrow.memory.RootAllocator
 import org.jetbrains.annotations.VisibleForTesting
-import org.jetbrains.bazel.jvm.*
+import org.jetbrains.bazel.jvm.ArgMap
+import org.jetbrains.bazel.jvm.WorkRequestExecutor
 import org.jetbrains.bazel.jvm.abi.JarContentToProcess
 import org.jetbrains.bazel.jvm.abi.writeAbi
-import org.jetbrains.bazel.jvm.jps.impl.*
+import org.jetbrains.bazel.jvm.emptyMap
+import org.jetbrains.bazel.jvm.hashMap
+import org.jetbrains.bazel.jvm.jps.impl.BazelBuildDataProvider
+import org.jetbrains.bazel.jvm.jps.impl.BazelBuildRootIndex
+import org.jetbrains.bazel.jvm.jps.impl.BazelBuildTargetIndex
+import org.jetbrains.bazel.jvm.jps.impl.BazelCompileContext
+import org.jetbrains.bazel.jvm.jps.impl.BazelCompileScope
+import org.jetbrains.bazel.jvm.jps.impl.BazelModuleBuildTarget
+import org.jetbrains.bazel.jvm.jps.impl.JpsTargetBuilder
+import org.jetbrains.bazel.jvm.jps.impl.NoopIgnoredFileIndex
+import org.jetbrains.bazel.jvm.jps.impl.NoopModuleExcludeIndex
+import org.jetbrains.bazel.jvm.jps.impl.RequestLog
+import org.jetbrains.bazel.jvm.jps.impl.createPathRelativizer
 import org.jetbrains.bazel.jvm.jps.java.BazelJavaBuilder
 import org.jetbrains.bazel.jvm.jps.kotlin.IncrementalKotlinBuilder
 import org.jetbrains.bazel.jvm.jps.kotlin.NonIncrementalKotlinBuilder
-import org.jetbrains.bazel.jvm.jps.state.*
+import org.jetbrains.bazel.jvm.jps.state.LoadStateResult
+import org.jetbrains.bazel.jvm.jps.state.TargetConfigurationDigestContainer
+import org.jetbrains.bazel.jvm.jps.state.TargetConfigurationDigestProperty
+import org.jetbrains.bazel.jvm.jps.state.createInitialSourceMap
+import org.jetbrains.bazel.jvm.jps.state.loadBuildState
+import org.jetbrains.bazel.jvm.jps.state.saveBuildState
 import org.jetbrains.bazel.jvm.kotlin.JvmBuilderFlags
 import org.jetbrains.bazel.jvm.kotlin.parseArgs
+import org.jetbrains.bazel.jvm.processRequests
+import org.jetbrains.bazel.jvm.span
+import org.jetbrains.bazel.jvm.use
 import org.jetbrains.jps.api.GlobalOptions
 import org.jetbrains.jps.backwardRefs.JavaBackwardReferenceIndexBuilder
 import org.jetbrains.jps.builders.logging.BuildLoggingManager
@@ -485,7 +513,7 @@ private suspend fun initAndBuild(
           .use { span ->
             val builders = arrayOf(
               if (compileScope.isIncrementalCompilation) {
-                IncrementalKotlinBuilder(isRebuild = isRebuild, span = span, dataManager = buildDataProvider)
+                IncrementalKotlinBuilder(isRebuild = isRebuild, span = span, dataManager = buildDataProvider, jpsTarget = moduleTarget)
               }
               else {
                 NonIncrementalKotlinBuilder(job = coroutineContext.job, span = span)

build/jvm-rules/src/jps-builder/impl/BazelKotlinFsOperationsHelper.kt

@@ -7,18 +7,12 @@ import com.intellij.openapi.util.io.FileUtilRt
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.Span
-import org.jetbrains.bazel.jvm.hashMap
+import org.jetbrains.bazel.jvm.jps.OutputSink
 import org.jetbrains.bazel.jvm.linkedSet
 import org.jetbrains.jps.ModuleChunk
-import org.jetbrains.jps.builders.FileProcessor
-import org.jetbrains.jps.builders.impl.DirtyFilesHolderBase
-import org.jetbrains.jps.builders.java.JavaSourceRootDescriptor
-import org.jetbrains.jps.incremental.BuildOperations
 import org.jetbrains.jps.incremental.CompileContext
 import org.jetbrains.jps.incremental.FSOperations.addCompletelyMarkedDirtyTarget
-import org.jetbrains.jps.incremental.ModuleBuildTarget
 import org.jetbrains.jps.incremental.fs.CompilationRound
-import org.jetbrains.kotlin.jps.build.KotlinDirtySourceFilesHolder
 import org.jetbrains.kotlin.jps.build.KotlinDirtySourceFilesHolder.TargetFiles
 import java.io.File
 import java.nio.file.Files
@@ -60,55 +54,40 @@ internal class BazelKotlinFsOperationsHelper(
     }
   }
 
-  internal fun markFilesForCurrentRound(files: Sequence<Path>, dirtyFilesHolder: KotlinDirtySourceFilesHolder) {
-    val buildRootIndex = context.projectDescriptor.buildRootIndex as BazelBuildRootIndex
-    for (file in files) {
-      val root = buildRootIndex.fileToDescriptors.get(file)
-      if (root != null) {
-        dirtyFilesHolder.byTarget.get(root.target)?._markDirty(file, root)
-      }
-    }
-
-    markFilesImpl(files, currentRound = true, span = span) { it.exists() }
-  }
-
-  /**
-   * Marks given [files] as dirty for current round and given [target] of [chunk].
-   */
-  fun markFilesForCurrentRound(target: ModuleBuildTarget, files: Collection<Path>, targetDirtyFiles: TargetFiles?) {
-    require(target in chunk.targets)
-
-    val dirtyFileToRoot = hashMap<Path, JavaSourceRootDescriptor>()
+  fun markFilesForCurrentRound(files: Sequence<Path>, targetDirtyFiles: TargetFiles?) {
     val buildRootIndex = context.projectDescriptor.buildRootIndex as BazelBuildRootIndex
     for (file in files) {
       val root = buildRootIndex.fileToDescriptors.get(file) ?: continue
-      targetDirtyFiles._markDirty(file, root)
-      dirtyFileToRoot.put(file, root)
+      targetDirtyFiles?._markDirty(file, root)
+    }
+
+    markFilesImpl(files = files, currentRound = true, span = span) { it.exists() }
+  }
+
+  /**
+   * Marks given [files] as dirty for current round.
+   */
+  fun markFilesForCurrentRound(
+    files: Collection<Path>,
+    targetDirtyFiles: TargetFiles?,
+    outputSink: OutputSink,
+    parentSpan: Span,
+  ) {
+    val buildRootIndex = context.projectDescriptor.buildRootIndex as BazelBuildRootIndex
+    for (file in files) {
+      targetDirtyFiles._markDirty(file, buildRootIndex.fileToDescriptors.get(file) ?: continue)
     }
 
     markFilesImpl(files.asSequence(), currentRound = true, span = span) { Files.exists(it) }
-    cleanOutputsForNewDirtyFilesInCurrentRound(target, dirtyFileToRoot)
-  }
-
-  private fun cleanOutputsForNewDirtyFilesInCurrentRound(target: ModuleBuildTarget, dirtyFiles: Map<Path, JavaSourceRootDescriptor>) {
-    val dirtyFilesHolder = object : DirtyFilesHolderBase<JavaSourceRootDescriptor, ModuleBuildTarget>(context) {
-      override fun processDirtyFiles(processor: FileProcessor<JavaSourceRootDescriptor, ModuleBuildTarget>) {
-        for ((file, root) in dirtyFiles) {
-          processor.apply(target, file.toFile(), root)
-        }
-      }
-
-      override fun hasDirtyFiles(): Boolean = dirtyFiles.isNotEmpty()
-    }
-    BuildOperations.cleanOutputsCorrespondingToChangedFiles(context, dirtyFilesHolder)
+    cleanOutputsCorrespondingToChangedFiles(files = files, dataManager = dataManager, outputSink = outputSink, parentSpan = parentSpan)
   }
 
   fun markFiles(files: Sequence<Path>) {
-    markFilesImpl(files, currentRound = false, span = span) { it.exists() }
+    markFilesImpl(files = files, currentRound = false, span = span) { it.exists() }
   }
 
   fun markInChunkOrDependents(files: Sequence<Path>, excludeFiles: Set<Path>) {
-    markFilesImpl(files, currentRound = false, span = span) {
+    markFilesImpl(files = files, currentRound = false, span = span) {
       !excludeFiles.contains(it) && it.exists()
     }
   }
@@ -133,18 +112,17 @@ internal class BazelKotlinFsOperationsHelper(
     }
 
     val projectDescriptor = context.projectDescriptor
+    val fileToDescriptors = (projectDescriptor.buildRootIndex as BazelBuildRootIndex).fileToDescriptors
     for (fileToMark in filesToMark) {
-      val rootDescriptor = (projectDescriptor.buildRootIndex as BazelBuildRootIndex).fileToDescriptors.get(fileToMark)
-      if (rootDescriptor != null) {
-        projectDescriptor.fsState.markDirty(
-          context,
-          compilationRound,
-          fileToMark,
-          rootDescriptor,
-          projectDescriptor.dataManager.getFileStampStorage(rootDescriptor.target),
-          false
-        )
-      }
+      val rootDescriptor = fileToDescriptors.get(fileToMark) ?: continue
+      projectDescriptor.fsState.markDirty(
+        /* context = */ context,
+        /* round = */ compilationRound,
+        /* file = */ fileToMark,
+        /* buildRootDescriptor = */ rootDescriptor,
+        /* stampStorage = */ projectDescriptor.dataManager.getFileStampStorage(rootDescriptor.target),
+        /* saveEventStamp = */ false,
+      )
     }
     span.addEvent("mark dirty", Attributes.of(
       AttributeKey.stringArrayKey("filesToMark"), filesToMark.map { it.toString() },

build/jvm-rules/src/jps-builder/impl/fsState.kt

@@ -71,6 +71,30 @@ internal fun cleanOutputsCorrespondingToChangedFiles(
   }
 }
 
+internal fun cleanOutputsCorrespondingToChangedFiles(
+  files: Collection<Path>,
+  dataManager: BazelBuildDataProvider,
+  outputSink: OutputSink,
+  parentSpan: Span,
+) {
+  val deletedOutputFiles = ArrayList<String>()
+  val sourceToOutputMapping = dataManager.sourceToOutputMapping
+  for (sourceFile in files) {
+    val outputs = sourceToOutputMapping.getAndClearOutputs(sourceFile)?.takeIf { it.isNotEmpty() } ?: continue
+    outputSink.removeAll(outputs)
+    deletedOutputFiles.addAll(outputs)
+  }
+
+  if (!deletedOutputFiles.isEmpty()) {
+    if (parentSpan.isRecording) {
+      parentSpan.addEvent("deletedOutputs", Attributes.of(
+        AttributeKey.stringArrayKey("deletedOutputFiles"), deletedOutputFiles,
+      ))
+    }
+    //context.processMessage(FileDeletedEvent(deletedOutputFiles.map { it.toString() }))
+  }
+}
+
 internal suspend fun markTargetUpToDate(
   context: CompileContext,
   target: ModuleBuildTarget,

build/jvm-rules/src/jps-builder/kotlin/IncrementalKotlinBuilder.kt

@@ -7,6 +7,10 @@ import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.Span
 import kotlinx.coroutines.ensureActive
+import org.jetbrains.bazel.jvm.concat
+import org.jetbrains.bazel.jvm.emptyList
+import org.jetbrains.bazel.jvm.emptySet
+import org.jetbrains.bazel.jvm.hashSet
 import org.jetbrains.bazel.jvm.jps.BazelConfigurationHolder
 import org.jetbrains.bazel.jvm.jps.OutputSink
 import org.jetbrains.bazel.jvm.jps.impl.BazelBuildDataProvider
@@ -30,7 +34,6 @@ import org.jetbrains.jps.incremental.BuilderCategory
 import org.jetbrains.jps.incremental.CompileContext
 import org.jetbrains.jps.incremental.ModuleBuildTarget
 import org.jetbrains.jps.incremental.ModuleLevelBuilder
-import org.jetbrains.jps.incremental.ModuleLevelBuilder.OutputConsumer
 import org.jetbrains.jps.incremental.RebuildRequestedException
 import org.jetbrains.jps.incremental.Utils
 import org.jetbrains.jps.model.module.JpsModule
@@ -43,7 +46,6 @@ import org.jetbrains.kotlin.build.report.ICReporter.ReportSeverity
 import org.jetbrains.kotlin.build.report.ICReporterBase
 import org.jetbrains.kotlin.cli.common.ExitCode
 import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
-import org.jetbrains.kotlin.compilerRunner.CompilerRunnerUtil
 import org.jetbrains.kotlin.compilerRunner.DummyKotlinPaths
 import org.jetbrains.kotlin.compilerRunner.JpsCompilerEnvironment
 import org.jetbrains.kotlin.compilerRunner.OutputItemsCollectorImpl
@@ -66,6 +68,7 @@ import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.incremental.getChangedAndImpactedSymbols
 import org.jetbrains.kotlin.incremental.mapClassesFqNamesToFiles
 import org.jetbrains.kotlin.incremental.mapLookupSymbolsToFiles
+import org.jetbrains.kotlin.incremental.parsing.classesFqNames
 import org.jetbrains.kotlin.jps.build.KotlinChunk
 import org.jetbrains.kotlin.jps.build.KotlinCompileContext
 import org.jetbrains.kotlin.jps.build.KotlinDirtySourceFilesHolder
@@ -98,6 +101,7 @@ private val classesToLoadByParent = ClassCondition { className ->
 
 internal class IncrementalKotlinBuilder(
   private val dataManager: BazelBuildDataProvider,
+  private val jpsTarget: BazelModuleBuildTarget,
   private val isRebuild: Boolean,
   private val span: Span,
 ) : BazelTargetBuilder(BuilderCategory.SOURCE_PROCESSOR) {
@@ -127,7 +131,7 @@ internal class IncrementalKotlinBuilder(
       }
 
       if (kotlinChunk.isEnabled) {
-        markAdditionalFilesForInitialRound(kotlinChunk, chunk, kotlinContext)
+        markAdditionalFilesForInitialRound(kotlinChunk, chunk, kotlinContext, jpsTarget)
       }
     }
   }
@@ -135,31 +139,10 @@ internal class IncrementalKotlinBuilder(
   private fun markAdditionalFilesForInitialRound(
     kotlinChunk: KotlinChunk,
     chunk: ModuleChunk,
-    kotlinContext: KotlinCompileContext
+    kotlinContext: KotlinCompileContext,
+    moduleTarget: BazelModuleBuildTarget,
   ) {
     val context = kotlinContext.jpsContext
-    val representativeTarget = kotlinContext.targetsBinding.get(chunk.representativeTarget()) ?: return
-
-    // dependent caches are not required, since we are not going to update caches
-    val incrementalCaches = kotlinChunk.loadCaches(loadDependent = false)
-
-    val messageCollector = MessageCollectorAdapter(context, span, representativeTarget)
-    val outputItemCollector = OutputItemsCollectorImpl()
-    val environment = createCompileEnvironment(
-      context = kotlinContext.jpsContext,
-      outputItemCollector = outputItemCollector,
-      kotlinModuleBuilderTarget = representativeTarget,
-      incrementalCaches = incrementalCaches,
-      lookupTracker = LookupTracker.DO_NOTHING,
-      exceptActualTracer = ExpectActualTracker.DoNothing,
-      inlineConstTracker = InlineConstTracker.DoNothing,
-      enumWhenTracker = EnumWhenTracker.DoNothing,
-      importTracker = ImportTracker.DoNothing,
-      chunk = chunk,
-      messageCollector = messageCollector,
-    ) ?: return
-
-    val removedClasses = HashSet<String>()
     val dirtyFilesHolder = KotlinDirtySourceFilesHolder(
       chunk = chunk,
       context = context,
@@ -169,12 +152,17 @@ internal class IncrementalKotlinBuilder(
         }
       }
     )
+
+    val removedClasses = hashSet<String>()
+    // dependent caches are not required, since we are not going to update caches
+    val incrementalCaches = kotlinChunk.loadCaches(loadDependent = false)
+    val targetDirtyFiles = dirtyFilesHolder.byTarget.get(moduleTarget)
     for (target in kotlinChunk.targets) {
       val cache = incrementalCaches.get(target) ?: continue
-      val dirtyFiles = dirtyFilesHolder.getDirtyFiles(target.jpsModuleBuildTarget).keys
-      val removedFiles = dirtyFilesHolder.getRemovedFiles(target.jpsModuleBuildTarget)
+      val dirtyFiles = targetDirtyFiles?.dirty?.keys ?: emptySet()
+      val removedFiles = targetDirtyFiles?.removed ?: emptyList()
 
-      val existingClasses = CompilerRunnerUtil.invokeClassesFqNames(environment, dirtyFiles)
+      val existingClasses = classesFqNames(dirtyFiles)
       val previousClasses = cache.classesFqNamesBySources(dirtyFiles + removedFiles)
       for (jvmClassName in previousClasses) {
         val fqName = jvmClassName.asString()
@@ -198,7 +186,7 @@ internal class IncrementalKotlinBuilder(
     val fsOperations = BazelKotlinFsOperationsHelper(context = context, chunk = chunk, span = span, dataManager = dataManager)
     fsOperations.markFilesForCurrentRound(
       files = affectedByRemovedClasses.dirtyFiles.asSequence().map { it.toPath() } + affectedByRemovedClasses.forceRecompileTogether.asSequence().map { it.toPath() },
-      dirtyFilesHolder = dirtyFilesHolder,
+      targetDirtyFiles = targetDirtyFiles,
     )
   }
 
@@ -245,7 +233,7 @@ internal class IncrementalKotlinBuilder(
     context: CompileContext,
     dirtyFilesHolder: BazelDirtyFileHolder,
     messageCollector: MessageCollectorAdapter,
-    outputConsumer: OutputConsumer,
+    outputConsumer: BazelTargetBuildOutputConsumer,
     fsOperations: BazelKotlinFsOperationsHelper,
   ): ExitCode {
     val kotlinContext = context.kotlin
@@ -398,8 +386,9 @@ private fun doCompileModuleChunk(
   environment: JpsCompilerEnvironment,
   incrementalCaches: Map<KotlinModuleBuildTarget<*>, JpsIncrementalCache>,
   messageCollector: MessageCollectorAdapter,
-  outputConsumer: OutputConsumer,
+  outputConsumer: BazelTargetBuildOutputConsumer,
   outputItemCollector: OutputItemsCollectorImpl,
+
   span: Span,
 ): List<GeneratedFile>? {
   val target = chunk.targets.single() as KotlinJvmModuleBuildTarget
@@ -409,11 +398,17 @@ private fun doCompileModuleChunk(
   val jpsTarget = target.jpsModuleBuildTarget
 
   if (cache != null && targetDirtyFiles != null) {
-    val complementaryFiles = cache.getComplementaryFilesRecursive(targetDirtyFiles.dirty.keys + targetDirtyFiles.removed)
+    val dirtyFiles = targetDirtyFiles.dirty.keys.concat(targetDirtyFiles.removed)
+    val complementaryFiles = cache.getComplementaryFilesRecursive(dirtyFiles)
     context.testingContext?.buildLogger?.markedAsComplementaryFiles(complementaryFiles.toList())
-    fsOperations.markFilesForCurrentRound(target = jpsTarget, files = complementaryFiles.map { it.toPath() }, targetDirtyFiles = targetDirtyFiles)
+    fsOperations.markFilesForCurrentRound(
+      files = complementaryFiles.map { it.toPath() },
+      targetDirtyFiles = targetDirtyFiles,
+      parentSpan = span,
+      outputSink = outputConsumer.outputSink,
+    )
 
-    cache.markDirty(targetDirtyFiles.dirty.keys + targetDirtyFiles.removed)
+    cache.markDirty(dirtyFiles)
   }
 
   if (targetDirtyFiles != null) {
@@ -462,7 +457,7 @@ private fun doCompileModuleChunk(
   var outputs: List<OutputFile> = emptyList()
   val pipeline = createJvmPipeline(config) {
     outputs = it.asList()
-    (outputConsumer as BazelTargetBuildOutputConsumer).registerKotlincOutput(context, outputs)
+    outputConsumer.registerKotlincOutput(context, outputs)
   }
 
   val exitCode = executeJvmPipeline(pipeline, bazelConfigurationHolder.kotlinArgs, environment.services, messageCollector)

build/jvm-rules/src/worker-util/collections.kt

@@ -21,9 +21,13 @@ private fun <T> slowEqualsAwareHashStrategy(): Hash.Strategy<T> {
 
 fun <T : Any> linkedSet(): ObjectLinkedOpenCustomHashSet<T> = ObjectLinkedOpenCustomHashSet(slowEqualsAwareHashStrategy())
 
+fun <T : Any> linkedSet(expectedSize: Int): ObjectLinkedOpenCustomHashSet<T> {
+  return ObjectLinkedOpenCustomHashSet(expectedSize, slowEqualsAwareHashStrategy())
+}
+
 fun <T : Any> hashSet(): ObjectOpenCustomHashSet<T> = ObjectOpenCustomHashSet(slowEqualsAwareHashStrategy())
 
-fun <T : Any> hashSet(size: Int): ObjectOpenCustomHashSet<T> = ObjectOpenCustomHashSet(size, slowEqualsAwareHashStrategy())
+fun <T : Any> hashSet(expectedSize: Int): ObjectOpenCustomHashSet<T> = ObjectOpenCustomHashSet(expectedSize, slowEqualsAwareHashStrategy())
 
 fun <K : Any, V : Any> hashMap(): Object2ObjectOpenCustomHashMap<K, V> {
   return Object2ObjectOpenCustomHashMap(slowEqualsAwareHashStrategy())
@@ -39,4 +43,13 @@ fun <K : Any, V : Any> hashMap(size: Int): Object2ObjectOpenCustomHashMap<K, V>
 
 fun <T : Any> emptyList(): List<T> = java.util.List.of()
 
-fun <K : Any, V : Any> emptyMap(): Map<K, V> = java.util.Map.of()
+fun <T : Any> emptySet(): Set<T> = java.util.Set.of()
+
+fun <K : Any, V : Any> emptyMap(): Map<K, V> = java.util.Map.of()
+
+fun <T> Set<T>.concat(collection: Collection<T>): Set<T> {
+  if (collection.isEmpty()) {
+    return this
+  }
+  return this + collection
+}