Rewrite NoComposableFileChecker logic

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

@@ -7,28 +7,29 @@
 
 package land.sungbin.composeinvestigator.compiler.frontend
 
-import androidx.compose.compiler.plugins.kotlin.k2.ComposableFunction
 import androidx.compose.compiler.plugins.kotlin.k2.hasComposableAnnotation
+import androidx.compose.compiler.plugins.kotlin.lower.fastForEach
 import land.sungbin.composeinvestigator.compiler.NO_INVESTIGATION_FQN
 import land.sungbin.composeinvestigator.compiler.lower.unsafeLazy
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationUseSiteTarget
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
+import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
 import org.jetbrains.kotlin.fir.analysis.checkers.MppCheckerKind
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.DeclarationCheckers
 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.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.smartPlus
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
-import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.constructClassLikeType
-import org.jetbrains.kotlin.fir.types.functionTypeKind
+import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
 
 public class InvalidationTraceTableInstantiationValidator(session: FirSession) : FirAdditionalCheckersExtension(session) {
@@ -45,27 +46,38 @@ private object NoComposableFileChecker : FirFileChecker(MppCheckerKind.Common) {
   }
 
   override fun check(declaration: FirFile, context: CheckerContext, reporter: DiagnosticReporter) {
-    if (
-      declaration.declarations.none { element ->
-        element.hasComposableAnnotation(context.session) ||
-          (
-            element is FirFunction &&
-              element.valueParameters.any { parameter ->
-                parameter.returnTypeRef.coneType.functionTypeKind(context.session) === ComposableFunction
-              }
-            )
+    var hasComposable = false
+
+    val composableCallVisitor = object : FirDefaultVisitorVoid() {
+      override fun visitElement(element: FirElement) {
+        if (hasComposable) return
+        element.acceptChildren(this)
       }
-    ) {
-      val noInvestigationAnnotation = buildAnnotation {
-        useSiteTarget = AnnotationUseSiteTarget.FILE
-        annotationTypeRef = noInvestigationType
-        argumentMapping = FirEmptyAnnotationArgumentMapping
+
+      override fun visitFunctionCall(functionCall: FirFunctionCall) {
+        if (functionCall.calleeReference.toResolvedFunctionSymbol()!!.hasComposableAnnotation(context.session))
+          hasComposable = true
+
+        if (hasComposable) return
+        super.visitFunctionCall(functionCall)
       }
+    }
 
-      declaration.replaceAnnotations(declaration.annotations.smartPlus(listOf(noInvestigationAnnotation)))
+    declaration.declarations.fastForEach { element ->
+      element.acceptChildren(composableCallVisitor)
+    }
+
+    if (hasComposable) return
 
-      // Validate that the new annotations do not break the file.
-      declaration.validate()
+    val noInvestigationAnnotation = buildAnnotation {
+      useSiteTarget = AnnotationUseSiteTarget.FILE
+      annotationTypeRef = noInvestigationType
+      argumentMapping = FirEmptyAnnotationArgumentMapping
     }
+
+    declaration.replaceAnnotations(declaration.annotations.smartPlus(listOf(noInvestigationAnnotation)))
+
+    // Validate that the new annotations do not break the file.
+    declaration.validate()
   }
 }

runtime/src/main/kotlin/land/sungbin/composeinvestigator/runtime/ComposableInvalidationTraceTable.kt

@@ -98,9 +98,8 @@ public class ComposableInvalidationTraceTable @ComposeInvestigatorCompilerApi pu
    * [Compose compiler's implementation](https://github.com/JetBrains/kotlin/blob/ede0373c4e5c0506b1491c6eb4c8bc0660ef7d21/plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/lower/DurableKeyTransformer.kt#L74)
    * of the AOSP.
    */
-  @property:ComposableScope
   public val currentComposableKeyName: String
-    @[Stable JvmSynthetic] get() = throw IntrinsicImplementedError()
+    @[Stable ComposableScope JvmSynthetic] get() = throw IntrinsicImplementedError()
 
   /**
    * Returns all arguments that were affected by the value change. This is useful for debugging and

runtime/src/main/kotlin/land/sungbin/composeinvestigator/runtime/ComposableScope.kt

@@ -10,5 +10,5 @@ package land.sungbin.composeinvestigator.runtime
 /** APIs annotated with this should only be used within Composable functions. */
 @MustBeDocumented
 @Retention(AnnotationRetention.SOURCE)
-@Target(AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY)
+@Target(AnnotationTarget.FUNCTION, AnnotationTarget.PROPERTY, AnnotationTarget.PROPERTY_GETTER)
 public annotation class ComposableScope