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

GitOrigin-RevId: 28a2244b1ec90904e13873e18b5284f8eb3c623d

.bazelversion

@@ -1 +1 @@
-8.1.0rc3
+8.1.0

build/jvm-rules/.bazelversion

@@ -1 +1 @@
-8.1.0rc3
+8.1.0

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

@@ -271,10 +271,7 @@ suspend fun buildUsingJps(
     }
   }
 
-  val buildState = if (isRebuild) null
-  else tracer.span("load and check state") { parentSpan ->
-    computeBuildState(parentSpan)
-  }
+  val buildState = if (isRebuild) null else tracer.span("load and check state") { parentSpan -> computeBuildState(parentSpan) }
 
   var exitCode = initAndBuild(
     compileScope = BazelCompileScope(isIncrementalCompilation = true, isRebuild = isRebuild),

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

@@ -100,6 +100,12 @@ internal class BazelStampStorage(private val map: Map<Path, SourceDescriptor>) :
     }
   }
 
+  fun markChanged(sourceFile: Path) {
+    synchronized(map) {
+      map.get(sourceFile)?.isChanged = true
+    }
+  }
+
   override fun getCurrentStampIfUpToDate(file: Path, buildTarget: BuildTarget<*>?, attrs: BasicFileAttributes?): ByteArray? {
     throw UnsupportedOperationException("Must not be used")
   }

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

@@ -3,7 +3,6 @@
 
 package org.jetbrains.bazel.jvm.jps.impl
 
-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
@@ -12,41 +11,50 @@ import org.jetbrains.bazel.jvm.linkedSet
 import org.jetbrains.jps.ModuleChunk
 import org.jetbrains.jps.incremental.CompileContext
 import org.jetbrains.jps.incremental.FSOperations.addCompletelyMarkedDirtyTarget
+import org.jetbrains.jps.incremental.fs.BuildFSState
 import org.jetbrains.jps.incremental.fs.CompilationRound
+import org.jetbrains.jps.incremental.fs.FilesDelta
 import org.jetbrains.kotlin.jps.build.KotlinDirtySourceFilesHolder.TargetFiles
 import java.io.File
 import java.nio.file.Files
 import java.nio.file.Path
