Rewrite the NoComposableFileChecker logic and test code

build.gradle.kts

@@ -9,6 +9,7 @@ import com.adarshr.gradle.testlogger.TestLoggerExtension
 import com.adarshr.gradle.testlogger.theme.ThemeType
 import com.diffplug.gradle.spotless.BaseKotlinExtension
 import com.diffplug.gradle.spotless.SpotlessExtension
+import org.jetbrains.kotlin.gradle.dsl.KotlinTopLevelExtension
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 
 plugins {
@@ -82,6 +83,12 @@ allprojects {
       )
     }
   }
+
+  afterEvaluate {
+    extensions.configure<KotlinTopLevelExtension> {
+      jvmToolchain(22)
+    }
+  }
 }
 
 subprojects {

compiler/src/main/kotlin/land/sungbin/composeinvestigator/compiler/ComposeInvestigatorPluginRegistrar.kt

@@ -22,7 +22,7 @@ import org.jetbrains.kotlin.config.messageCollector
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 
 public class ComposeInvestigatorPluginRegistrar : ComponentRegistrar {
-  override val supportsK2: Boolean = true
+  override val supportsK2: Boolean get() = true
 
   // This deprecated override is safe to use up to Kotlin 2.1.0 by KT-55300.
   // Also see: https://youtrack.jetbrains.com/issue/KT-52665/Deprecate-ComponentRegistrar#focus=Change-27-7999959.0-0

compiler/src/main/kotlin/land/sungbin/composeinvestigator/compiler/frontend/InvalidationTraceTableInstantiationValidator.kt

@@ -8,6 +8,7 @@
 package land.sungbin.composeinvestigator.compiler.frontend
 
 import androidx.compose.compiler.plugins.kotlin.k2.hasComposableAnnotation
+import androidx.compose.compiler.plugins.kotlin.k2.isComposable
 import androidx.compose.compiler.plugins.kotlin.lower.fastForEach
 import land.sungbin.composeinvestigator.compiler.NO_INVESTIGATION_FQN
 import land.sungbin.composeinvestigator.compiler.lower.unsafeLazy
@@ -21,11 +22,13 @@ import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationChecker
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.FirFileChecker
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
 import org.jetbrains.kotlin.fir.declarations.FirFile
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.declarations.validate
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.builder.buildAnnotation
 import org.jetbrains.kotlin.fir.expressions.impl.FirEmptyAnnotationArgumentMapping
-import org.jetbrains.kotlin.fir.references.toResolvedFunctionSymbol
+import org.jetbrains.kotlin.fir.references.toResolvedCallableSymbol
 import org.jetbrains.kotlin.fir.smartPlus
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.constructClassLikeType
@@ -39,13 +42,15 @@ public class InvalidationTraceTableInstantiationValidator(session: FirSession) :
 }
 
 private object NoComposableFileChecker : FirFileChecker(MppCheckerKind.Common) {
+  private val NO_INVESTIGATION = ClassId.topLevel(NO_INVESTIGATION_FQN)
   private val noInvestigationType by unsafeLazy {
     buildResolvedTypeRef {
-      coneType = ClassId.topLevel(NO_INVESTIGATION_FQN).constructClassLikeType()
+      coneType = NO_INVESTIGATION.constructClassLikeType()
     }
   }
 
   override fun check(declaration: FirFile, context: CheckerContext, reporter: DiagnosticReporter) {
+    if (declaration.hasAnnotation(NO_INVESTIGATION, context.session)) return
     var hasComposable = false
 
     val composableCallVisitor = object : FirDefaultVisitorVoid() {
@@ -54,17 +59,25 @@ private object NoComposableFileChecker : FirFileChecker(MppCheckerKind.Common) {
         element.acceptChildren(this)
       }
 
-      override fun visitFunctionCall(functionCall: FirFunctionCall) {
-        if (functionCall.calleeReference.toResolvedFunctionSymbol()!!.hasComposableAnnotation(context.session))
+      override fun visitFunctionCall(call: FirFunctionCall) {
+        if (call.calleeReference.toResolvedCallableSymbol()!!.isComposable(context.session))
           hasComposable = true
 
         if (hasComposable) return
-        super.visitFunctionCall(functionCall)
+
+        super.visitFunctionCall(call)
       }
     }
 
     declaration.declarations.fastForEach { element ->
-      element.acceptChildren(composableCallVisitor)
+      // fast path -- 1
+      if (element.hasComposableAnnotation(context.session))
+        return@check // early return if the file has composable functions
+
+      if (element is FirFunction) // fast path -- 2
+        element.body?.acceptChildren(composableCallVisitor)
+      else // slow path
+        element.acceptChildren(composableCallVisitor)
     }
 
     if (hasComposable) return

compiler/src/test/kotlin/land/sungbin/composeinvestigator/compiler/_compilation/AbstractCompilerTest.kt

@@ -18,6 +18,8 @@ import land.sungbin.composeinvestigator.compiler.ComposeInvestigatorFirExtension
 import land.sungbin.composeinvestigator.compiler.ComposeInvestigatorFirstPhaseExtension
 import land.sungbin.composeinvestigator.compiler.ComposeInvestigatorLastPhaseExtension
 import land.sungbin.composeinvestigator.compiler.FeatureFlag
+import land.sungbin.composeinvestigator.compiler._source.sourcePath
+import land.sungbin.composeinvestigator.compiler._source.sourceString
 import land.sungbin.composeinvestigator.runtime.ComposeInvestigatorConfig
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.cli.jvm.config.addJvmClasspathRoots
@@ -38,41 +40,12 @@ import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.config.messageCollector
 import org.jetbrains.kotlin.fir.extensions.FirExtensionRegistrarAdapter
 import org.jetbrains.kotlin.fir.pipeline.Fir2IrActualizedResult
+import org.jetbrains.kotlin.ir.declarations.IrFile
 
 abstract class AbstractCompilerTest(
-  private val features: EnumSet<FeatureFlag> = EnumSet.noneOf(FeatureFlag::class.java),
+  private val features: EnumSet<FeatureFlag> = NO_FEATURES,
+  private val sourceRoot: String? = null,
 ) {
-  companion object {
-    private fun File.applyExistenceCheck(): File = apply {
-      if (!exists()) throw NoSuchFileException(this)
-    }
-
-    private val homeDir: String = run {
-      val userDir = System.getProperty("user.dir")
-      val dir = File(userDir ?: ".")
-      val path = FileUtil.toCanonicalPath(dir.absolutePath)
-      File(path).applyExistenceCheck().absolutePath
-    }
-
-    // https://github.com/JetBrains/kotlin/blob/bb25d2f8aa74406ff0af254b2388fd601525386a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/AbstractCompilerTest.kt#L212-L228
-    val defaultClassPath by lazy {
-      listOf(
-        jarFor<Unit>(),
-        jarFor<kotlinx.coroutines.CoroutineScope>(),
-        jarFor<androidx.compose.runtime.Composable>(),
-        jarFor<androidx.compose.animation.EnterTransition>(),
-        jarFor<androidx.compose.ui.Modifier>(),
-        jarFor<androidx.compose.ui.graphics.ColorProducer>(),
-        jarFor<androidx.compose.ui.unit.Dp>(),
-        jarFor<androidx.compose.ui.text.input.TextFieldValue>(),
-        jarFor<androidx.compose.foundation.Indication>(),
-        jarFor<androidx.compose.foundation.text.KeyboardActions>(),
-        jarFor<androidx.compose.foundation.layout.RowScope>(),
-        jarFor<ComposeInvestigatorConfig>(),
-      )
-    }
-  }
-
   private val disposable = Disposer.newDisposable()
 
   @BeforeTest fun setSystemProperties() {
@@ -146,12 +119,197 @@ abstract class AbstractCompilerTest(
     createK2Compiler().compile(file)
 
   protected fun irTest(file: SourceFile, expect: () -> String) {
-    val actual = createK2Compiler().compile(file)
-      .irModuleFragment.files
-      .single()
+    val actual = createK2Compiler().compile(file).irModuleFragment.files.single()
+    assertEquals(expect().trim(), with(GoldenUtil) { actual.dumpSrcForGolden(code = file.source) })
+  }
+
+  protected fun source(filename: String): SourceFile {
+    val resolvedFilename = sourceRoot?.let { "$it/$filename" } ?: filename
+    return SourceFile(
+      name = resolvedFilename.substringAfterLast('/'),
+      source = sourceString(resolvedFilename),
+      path = sourcePath(resolvedFilename).substringBeforeLast('/'),
+    )
+  }
+
+  companion object {
+    private fun File.applyExistenceCheck(): File = apply {
+      if (!exists()) throw NoSuchFileException(this)
+    }
+
+    private val homeDir: String = run {
+      val userDir = System.getProperty("user.dir")
+      val dir = File(userDir ?: ".")
+      val path = FileUtil.toCanonicalPath(dir.absolutePath)
+      File(path).applyExistenceCheck().absolutePath
+    }
+
+    private val NO_FEATURES = EnumSet.noneOf(FeatureFlag::class.java)
+
+    // https://github.com/JetBrains/kotlin/blob/bb25d2f8aa74406ff0af254b2388fd601525386a/plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/AbstractCompilerTest.kt#L212-L228
+    private val defaultClassPath by lazy {
+      listOf(
+        jarFor<Unit>(),
+        jarFor<kotlinx.coroutines.CoroutineScope>(),
+        jarFor<androidx.compose.runtime.Composable>(),
+        jarFor<androidx.compose.animation.EnterTransition>(),
+        jarFor<androidx.compose.ui.Modifier>(),
+        jarFor<androidx.compose.ui.graphics.ColorProducer>(),
+        jarFor<androidx.compose.ui.unit.Dp>(),
+        jarFor<androidx.compose.ui.text.input.TextFieldValue>(),
+        jarFor<androidx.compose.foundation.Indication>(),
+        jarFor<androidx.compose.foundation.text.KeyboardActions>(),
+        jarFor<androidx.compose.foundation.layout.RowScope>(),
+        jarFor<ComposeInvestigatorConfig>(),
+      )
+    }
+  }
+}
+
+private object GoldenUtil {
+  private val MatchResult.text get() = groupValues[0]
+  private fun MatchResult.number() = groupValues[1].toInt()
+  private fun MatchResult.isChar(c: String) = text == c
+  private fun MatchResult.isNumber() = groupValues[1].isNotEmpty()
+  private fun MatchResult.isFileName() = groups[4] != null
+
+  fun IrFile.dumpSrcForGolden(code: String): String {
+    val keySet = mutableListOf<Int>()
+    return this
       .dumpSrc(useFir = true)
-    assertEquals(expect().trim(), actual)
+      .replace('$', '%')
+      // replace source keys for start group calls
+      .replace(Regex("(%composer\\.start(Restart|Movable|Replaceable|Replace)Group\\()-?((0b)?[-\\d]+)")) { match ->
+        val stringKey = match.groupValues[3]
+        val key = if (stringKey.startsWith("0b")) Integer.parseInt(stringKey.drop(2), 2) else stringKey.toInt()
+        if (key in keySet) {
+          "${match.groupValues[1]}<!DUPLICATE KEY: $key!>"
+        } else {
+          keySet.add(key)
+          "${match.groupValues[1]}<>"
+        }
+      }
+      .replace(Regex("(sourceInformationMarkerStart\\(%composer, )([-\\d]+)")) { match ->
+        "${match.groupValues[1]}<>"
+      }
+      // replace traceEventStart values with a token
+      .replace(Regex("traceEventStart\\(-?\\d+, (%dirty|%changed|-1), (%dirty1|%changed1|-1), (.*)")) { match ->
+        "traceEventStart(<>, ${match.groupValues[1]}, ${match.groupValues[2]}, <>)"
+      }
+      // replace source information with source it references
+      .replace(Regex("(%composer\\.start(Restart|Movable|Replaceable|Replace)Group\\([^\"\\n]*)\"(.*)\"\\)")) { match ->
+        "${match.groupValues[1]}\"${generateSourceInfo(match.groupValues[4], code)}\")"
+      }
+      .replace(Regex("(sourceInformation(MarkerStart)?\\(.*)\"(.*)\"\\)")) { match ->
+        "${match.groupValues[1]}\"${generateSourceInfo(match.groupValues[3], code)}\")"
+      }
+      .replace(Regex("(composableLambda[N]?\\([^\"\\n]*)\"(.*)\"\\)")) { match ->
+        "${match.groupValues[1]}\"${generateSourceInfo(match.groupValues[2], code)}\")"
+      }
+      .replace(Regex("(rememberComposableLambda[N]?)\\((-?\\d+)")) { match ->
+        "${match.groupValues[1]}(<>"
+      }
+      // replace source keys for joinKey calls
+      .replace(Regex("(%composer\\.joinKey\\()([-\\d]+)")) { match ->
+        "${match.groupValues[1]}<>"
+      }
+      // composableLambdaInstance(<>, true)
+      .replace(Regex("(composableLambdaInstance\\()([-\\d]+, (true|false))")) { match ->
+        "${match.groupValues[1]}<>, ${match.groupValues[3]}"
+      }
+      // composableLambda(%composer, <>, true)
+      .replace(Regex("(composableLambda\\(%composer,\\s)([-\\d]+)")) { match ->
+        "${match.groupValues[1]}<>"
+      }
+      .trimIndent()
+      .trimTrailingWhitespaces()
+  }
+
+  @Suppress("RegExpSimplifiable")
+  private fun generateSourceInfo(sourceInfo: String, source: String): String {
+    val regex = Regex("(\\d+)|([,])|([*])|([:])|C(\\(.*\\))?|L|(P\\(*\\))|@")
+
+    var current = 0
+    var currentResult = regex.find(sourceInfo, current)
+    var result = ""
+
+    fun next(): MatchResult? {
+      currentResult?.let { match ->
+        current = match.range.last + 1
+        currentResult = match.next()
+      }
+      return currentResult
+    }
+
+    // A location has the format: [<line-number>]['@' <offset> ['L' <length>]]
+    // where the named productions are numbers
+    fun parseLocation(): String? {
+      var mr = currentResult
+      if (mr != null && mr.isNumber()) {
+        // line number, we ignore the value in during testing.
+        mr = next()
+      }
+      if (mr != null && mr.isChar("@")) {
+        // Offset
+        mr = next()
+        if (mr == null || !mr.isNumber()) {
+          return null
+        }
+        val offset = mr.number()
+        mr = next()
+        var ellipsis = ""
+        val maxFragment = 6
+        val rawLength = if (mr != null && mr.isChar("L")) {
+          mr = next()
+          if (mr == null || !mr.isNumber()) {
+            return null
+          }
+          mr.number().also { next() }
+        } else {
+          maxFragment
+        }
+        val eol = source.indexOf('\n', offset).let {
+          if (it < 0) source.length else it
+        }
+        val space = source.indexOf(' ', offset).let {
+          if (it < 0) source.length else it
+        }
+        val maxEnd = offset + maxFragment
+        if (eol > maxEnd && space > maxEnd) ellipsis = "..."
+        val length = minOf(maxEnd, minOf(offset + rawLength, space, eol)) - offset
+        return "<${source.substring(offset, offset + length)}$ellipsis>"
+      }
+      return null
+    }
+
+    while (currentResult != null) {
+      val mr = currentResult!!
+      if (mr.range.first != current) {
+        return "invalid source info at $current: '$sourceInfo'"
+      }
+      when {
+        mr.isNumber() || mr.isChar("@") -> {
+          val fragment = parseLocation() ?: return "invalid source info at $current: '$sourceInfo'"
+          result += fragment
+        }
+        mr.isFileName() -> {
+          return result + ":" + sourceInfo.substring(mr.range.last + 1)
+        }
+        else -> {
+          result += mr.text
+          next()
+        }
+      }
+      require(mr != currentResult) { "regex didn't advance" }
+    }
+
+    if (current != sourceInfo.length) return "invalid source info at $current: '$sourceInfo'"
+
+    return result
   }
+
+  private fun String.trimTrailingWhitespaces(): String =
+    split('\n').joinToString("\n", transform = String::trimEnd)
 }
 
 private inline fun <reified T> jarFor() = File(PathUtil.getJarPathForClass(T::class.java))

compiler/src/test/kotlin/land/sungbin/composeinvestigator/compiler/_source/utils.kt

@@ -8,14 +8,6 @@
 package land.sungbin.composeinvestigator.compiler._source
 
 import java.io.File
-import land.sungbin.composeinvestigator.compiler._compilation.SourceFile
-
-fun source(filename: String): SourceFile =
-  SourceFile(
-    name = filename.substringAfterLast('/'),
-    source = sourceString(filename),
-    path = sourcePath(filename).substringBeforeLast('/'),
-  )
 
 fun sourcePath(filename: String): String =
   "src/test/kotlin/land/sungbin/composeinvestigator/compiler/_source/$filename"

compiler/src/test/kotlin/land/sungbin/composeinvestigator/compiler/analysis/DurableComposableKeyAnalyzerTest.kt

@@ -12,16 +12,18 @@ import kotlin.test.Test
 import kotlin.test.assertEquals
 import land.sungbin.composeinvestigator.compiler.FeatureFlag
 import land.sungbin.composeinvestigator.compiler._compilation.AbstractCompilerTest
-import land.sungbin.composeinvestigator.compiler._source.source
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.expressions.IrConst
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
 import org.jetbrains.kotlin.utils.addToStdlib.enumSetOf
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
-class DurableComposableKeyAnalyzerTest : AbstractCompilerTest(enumSetOf(FeatureFlag.InvalidationProcessTracing)) {
+class DurableComposableKeyAnalyzerTest : AbstractCompilerTest(
+  enumSetOf(FeatureFlag.InvalidationProcessTracing),
+  sourceRoot = "analysis/durableComposableKey",
+) {
   @Test fun generates_a_unique_key_for_the_same_function_name() {
-    val result = compile(source("analysis/durableComposableKey/sameFunctionNames.kt"))
+    val result = compile(source("sameFunctionNames.kt"))
     val trace = result.pluginContext.irTrace
 
     val functions = result.irModuleFragment