Add cancelation support for Kotlin frontend.

This currently needs to be wired in via ProgressIndicatorAndCompilationCanceledStatus. Noticeably they are missing a lot of checks in between works that causes many large 1000ms+ delays in my local testing. If it turns out to be a problem, we will need to change couple of places to add checks in more places within kotlinc.

Also unfortunately the mechanism they have relies on static state so doesn't support multiplex workers. In the future I might workaround that using thread locals but for now I kept it simple.

PiperOrigin-RevId: 728739574

transpiler/java/com/google/j2cl/transpiler/frontend/kotlin/CompilationUnitBuilder.kt

@@ -189,6 +189,7 @@ import org.jetbrains.kotlin.ir.util.parentAsClass
 import org.jetbrains.kotlin.ir.util.render
 import org.jetbrains.kotlin.ir.util.statements
 import org.jetbrains.kotlin.name.SpecialNames
+import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 
 /** Creates a J2CL Java AST from Kotlin IR. */
 class CompilationUnitBuilder(
@@ -224,6 +225,8 @@ class CompilationUnitBuilder(
   private fun convertClass(irClass: IrClass): Type {
     val type = Type(getNameSourcePosition(irClass), environment.getDeclarationForType(irClass))
     processEnclosedBy(type) {
+      ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+
       // Skip synthetic declarations. Kotlinc adds synthetic declarations like (fake) override
       // members
       // to help with bridge synthesis and the resolution phase.
@@ -245,8 +248,9 @@ class CompilationUnitBuilder(
     return type
   }
 
-  private fun convertDeclaration(irDeclaration: IrDeclaration): List<Member> =
-    when (irDeclaration) {
+  private fun convertDeclaration(irDeclaration: IrDeclaration): List<Member> {
+    ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+    return when (irDeclaration) {
       is IrEnumEntry -> listOf(convertEnumEntry(irDeclaration))
       is IrProperty -> convertProperty(irDeclaration)
       // Lowering passes can add field on object classes.
@@ -255,6 +259,7 @@ class CompilationUnitBuilder(
       is IrAnonymousInitializer -> listOf(convertAnonymousInitializer(irDeclaration))
       else -> throw NotImplementedError("Declaration not yet supported: $irDeclaration")
     }
+  }
 
   private fun convertEnumEntry(irEnumEntry: IrEnumEntry): Field {
     val initializerExpression = requireNotNull(irEnumEntry.initializerExpression).expression

transpiler/java/com/google/j2cl/transpiler/frontend/kotlin/KotlinParser.kt

@@ -77,27 +77,46 @@ import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmMetadataVersion
 import org.jetbrains.kotlin.metadata.jvm.deserialization.JvmProtoBufUtil
 import org.jetbrains.kotlin.modules.TargetId
+import org.jetbrains.kotlin.progress.CompilationCanceledStatus
+import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 
 /** A parser for Kotlin sources that builds {@code CompilationtUnit}s. */
 class KotlinParser(private val problems: Problems) {
 
   /** Returns a list of compilation units after Kotlinc parsing. */
   fun parseFiles(options: FrontendOptions): Library {
     val packageInfoCache = PackageInfoCache(options.classpaths, problems)
+    problems.abortIfCancelled()
     val packageAnnotationResolver = getPackageAnnotationResolver(options, packageInfoCache)
+    problems.abortIfCancelled()
 
     val kotlincDisposable = Disposer.newDisposable("J2CL Root Disposable")
     val compilerConfiguration = createCompilerConfiguration(options, packageInfoCache)
+    problems.abortIfCancelled()
 
     check(compilerConfiguration.getBoolean(USE_FIR)) { "Kotlin/Closure only supports > K2" }
 
-    val compilationUnits =
-      parseFiles(compilerConfiguration, kotlincDisposable, packageAnnotationResolver)
+    try {
+      // TODO(b/148292139): This assumes single worker for given time and doesn't work with
+      // multiplex workers.
+      ProgressIndicatorAndCompilationCanceledStatus.setCompilationCanceledStatus(
+        object : CompilationCanceledStatus {
+          override fun checkCanceled() {
+            problems.abortIfCancelled()
+          }
+        }
+      )
+
+      val compilationUnits =
+        parseFiles(compilerConfiguration, kotlincDisposable, packageAnnotationResolver)
 
-    return Library.newBuilder()
-      .setCompilationUnits(compilationUnits)
-      .setDisposableListener { Disposer.dispose(kotlincDisposable) }
-      .build()
+      return Library.newBuilder()
+        .setCompilationUnits(compilationUnits)
+        .setDisposableListener { Disposer.dispose(kotlincDisposable) }
+        .build()
+    } finally {
+      ProgressIndicatorAndCompilationCanceledStatus.setCompilationCanceledStatus(null)
+    }
   }
 
   private fun parseFiles(
@@ -114,21 +133,20 @@ class KotlinParser(private val problems: Problems) {
         EnvironmentConfigFiles.JVM_CONFIG_FILES,
         messageCollector,
       )
+    problems.abortIfCancelled()
 
     val module = compilerConfiguration.get(MODULES)!![0]
     val diagnosticsReporter = DiagnosticReporterFactory.createPendingReporter(messageCollector)
+    val sources = collectSources(compilerConfiguration, projectEnvironment, messageCollector)
+    problems.abortIfCancelled()
 
     val analysisResults =
       @OptIn(IncrementalCompilationApi::class)
       compileModuleToAnalyzedFirViaLightTreeIncrementally(
         projectEnvironment,
         messageCollector,
         compilerConfiguration,
-        ModuleCompilerInput(
-          TargetId(module),
-          collectSources(compilerConfiguration, projectEnvironment, messageCollector),
-          compilerConfiguration,
-        ),
+        ModuleCompilerInput(TargetId(module), sources, compilerConfiguration),
         diagnosticsReporter,
         incrementalExcludesScope = null,
       )
@@ -150,6 +168,7 @@ class KotlinParser(private val problems: Problems) {
         .isIrBackend(true)
         .diagnosticReporter(diagnosticsReporter)
         .build()
+    problems.abortIfCancelled()
 
     val compilationUnitBuilderExtension =
       createAndRegisterCompilationUnitBuilder(
@@ -158,6 +177,7 @@ class KotlinParser(private val problems: Problems) {
         state,
         packageAnnotationResolver,
       )
+    problems.abortIfCancelled()
 
     val unused =
       analysisResults.convertToIrAndActualizeForJvm(

transpiler/java/com/google/j2cl/transpiler/frontend/kotlin/lower/LoweringPasses.kt

@@ -53,6 +53,7 @@ import org.jetbrains.kotlin.ir.backend.js.lower.inline.RemoveInlineDeclarationsW
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
 import org.jetbrains.kotlin.ir.util.SymbolTable
+import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 
 /** The list of lowering passes to run in execution order. */
 private val loweringPassFactories: List<J2clLoweringPassFactory> = buildList {
@@ -262,6 +263,7 @@ private fun IrModuleFragment.lower(
 ) {
   for (f in files) {
     try {
+      ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
       loweringFactory(context).lower(f)
     } catch (e: CompilationException) {
       e.initializeFileDetails(f)

transpiler/javatests/com/google/j2cl/transpiler/TranspilerTester.java

@@ -550,6 +550,7 @@ public boolean isCancelled() {
 
     final String[] knownDelayedCalls = {
       "com.google.j2cl.transpiler.frontend.javac.JavacParser.parseFiles",
+      // Kotlin frontend is currently missing a lot of checks in between large chunks of works.
       "com.google.j2cl.transpiler.frontend.kotlin.KotlinParser.parseFiles",
     };
     for (String knownDelayedCall : knownDelayedCalls) {