-import kotlin.io.path.exists
 
 internal class BazelKotlinFsOperationsHelper(
   private val context: CompileContext,
   private val chunk: ModuleChunk,
-  private val span: Span,
-  private val dataManager: BazelBuildDataProvider,
 ) {
   internal var hasMarkedDirty = false
     private set
 
-  fun markChunk(excludeFiles: Set<File>) {
+  fun markChunk(context: CompileContext, excludeFiles: Set<File>, dataManager: BazelBuildDataProvider) {
     val target = chunk.targets.single()
     var completelyMarkedDirty = true
-    val stampStorage = dataManager.getFileStampStorage(target)
-    for (rootDescriptor in (context.projectDescriptor.buildRootIndex as BazelBuildRootIndex).descriptors) {
+    val stampStorage = if (dataManager.isCleanBuild) null else dataManager.stampStorage
+    val projectDescriptor = context.projectDescriptor
+    for (rootDescriptor in (projectDescriptor.buildRootIndex as BazelBuildRootIndex).descriptors) {
       val file = rootDescriptor.rootFile
       val filePath = file.toString()
-      if (!(FileUtilRt.extensionEquals(filePath, "kt") || FileUtilRt.extensionEquals(filePath, "kts")) ||
-        excludeFiles.contains(file.toFile())) {
+      if ((!filePath.endsWith(".kt") && !filePath.endsWith(".kts")) || excludeFiles.contains(file.toFile())) {
         completelyMarkedDirty = false
         continue
       }
 
       hasMarkedDirty = true
 
-      // if it is a full project rebuild, all storages are already completely cleared;
-      // so passing null because there is no need to access the storage to clear non-existing data
-      val marker = if (dataManager.isCleanBuild) null else stampStorage
-      context.projectDescriptor.fsState.markDirty(context, CompilationRound.NEXT, file, rootDescriptor, marker, false)
+      val roundDelta = context.getUserData(BuildFSState.NEXT_ROUND_DELTA_KEY)
+      roundDelta?.markRecompile(rootDescriptor, file)
+
+      val filesDelta = projectDescriptor.fsState.getDelta(target)
+      filesDelta.lockData()
+      try {
+        val marked = filesDelta.markRecompile(rootDescriptor, file)
+        if (marked) {
+          stampStorage?.markChanged(file)
+        }
+      }
+      finally {
+        filesDelta.unlockData()
+      }
     }
 
     if (completelyMarkedDirty) {
@@ -54,16 +62,6 @@ internal class BazelKotlinFsOperationsHelper(
     }
   }
 
-  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)
-    }
-
-    markFilesImpl(files = files, currentRound = true, span = span) { it.exists() }
-  }
-
   /**
    * Marks given [files] as dirty for current round.
    */
@@ -72,61 +70,116 @@ internal class BazelKotlinFsOperationsHelper(
     targetDirtyFiles: TargetFiles?,
     outputSink: OutputSink,
     parentSpan: Span,
+    target: BazelModuleBuildTarget,
+    dataManager: BazelBuildDataProvider,
   ) {
-    val buildRootIndex = context.projectDescriptor.buildRootIndex as BazelBuildRootIndex
+    val fileToDescriptors = (context.projectDescriptor.buildRootIndex as BazelBuildRootIndex).fileToDescriptors
     for (file in files) {
-      targetDirtyFiles._markDirty(file, buildRootIndex.fileToDescriptors.get(file) ?: continue)
+      targetDirtyFiles._markDirty(file.toFile(), fileToDescriptors.get(file) ?: continue)
     }
 
-    markFilesImpl(files.asSequence(), currentRound = true, span = span) { Files.exists(it) }
+    markFiles(
+      files = files.filterTo(linkedSet()) { Files.exists(it) },
+      currentRound = true,
+      dataManager = dataManager,
+      target = target,
+      span = parentSpan,
+    )
     cleanOutputsCorrespondingToChangedFiles(files = files, dataManager = dataManager, outputSink = outputSink, parentSpan = parentSpan)
   }
 
-  fun markFiles(files: Sequence<Path>) {
-    markFilesImpl(files = files, currentRound = false, span = span) { it.exists() }
-  }
-
-  fun markInChunkOrDependents(files: Sequence<Path>, excludeFiles: Set<Path>) {
-    markFilesImpl(files = files, currentRound = false, span = span) {
-      !excludeFiles.contains(it) && it.exists()
-    }
-  }
-
-  private inline fun markFilesImpl(
-    files: Sequence<Path>,
+  fun markFiles(
+    files: Collection<Path>,
     currentRound: Boolean,
+    target: BazelModuleBuildTarget,
+    dataManager: BazelBuildDataProvider,
     span: Span,
-    shouldMark: (Path) -> Boolean
   ) {
-    val filesToMark = files.filterTo(linkedSet(), shouldMark)
-    if (filesToMark.isEmpty()) {
+    if (files.isEmpty()) {
       return
     }
 
+    val roundDelta: FilesDelta?
     val compilationRound = if (currentRound) {
+      roundDelta = context.getUserData(BuildFSState.CURRENT_ROUND_DELTA_KEY)
       CompilationRound.CURRENT
     }
     else {
+      roundDelta = context.getUserData(BuildFSState.NEXT_ROUND_DELTA_KEY)
       hasMarkedDirty = true
       CompilationRound.NEXT
     }
 
     val projectDescriptor = context.projectDescriptor
+    val stampStorage = dataManager.stampStorage
     val fileToDescriptors = (projectDescriptor.buildRootIndex as BazelBuildRootIndex).fileToDescriptors
-    for (fileToMark in filesToMark) {
-      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,
-      )
+
+    val filesDelta = projectDescriptor.fsState.getDelta(target)
+    filesDelta.lockData()
+    try {
+      for (fileToMark in files) {
+        val rootDescriptor = fileToDescriptors.get(fileToMark) ?: continue
+        roundDelta?.markRecompile(rootDescriptor, fileToMark)
+        val marked = filesDelta.markRecompile(rootDescriptor, fileToMark)
+        if (marked) {
+          stampStorage.markChanged(fileToMark)
+        }
+      }
     }
+    finally {
+      filesDelta.unlockData()
+    }
+
+    if (span.isRecording) {
+      span.addEvent("mark dirty", Attributes.of(
+        AttributeKey.stringArrayKey("filesToMark"), files.map { it.toString() },
+        AttributeKey.stringKey("compilationRound"), compilationRound.name,
+      ))
+    }
+  }
+}
+
+internal fun markFilesForCurrentRound(
+  context: CompileContext,
+  files: Set<File>,
+  targetDirtyFiles: TargetFiles?,
+  span: Span,
+  target: BazelModuleBuildTarget,
+  dataManager: BazelBuildDataProvider,
+) {
+  if (files.isEmpty()) {
+    return
+  }
+
+  val buildRootIndex = context.projectDescriptor.buildRootIndex as BazelBuildRootIndex
+  val fileToDescriptors = buildRootIndex.fileToDescriptors
+  for (file in files) {
+    val root = fileToDescriptors.get(file.toPath()) ?: continue
+    targetDirtyFiles?._markDirty(file, root)
+  }
+
+  val stampStorage = dataManager.stampStorage
+  val roundDelta = context.getUserData(BuildFSState.CURRENT_ROUND_DELTA_KEY)
+  val fileDelta = context.projectDescriptor.fsState.getDelta(target)
+  fileDelta.lockData()
+  try {
+    for (ioFile in files) {
+      val file = ioFile.toPath()
+      val rootDescriptor = fileToDescriptors.get(file) ?: continue
+      roundDelta?.markRecompile(rootDescriptor, file)
+      val marked = fileDelta.markRecompile(rootDescriptor, file)
+      if (marked) {
+        stampStorage.markChanged(file)
+      }
+    }
+  }
+  finally {
+    fileDelta.unlockData()
+  }
+  if (span.isRecording) {
     span.addEvent("mark dirty", Attributes.of(
-      AttributeKey.stringArrayKey("filesToMark"), filesToMark.map { it.toString() },
-      AttributeKey.stringKey("compilationRound"), compilationRound.name,
+      AttributeKey.stringArrayKey("filesToMark"), files.map { it.toString() },
+      AttributeKey.stringKey("compilationRound"), "CURRENT",
     ))
   }
-}
+}

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

@@ -14,10 +14,9 @@ import org.jetbrains.jps.incremental.ModuleLevelBuilder.OutputConsumer
 import org.jetbrains.kotlin.backend.common.output.OutputFile
 import java.io.File
 import java.nio.file.Path
-import java.util.*
 
 internal class BazelTargetBuildOutputConsumer(
-  private val dataManager: BazelBuildDataProvider?,
+  @JvmField val dataManager: BazelBuildDataProvider?,
   @JvmField val outputSink: OutputSink,
 ) : OutputConsumer {
   private var registeredSourceCount = 0

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

@@ -1,5 +1,5 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-@file:Suppress("HardCodedStringLiteral", "INVISIBLE_REFERENCE", "INVISIBLE_MEMBER", "DialogTitleCapitalization", "UnstableApiUsage", "ReplaceGetOrSet")
+@file:Suppress("HardCodedStringLiteral", "DialogTitleCapitalization", "UnstableApiUsage", "ReplaceGetOrSet", "PackageDirectoryMismatch")
 
 package org.jetbrains.bazel.jvm.jps.kotlin
 
@@ -19,6 +19,7 @@ import org.jetbrains.bazel.jvm.jps.impl.BazelKotlinFsOperationsHelper
 import org.jetbrains.bazel.jvm.jps.impl.BazelModuleBuildTarget
 import org.jetbrains.bazel.jvm.jps.impl.BazelTargetBuildOutputConsumer
 import org.jetbrains.bazel.jvm.jps.impl.BazelTargetBuilder
+import org.jetbrains.bazel.jvm.jps.impl.markFilesForCurrentRound
 import org.jetbrains.bazel.jvm.kotlin.configureModule
 import org.jetbrains.bazel.jvm.kotlin.createJvmPipeline
 import org.jetbrains.bazel.jvm.kotlin.executeJvmPipeline
@@ -34,7 +35,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.RebuildRequestedException
 import org.jetbrains.jps.incremental.Utils
 import org.jetbrains.jps.model.module.JpsModule
 import org.jetbrains.kotlin.backend.common.output.OutputFile
@@ -73,9 +73,7 @@ import org.jetbrains.kotlin.jps.build.KotlinChunk
 import org.jetbrains.kotlin.jps.build.KotlinCompileContext
 import org.jetbrains.kotlin.jps.build.KotlinDirtySourceFilesHolder
 import org.jetbrains.kotlin.jps.build.KotlinDirtySourceFilesHolder.TargetFiles
-import org.jetbrains.kotlin.jps.build.kotlin
 import org.jetbrains.kotlin.jps.build.testingContext
-import org.jetbrains.kotlin.jps.incremental.CacheStatus
 import org.jetbrains.kotlin.jps.incremental.JpsIncrementalCache
 import org.jetbrains.kotlin.jps.incremental.JpsLookupStorageManager
 import org.jetbrains.kotlin.jps.targets.KotlinJvmModuleBuildTarget
@@ -87,7 +85,6 @@ import org.jetbrains.kotlin.metadata.deserialization.MetadataVersion
 import org.jetbrains.kotlin.modules.TargetId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.preloading.ClassCondition
-import org.jetbrains.kotlin.progress.CompilationCanceledException
 import org.jetbrains.kotlin.progress.CompilationCanceledStatus
 import java.io.File
 import java.nio.file.Files
@@ -115,24 +112,19 @@ internal class IncrementalKotlinBuilder(
 
   override fun chunkBuildStarted(context: CompileContext, chunk: ModuleChunk) {
     val kotlinContext = ensureKotlinContextInitialized(context, span)
-    if (isRebuild) {
-      return
-    }
-
-    if (chunk.targets.none { kotlinContext.hasKotlinMarker.get(it) == true }) {
+    if (isRebuild || kotlinContext.hasKotlinMarker.get(chunk.targets.single()) != true) {
       return
     }
 
     val kotlinChunk = kotlinContext.getChunk(chunk) ?: return
-    if (!kotlinContext.rebuildingAllKotlin) {
-      val target = kotlinChunk.targets.single()
-      if (target.initialLocalCacheAttributesDiff.status == CacheStatus.INVALID) {
-        throw RebuildRequestedException(RuntimeException("cache is invalid, rebuilding (diff=${target.initialLocalCacheAttributesDiff}"))
-      }
-
-      if (kotlinChunk.isEnabled) {
-        markAdditionalFilesForInitialRound(kotlinChunk, chunk, kotlinContext, jpsTarget)
-      }
+    if (kotlinChunk.isEnabled) {
+      markAdditionalFilesForInitialRound(
+        kotlinChunk = kotlinChunk,
+        chunk = chunk,
+        kotlinContext = kotlinContext,
+        moduleTarget = jpsTarget,
+        span = span,
+      )
     }
   }
 
@@ -141,6 +133,7 @@ internal class IncrementalKotlinBuilder(
     chunk: ModuleChunk,
     kotlinContext: KotlinCompileContext,
     moduleTarget: BazelModuleBuildTarget,
+    span: Span,
   ) {
     val context = kotlinContext.jpsContext
     val dirtyFilesHolder = KotlinDirtySourceFilesHolder(
@@ -150,7 +143,7 @@ internal class IncrementalKotlinBuilder(
         override fun processDirtyFiles(processor: FileProcessor<JavaSourceRootDescriptor, ModuleBuildTarget>) {
           context.projectDescriptor.fsState.processFilesToRecompile(context, chunk.targets.single(), processor)
         }
-      }
+      },
     )
 
     val removedClasses = hashSet<String>()
@@ -183,10 +176,13 @@ internal class IncrementalKotlinBuilder(
       reporter = BazelJpsICReporter(span),
     )
 
-    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() },
+    markFilesForCurrentRound(
+      files = affectedByRemovedClasses.dirtyFiles.concat(affectedByRemovedClasses.forceRecompileTogether),
       targetDirtyFiles = targetDirtyFiles,
+      span = span,
+      context = context,
+      target = moduleTarget,
+      dataManager = dataManager,
     )
   }
 
@@ -199,30 +195,26 @@ internal class IncrementalKotlinBuilder(
     outputConsumer: BazelTargetBuildOutputConsumer,
     outputSink: OutputSink
   ): ModuleLevelBuilder.ExitCode {
-    val kotlinTarget = context.kotlin.targetsBinding.get(chunk.representativeTarget()) ?: return ModuleLevelBuilder.ExitCode.OK
-    val fsOperations = BazelKotlinFsOperationsHelper(
-      context = context,
+    val kotlinContext = getKotlinCompileContext(context)
+    val kotlinTarget = kotlinContext.targetsIndex.byJpsTarget.get(jpsTarget) ?: return ModuleLevelBuilder.ExitCode.OK
+    val fsOperations = BazelKotlinFsOperationsHelper(context = context, chunk = chunk)
+    val proposedExitCode = doBuild(
       chunk = chunk,
-      dataManager = dataManager,
-      span = span,
+      representativeTarget = kotlinTarget,
+      context = context,
+      dirtyFilesHolder = dirtyFilesHolder,
+      messageCollector = MessageCollectorAdapter(context, span, kotlinTarget),
+      outputConsumer = outputConsumer,
+      fsOperations = fsOperations,
+      kotlinContext = kotlinContext,
     )
-    val proposedExitCode = try {
-      doBuild(
-        chunk = chunk,
-        representativeTarget = kotlinTarget,
-        context = context,
-        dirtyFilesHolder = dirtyFilesHolder,
-        messageCollector = MessageCollectorAdapter(context, span, kotlinTarget),
-        outputConsumer = outputConsumer,
-        fsOperations = fsOperations,
-      )
-    }
-    catch (e: CompilationCanceledException) {
-      // https://youtrack.jetbrains.com/issue/KTI-2139
-      throw CancellationException(e)
-    }
 
-    val actualExitCode = if (proposedExitCode == ModuleLevelBuilder.ExitCode.OK && fsOperations.hasMarkedDirty) ModuleLevelBuilder.ExitCode.ADDITIONAL_PASS_REQUIRED else proposedExitCode
+    val actualExitCode = if (proposedExitCode == ModuleLevelBuilder.ExitCode.OK && fsOperations.hasMarkedDirty) {
+      ModuleLevelBuilder.ExitCode.ADDITIONAL_PASS_REQUIRED
+    }
+    else {
+      proposedExitCode
+    }
     context.testingContext?.buildLogger?.buildFinished(actualExitCode)
     return actualExitCode
   }
@@ -235,8 +227,8 @@ internal class IncrementalKotlinBuilder(
     messageCollector: MessageCollectorAdapter,
     outputConsumer: BazelTargetBuildOutputConsumer,
     fsOperations: BazelKotlinFsOperationsHelper,
+    kotlinContext: KotlinCompileContext,
   ): ExitCode {
-    val kotlinContext = context.kotlin
     val kotlinChunk = kotlinContext.getChunk(chunk)!!
     if (!kotlinChunk.isEnabled) {
       return ModuleLevelBuilder.ExitCode.NOTHING_DONE
@@ -247,7 +239,7 @@ internal class IncrementalKotlinBuilder(
     val isChunkRebuilding = isRebuild || kotlinContext.rebuildAfterCacheVersionChanged.get(target) == true
 
     val kotlinDirtyFilesHolder = KotlinDirtySourceFilesHolder(chunk, context, dirtyFilesHolder)
-    val dirtyByTarget = kotlinDirtyFilesHolder.byTarget.get(dirtyFilesHolder.target)
+    val dirtyByTarget = kotlinDirtyFilesHolder.byTarget.get(jpsTarget)
     if (dirtyByTarget == null || (dirtyByTarget.removed.isEmpty() && dirtyByTarget.dirty.isEmpty())) {
       if (isChunkRebuilding) {
         kotlinContext.hasKotlinMarker.set(target, false)
@@ -280,7 +272,6 @@ internal class IncrementalKotlinBuilder(
     val generatedFiles = doCompileModuleChunk(
       chunk = kotlinChunk,
       outputItemCollector = outputItemCollector,
-      representativeTarget = representativeTarget,
       context = context,
       targetDirtyFiles = dirtyByTarget,
       fsOperations = fsOperations,
@@ -304,17 +295,22 @@ internal class IncrementalKotlinBuilder(
       JavaBuilderUtil.registerSuccessfullyCompiled(context, dirtyByTarget.dirty.keys.toList())
     }
 
-    markDirtyComplementaryMultifileClasses(
-      files = generatedFiles,
-      kotlinModuleBuilderTarget = representativeTarget,
-      incrementalCaches = incrementalCaches,
-      fsOperations = fsOperations,
-    )
+    val cache = incrementalCaches.get(representativeTarget)
+    if (cache is IncrementalJvmCache) {
+      markDirtyComplementaryMultifileClasses(
+        files = generatedFiles,
+        cache = cache,
+        fsOperations = fsOperations,
+        target = jpsTarget,
+        dataManager = dataManager,
+        span = span,
+      )
+    }
 
     // we do not save cache version - TargetConfigurationDigestProperty.KOTLIN_VERSION is used to rebuild in case of kotlinc update
 
     if (kotlinContext.hasKotlinMarker.get(target) == null) {
-      fsOperations.markChunk(excludeFiles = dirtyByTarget.dirty.keys)
+      fsOperations.markChunk(context = context, excludeFiles = dirtyByTarget.dirty.keys, dataManager = dataManager)
     }
 
     kotlinContext.hasKotlinMarker.set(target, true)
@@ -363,11 +359,14 @@ internal class IncrementalKotlinBuilder(
     if (!isChunkRebuilding) {
       doProcessChangesUsingLookups(
         collector = changeCollector,
-        compiledFiles = dirtyByTarget.dirty.keys.mapTo(linkedSet()) { it.toPath() },
+        compiledFiles = dirtyByTarget.dirty.keys.mapTo(hashSet()) { it.toPath() },
         lookupStorageManager = kotlinContext.lookupStorageManager,
         fsOperations = fsOperations,
         caches = incrementalCaches.values,
         reporter = BazelJpsICReporter(span),
+        target = jpsTarget,
+        dataManager = dataManager,
+        span = span,
       )
     }
 
@@ -379,7 +378,6 @@ internal class IncrementalKotlinBuilder(
 // todo(1.2.80): introduce KotlinRoundCompileContext, move dirtyFilesHolder, fsOperations, environment to it
 private fun doCompileModuleChunk(
   chunk: KotlinChunk,
-  representativeTarget: KotlinModuleBuildTarget<*>,
   context: CompileContext,
   targetDirtyFiles: TargetFiles?,
   fsOperations: BazelKotlinFsOperationsHelper,
@@ -388,14 +386,13 @@ private fun doCompileModuleChunk(
   messageCollector: MessageCollectorAdapter,
   outputConsumer: BazelTargetBuildOutputConsumer,
   outputItemCollector: OutputItemsCollectorImpl,
-
   span: Span,
 ): List<GeneratedFile>? {
   val target = chunk.targets.single() as KotlinJvmModuleBuildTarget
   target.nextRound(context)
 
   val cache = incrementalCaches.get(target)
-  val jpsTarget = target.jpsModuleBuildTarget
+  val jpsTarget = target.jpsModuleBuildTarget as BazelModuleBuildTarget
 
   if (cache != null && targetDirtyFiles != null) {
     val dirtyFiles = targetDirtyFiles.dirty.keys.concat(targetDirtyFiles.removed)
@@ -406,6 +403,8 @@ private fun doCompileModuleChunk(
       targetDirtyFiles = targetDirtyFiles,
       parentSpan = span,
       outputSink = outputConsumer.outputSink,
+      dataManager = outputConsumer.dataManager!!,
+      target = jpsTarget,
     )
 
     cache.markDirty(dirtyFiles)
@@ -419,7 +418,6 @@ private fun doCompileModuleChunk(
     registerFilesToCompile(context, allDirtyFiles)
   }
 
-  require(chunk.representativeTarget == representativeTarget)
   val filesSet = targetDirtyFiles?.dirty?.keys ?: emptySet()
   val changedSources = targetDirtyFiles?.dirty?.values ?: emptyList()
   val removedFiles = targetDirtyFiles?.removed ?: emptyList()
@@ -546,27 +544,41 @@ private fun updateLookupStorage(
 
 private fun markDirtyComplementaryMultifileClasses(
   files: List<GeneratedFile>,
-  incrementalCaches: Map<KotlinModuleBuildTarget<*>, JpsIncrementalCache>,
+  cache: IncrementalJvmCache,
   fsOperations: BazelKotlinFsOperationsHelper,
-  kotlinModuleBuilderTarget: KotlinModuleBuildTarget<*>,
+  target: BazelModuleBuildTarget,
+  dataManager: BazelBuildDataProvider,
+  span: Span,
 ) {
-  val cache = incrementalCaches.get(kotlinModuleBuilderTarget) as? IncrementalJvmCache ?: return
-  val generated = files.filterIsInstance<GeneratedJvmClass>()
-  val multifileClasses = generated.filter { it.outputClass.classHeader.kind == KotlinClassHeader.Kind.MULTIFILE_CLASS }
-  val expectedAllParts = multifileClasses.flatMap { cache.getAllPartsOfMultifileFacade(it.outputClass.className).orEmpty() }
+  val multifileClasses = files
+    .asSequence()
+    .filterIsInstance<GeneratedJvmClass>()
+    .filter { it.outputClass.classHeader.kind == KotlinClassHeader.Kind.MULTIFILE_CLASS }
+    .toList()
   if (multifileClasses.isEmpty()) {
     return
   }
 
-  val actualParts = generated
+  val actualParts = files
     .asSequence()
+    .filterIsInstance<GeneratedJvmClass>()
     .filter { it.outputClass.classHeader.kind == KotlinClassHeader.Kind.MULTIFILE_CLASS_PART }
     .map { it.outputClass.className.toString() }
     .toList()
+
+  val expectedAllParts = multifileClasses.flatMap { cache.getAllPartsOfMultifileFacade(it.outputClass.className).orEmpty() }
   if (!actualParts.containsAll(expectedAllParts)) {
     fsOperations.markFiles(
-      expectedAllParts.asSequence().flatMap { cache.sourcesByInternalName(it) }.map { it.toPath() }
-        + multifileClasses.asSequence().flatMap { it.sourceFiles }.map { it.toPath() }
+      files = (
+        expectedAllParts.asSequence().flatMap { cache.sourcesByInternalName(it) }
+          + multifileClasses.asSequence().flatMap { it.sourceFiles }
+        )
+        .map { it.toPath() }
+        .filterTo(linkedSet()) { Files.exists(it) },
+      currentRound = false,
+      target = target,
+      dataManager = dataManager,
+      span = span,
     )
   }
 }
@@ -587,26 +599,38 @@ private fun doProcessChangesUsingLookups(
   fsOperations: BazelKotlinFsOperationsHelper,
   caches: Iterable<JpsIncrementalCache>,
   reporter: ICReporter,
+  target: BazelModuleBuildTarget,
+  dataManager: BazelBuildDataProvider,
+  span: Span,
 ) {
-  val allCaches = caches.flatMap { it.thisWithDependentCaches }
-
   val dirtyFiles = getDirtyFiles(
     changeCollector = collector,
-    caches = allCaches,
+    caches = caches.flatMap { it.thisWithDependentCaches },
     lookupStorageManager = lookupStorageManager,
     reporter = reporter,
   )
-  // if a list of inheritors of sealed class has changed it should be recompiled with all the inheritors
+  // If a list of inheritors of sealed class has changed, it should be recompiled with all the inheritors
   // Here we have a small optimization. Do not recompile the bunch if ALL these files were recompiled during the previous round.
-  val excludeFiles = if (compiledFiles.containsAll(dirtyFiles.forceRecompileTogether.map { it.toPath() })) {
+  val forceRecompileTogether = dirtyFiles.forceRecompileTogether
+  val excludeFiles = if (forceRecompileTogether.isEmpty() || compiledFiles.containsAll(forceRecompileTogether.map { it.toPath() })) {
     compiledFiles
   }
   else {
-    compiledFiles.minus(dirtyFiles.forceRecompileTogether.map { it.toPath() })
+    val result = HashSet(compiledFiles)
+    for (file in forceRecompileTogether) {
+      result.remove(file.toPath())
+    }
+    result
   }
-  fsOperations.markInChunkOrDependents(
-    files = (dirtyFiles.dirtyFiles.asSequence().map { it.toPath() } + dirtyFiles.forceRecompileTogether.asSequence().map { it.toPath() }),
-    excludeFiles = excludeFiles,
+
+  fsOperations.markFiles(
+    files = (dirtyFiles.dirtyFiles.asSequence() + forceRecompileTogether.asSequence())
+      .map { it.toPath() }
+      .filterTo(linkedSet()) { !excludeFiles.contains(it) && Files.exists(it) },
+    currentRound = false,
+    dataManager = dataManager,
+    target = target,
+    span = span,
   )
 }
 
@@ -624,6 +648,6 @@ private fun getDirtyFiles(
   }
   return FilesToRecompile(
     dirtyFiles = dirtyFilesFromLookups + mapClassesFqNamesToFiles(caches, dirtyClassFqNames, reporter),
-    forceRecompileTogether = mapClassesFqNamesToFiles(caches, forceRecompile, reporter)
+    forceRecompileTogether = mapClassesFqNamesToFiles(caches, forceRecompile, reporter),
   )
 }

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

@@ -0,0 +1,33 @@
+// Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+@file:Suppress("unused", "PackageDirectoryMismatch", "INVISIBLE_REFERENCE", "INVISIBLE_MEMBER", "ReplaceJavaStaticMethodWithKotlinAnalog")
+
+package org.jetbrains.kotlin.jps.targets
+
+import org.jetbrains.bazel.jvm.jps.impl.BazelBuildTargetIndex
+import org.jetbrains.jps.incremental.ModuleBuildTarget
+import org.jetbrains.kotlin.jps.build.KotlinChunk
+import org.jetbrains.kotlin.jps.build.KotlinCompileContext
+
+class KotlinTargetsIndex(
+  val byJpsTarget: Map<ModuleBuildTarget, KotlinModuleBuildTarget<*>>,
+  val chunks: List<KotlinChunk>,
+  val chunksByJpsRepresentativeTarget: Map<ModuleBuildTarget, KotlinChunk>
+)
+
+internal class KotlinTargetsIndexBuilder internal constructor(
+  private val uninitializedContext: KotlinCompileContext
+) {
+  fun build(): KotlinTargetsIndex {
+    val target = (uninitializedContext.jpsContext.projectDescriptor.buildTargetIndex as BazelBuildTargetIndex).moduleTarget
+    val moduleBuildTarget = KotlinJvmModuleBuildTarget(uninitializedContext, target)
+    val kotlinChunk = KotlinChunk(uninitializedContext, java.util.List.of(moduleBuildTarget))
+    moduleBuildTarget.chunk = kotlinChunk
+    val chunks = listOf(kotlinChunk)
+    KotlinChunk.calculateChunkDependencies(chunks, java.util.Map.of<ModuleBuildTarget, KotlinModuleBuildTarget<*>>() as MutableMap<ModuleBuildTarget, KotlinModuleBuildTarget<*>>)
+    return KotlinTargetsIndex(
+      byJpsTarget = java.util.Map.of(target, moduleBuildTarget),
+      chunks = chunks,
+      chunksByJpsRepresentativeTarget = java.util.Map.of(target, kotlinChunk)
+    )
+  }
+}

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

@@ -10,54 +10,41 @@ import org.jetbrains.kotlin.jps.build.kotlinCompileContextKey
 import org.jetbrains.kotlin.jps.build.testingContext
 import kotlin.system.measureTimeMillis
 
-/**
- * Ensure Kotlin Context initialized.
- * Kotlin Context should be initialized only when required (before the first kotlin chunk build).
- */
+internal fun getKotlinCompileContext(context: CompileContext): KotlinCompileContext = context.getUserData(kotlinCompileContextKey)!!
+
 internal fun ensureKotlinContextInitialized(context: CompileContext, span: Span): KotlinCompileContext {
   context.getUserData(kotlinCompileContextKey)?.let {
     return it
   }
 
-  // don't synchronize on context, since it is chunk local only
-  synchronized(kotlinCompileContextKey) {
-    context.getUserData(kotlinCompileContextKey)?.let {
-      return it
-    }
-    return initializeKotlinContext(context, span)
+  synchronized(context) {
+    return context.getUserData(kotlinCompileContextKey) ?: initializeKotlinContext(context, span)
   }
 }
 
 internal fun ensureKotlinContextDisposed(context: CompileContext) {
-  if (context.getUserData(kotlinCompileContextKey) != null) {
-    synchronized(kotlinCompileContextKey) {
-      val kotlinCompileContext = context.getUserData(kotlinCompileContextKey)
-      if (kotlinCompileContext != null) {
-        kotlinCompileContext.dispose()
-        context.putUserData(kotlinCompileContextKey, null)
-      }
-    }
+  if (context.getUserData(kotlinCompileContextKey) == null) {
+    return
+  }
+
+  synchronized(context) {
+    val kotlinCompileContext = context.getUserData(kotlinCompileContextKey) ?: return
+    kotlinCompileContext.dispose()
+    context.putUserData(kotlinCompileContextKey, null)
   }
 }
 
 private fun initializeKotlinContext(context: CompileContext, span: Span): KotlinCompileContext {
   lateinit var kotlinContext: KotlinCompileContext
-
   val time = measureTimeMillis {
     kotlinContext = KotlinCompileContext(context)
 
     context.putUserData(kotlinCompileContextKey, kotlinContext)
     context.testingContext?.kotlinCompileContext = kotlinContext
-
-    if (kotlinContext.shouldCheckCacheVersions && kotlinContext.hasKotlin()) {
-      kotlinContext.checkCacheVersions()
-    }
-
     kotlinContext.cleanupCaches()
     kotlinContext.reportUnsupportedTargets()
   }
 
   span.addEvent("total Kotlin global compile context initialization time: $time ms")
-
   return kotlinContext
 }

lib/.bazelversion

@@ -1 +1 @@
-8.1.0rc3
+8.1.0