Reimplementing InvalidationTraceTableApiChecker (#215)

jisungbin · Sep 30, 2024 · 47bc427 · 47bc427
1 parent 841569d
commit 47bc427
Show file tree

Hide file tree

Showing 17 changed files with 225 additions and 225 deletions.
diff --git a/...in/land/sungbin/composeinvestigator/compiler/frontend/InvalidationTraceTableApiChecker.kt b/...in/land/sungbin/composeinvestigator/compiler/frontend/InvalidationTraceTableApiChecker.kt
@@ -9,100 +9,92 @@ 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.COMPOSABLE_INVALIDATION_TRACE_TABLE_FQN
 import land.sungbin.composeinvestigator.compiler.COMPOSABLE_NAME_FQN
+import land.sungbin.composeinvestigator.compiler.COMPOSABLE_SCOPE_FQN
+import land.sungbin.composeinvestigator.compiler.CURRENT_COMPOSABLE_INVALIDATION_TRACER_FQN
 import land.sungbin.composeinvestigator.compiler.NO_INVESTIGATION_FQN
-import org.jetbrains.kotlin.KtSourceElement
-import org.jetbrains.kotlin.diagnostics.AbstractKtDiagnosticFactory
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
-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.FirFunctionChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.context.findClosest
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.ExpressionCheckers
 import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirFunctionCallChecker
+import org.jetbrains.kotlin.fir.analysis.checkers.expression.FirPropertyAccessExpressionChecker
 import org.jetbrains.kotlin.fir.analysis.extensions.FirAdditionalCheckersExtension
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.fir.declarations.hasAnnotation
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirLiteralExpression
 import org.jetbrains.kotlin.fir.expressions.FirPropertyAccessExpression
 import org.jetbrains.kotlin.fir.expressions.argument
+import org.jetbrains.kotlin.fir.references.toResolvedPropertySymbol
+import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
 import org.jetbrains.kotlin.fir.types.classId
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.fir.types.functionTypeKind
 import org.jetbrains.kotlin.fir.types.resolvedType
-import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitorVoid
 import org.jetbrains.kotlin.name.ClassId
+import org.jetbrains.kotlin.utils.addToStdlib.unreachableBranch
 
-// TODO add top-level usage checkers
 public class InvalidationTraceTableApiChecker(session: FirSession) : FirAdditionalCheckersExtension(session) {
-  override val declarationCheckers: DeclarationCheckers = object : DeclarationCheckers() {
-    override val functionCheckers = setOf(InvalidationTraceTableApiUsageCheck)
-  }
-
   override val expressionCheckers: ExpressionCheckers = object : ExpressionCheckers() {
+    override val propertyAccessExpressionCheckers = setOf(TraceTableApiAccessChecker)
     override val functionCallCheckers = setOf(ComposableNameExpressionChecker)
   }
 }
 
-private object InvalidationTraceTableApiUsageCheck : FirFunctionChecker(MppCheckerKind.Common) {
-  // TODO we need to fix the real reason for the duplicate diagnostics.
-  private val handled = mutableMapOf<KtSourceElement?, MutableSet<AbstractKtDiagnosticFactory>>()
+private object TraceTableApiAccessChecker : FirPropertyAccessExpressionChecker(MppCheckerKind.Common) {
+  private val NO_INVESTIGATION = ClassId.topLevel(NO_INVESTIGATION_FQN)
+  private val COMPOSABLE_SCOPE = ClassId.topLevel(COMPOSABLE_SCOPE_FQN)
 
-  override fun check(declaration: FirFunction, context: CheckerContext, reporter: DiagnosticReporter) {
-    val hasNoInvestigation = context.containingFile!!.hasAnnotation(ClassId.topLevel(NO_INVESTIGATION_FQN), context.session)
-    val isComposableScope = when (declaration) {
-      is FirAnonymousFunction -> declaration.typeRef.coneType.functionTypeKind(context.session) === ComposableFunction
-      else -> declaration.hasComposableAnnotation(context.session)
-    }
+  override fun check(expression: FirPropertyAccessExpression, context: CheckerContext, reporter: DiagnosticReporter) {
+    // TODO when accessed by `it`, such as `traceTable.let { it.action() }`,
+    //  Symbol is `FirValueParameterSymbol`. These variants need to be handled separately.
+    val callee = expression.calleeReference.toResolvedPropertySymbol() ?: return
+
+    if (
+      callee.callableId.asSingleFqName() == CURRENT_COMPOSABLE_INVALIDATION_TRACER_FQN &&
+      context.isNoInvestigationFile()
+    )
+      reporter.reportOn(expression.source, ComposeInvestigatorErrors.API_ACCESS_IN_NO_INVESTIGATION_FILE, context)
 
-    val visitor = TraceTableApiAccessVisiter(hasNoInvestigation, isComposableScope, context, reporter)
-    declaration.body?.statements?.fastForEach { statement -> statement.accept(visitor) }
+    checkComposableScopeCall(expression, callee, context, reporter)
   }
 
-  private class TraceTableApiAccessVisiter(
-    private val hasNoInvestigation: Boolean,
-    @Suppress("unused") private val isComposableScope: Boolean,
-    private val context: CheckerContext,
-    private val reporter: DiagnosticReporter,
-  ) : FirDefaultVisitorVoid() {
-    override fun visitElement(element: FirElement) {
-      element.acceptChildren(this)
-    }
+  private fun checkComposableScopeCall(
+    expression: FirPropertyAccessExpression,
+    callee: FirPropertySymbol,
+    context: CheckerContext,
+    reporter: DiagnosticReporter,
+  ) {
+    // FIXME the annotation for callee was not being looked up correctly.
+    //  I don't know the cause yet, so I'm temporarily disabling this check.
+    return
+
+    if (
+      !callee.hasAnnotation(COMPOSABLE_SCOPE, context.session) &&
+      callee.getterSymbol?.hasAnnotation(COMPOSABLE_SCOPE, context.session) != true &&
+      callee.setterSymbol?.hasAnnotation(COMPOSABLE_SCOPE, context.session) != true
+    )
+      return
 
-    override fun visitPropertyAccessExpression(expression: FirPropertyAccessExpression) {
-      if (
-        (expression.dispatchReceiver?.resolvedType ?: expression.extensionReceiver?.resolvedType ?: expression.resolvedType)
-          .classId?.asSingleFqName() != COMPOSABLE_INVALIDATION_TRACE_TABLE_FQN
-      )
-        return
+    if (!context.isComposableScope())
+      reporter.reportOn(expression.source, ComposeInvestigatorErrors.ILLEGAL_COMPOSABLE_SCOPE_CALL, context)
+  }
 
-      if (hasNoInvestigation) {
-        if (handled.getOrPut(expression.source, ::mutableSetOf).add(ComposeInvestigatorErrors.API_ACCESS_IN_NO_INVESTIGATION_FILE))
-          reporter.reportOn(expression.source, ComposeInvestigatorErrors.API_ACCESS_IN_NO_INVESTIGATION_FILE, context)
-      }
+  private fun CheckerContext.isNoInvestigationFile() =
+    containingFile!!.hasAnnotation(NO_INVESTIGATION, session)
 
-// TODO need to rewrite logic to check if current function is in Composable scope.
-//      if (
-//        expression.calleeReference
-//          .toResolvedPropertySymbol()
-//          ?.hasAnnotation(ClassId.topLevel(COMPOSABLE_SCOPE_FQN), context.session)
-//        != true
-//      )
-//        return
-//
-//      if (!isComposableScope) {
-//        if (handled.getOrPut(expression.source, ::mutableSetOf).add(ComposeInvestigatorErrors.ILLEGAL_COMPOSABLE_SCOPE_CALL))
-//          reporter.reportOn(expression.source, ComposeInvestigatorErrors.ILLEGAL_COMPOSABLE_SCOPE_CALL, context)
-//      }
+  fun CheckerContext.isComposableScope(): Boolean =
+    when (val declaration = findClosest<FirFunction>() ?: return false) {
+      is FirAnonymousFunction -> declaration.typeRef.coneType.functionTypeKind(session) === ComposableFunction
+      is FirSimpleFunction -> declaration.hasComposableAnnotation(session)
+      else -> unreachableBranch(declaration::class.simpleName)
     }
-  }
 }
 
 private object ComposableNameExpressionChecker : FirFunctionCallChecker(MppCheckerKind.Common) {

diff --git a/.../sungbin/composeinvestigator/compiler/lower/InvalidationTraceTableIntrinsicTransformer.kt b/.../sungbin/composeinvestigator/compiler/lower/InvalidationTraceTableIntrinsicTransformer.kt
@@ -62,6 +62,8 @@ public class InvalidationTraceTableIntrinsicTransformer(
     tableSymbol.getPropertyGetter(ComposableInvalidationTraceTable_CURRENT_COMPOSABLE_KEY_NAME.asString())!!.owner
 
   override fun visitCall(expression: IrCall): IrExpression {
+    // TODO generating `throw Exception(NO_TABLE)` code to the target IR instead
+    //  of throwing it here if no table exists
     val table = tables[currentFile]
     return when (expression.symbol.owner.kotlinFqName) {
       currentTableGetterSymbol.kotlinFqName -> {

diff --git a/...est/kotlin/land/sungbin/composeinvestigator/compiler/_assert/DiagnosticsResultAsserter.kt b/...est/kotlin/land/sungbin/composeinvestigator/compiler/_assert/DiagnosticsResultAsserter.kt
@@ -13,9 +13,14 @@ import kotlin.test.fail
 import land.sungbin.composeinvestigator.compiler._compilation.FirAnalysisResult
 import org.jetbrains.kotlin.diagnostics.AbstractKtDiagnosticFactory
 
+private const val SEPARATOR = "\n=====\n"
+
 fun FirAnalysisResult.assertNoDiagnostic(diagnostic: AbstractKtDiagnosticFactory) {
   val results = diagnostics.getOrElse(diagnostic.name, ::emptyList)
-  assertTrue(results.isEmpty(), "Expected no diagnostic but found diagnostic(s).")
+  assertTrue(
+    results.isEmpty(),
+    "Expected no diagnostic but found diagnostics:\n${results.joinToString(SEPARATOR)}",
+  )
 }
 
 fun FirAnalysisResult.assertDiagnostics(
@@ -26,6 +31,6 @@ fun FirAnalysisResult.assertDiagnostics(
   if (results.isEmpty())
     fail("Expected diagnostics message but no diagnostic was found.")
 
-  val actualMessages = results.joinToString("\n=====\n")
+  val actualMessages = results.joinToString(SEPARATOR)
   assertEquals(expectMessages().trim(), actualMessages)
 }
diff --git a/...in/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/composableFunction.kt b/...in/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/composableFunction.kt
@@ -14,9 +14,8 @@ import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
 @Composable private fun composableFunction() {
-  with(currentComposableInvalidationTracer) {
-    currentComposableName
-    currentComposableName = ComposableName("")
-    currentComposableKeyName
-  }
+  val t = currentComposableInvalidationTracer
+  t.currentComposableName
+  t.currentComposableName = ComposableName("")
+  currentComposableInvalidationTracer.currentComposableKeyName
 }
diff --git a/...gbin/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/composableLambda.kt b/...gbin/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/composableLambda.kt
@@ -10,17 +10,20 @@
 package land.sungbin.composeinvestigator.compiler._source.frontend.traceTableApiUsage
 
 import androidx.compose.runtime.Composable
+import land.sungbin.composeinvestigator.runtime.ComposableInvalidationTraceTable
 import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
-private fun composableLambda() {
-  l {
-    with(currentComposableInvalidationTracer) {
-      currentComposableName
-      currentComposableName = ComposableName("")
-      currentComposableKeyName
-    }
+@Composable private fun composableLambda() {
+  val t = currentComposableInvalidationTracer
+  l(t) {
+    currentComposableName
+    it.currentComposableKeyName
+    currentComposableInvalidationTracer.currentComposableName = ComposableName("")
   }
 }
 
-private fun l(l: @Composable () -> Unit) = Unit
+private fun l(
+  t: ComposableInvalidationTraceTable,
+  b: @Composable ComposableInvalidationTraceTable.(ComposableInvalidationTraceTable) -> Unit,
+) = Unit
diff --git a/...poseinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineComposableFunction.kt b/...poseinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineComposableFunction.kt
@@ -14,9 +14,8 @@ import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
 @Composable private inline fun inlineComposableFunction() {
-  with(currentComposableInvalidationTracer) {
-    currentComposableName
-    currentComposableName = ComposableName("")
-    currentComposableKeyName
-  }
+  val t = currentComposableInvalidationTracer
+  t.currentComposableName
+  t.currentComposableName = ComposableName("")
+  currentComposableInvalidationTracer.currentComposableKeyName
 }
diff --git a/...omposeinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineComposableLambda.kt b/...omposeinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineComposableLambda.kt
@@ -10,17 +10,20 @@
 package land.sungbin.composeinvestigator.compiler._source.frontend.traceTableApiUsage
 
 import androidx.compose.runtime.Composable
+import land.sungbin.composeinvestigator.runtime.ComposableInvalidationTraceTable
 import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
-private fun inlineComposableLambda() {
-  l {
-    with(currentComposableInvalidationTracer) {
-      currentComposableName
-      currentComposableName = ComposableName("")
-      currentComposableKeyName
-    }
+@Composable private inline fun inlineComposableLambda() {
+  val t = currentComposableInvalidationTracer
+  l(t) {
+    currentComposableName
+    it.currentComposableKeyName
+    currentComposableInvalidationTracer.currentComposableName = ComposableName("")
   }
 }
 
-private inline fun l(b: @Composable () -> Unit) = Unit
+private inline fun l(
+  t: ComposableInvalidationTraceTable,
+  b: @Composable ComposableInvalidationTraceTable.(ComposableInvalidationTraceTable) -> Unit,
+) = Unit
diff --git a/.../composeinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineNormalFunction.kt b/.../composeinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineNormalFunction.kt
@@ -13,9 +13,8 @@ import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
 private inline fun inlineNormalFunction() {
-  with(currentComposableInvalidationTracer) {
-    currentComposableName
-    currentComposableName = ComposableName("")
-    currentComposableKeyName
-  }
+  val t = currentComposableInvalidationTracer
+  t.currentComposableName
+  t.currentComposableName = ComposableName("")
+  currentComposableInvalidationTracer.currentComposableKeyName
 }
diff --git a/...in/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineNormalLambda.kt b/...in/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/inlineNormalLambda.kt
@@ -9,17 +9,20 @@
 
 package land.sungbin.composeinvestigator.compiler._source.frontend.traceTableApiUsage
 
+import land.sungbin.composeinvestigator.runtime.ComposableInvalidationTraceTable
 import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
-private fun inlineNormalLambda() {
-  l {
-    with(currentComposableInvalidationTracer) {
-      currentComposableName
-      currentComposableName = ComposableName("")
-      currentComposableKeyName
-    }
+private inline fun inlineNormalLambda() {
+  val t = currentComposableInvalidationTracer
+  l(t) {
+    currentComposableName
+    it.currentComposableKeyName
+    currentComposableInvalidationTracer.currentComposableName = ComposableName("")
   }
 }
 
-private inline fun l(b: () -> Unit) = Unit
+private inline fun l(
+  t: ComposableInvalidationTraceTable,
+  b: ComposableInvalidationTraceTable.(ComposableInvalidationTraceTable) -> Unit,
+) = Unit
diff --git a/...ungbin/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/normalFunction.kt b/...ungbin/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/normalFunction.kt
@@ -13,9 +13,8 @@ import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
 private fun normalFunction() {
-  with(currentComposableInvalidationTracer) {
-    currentComposableName
-    currentComposableName = ComposableName("")
-    currentComposableKeyName
-  }
+  val t = currentComposableInvalidationTracer
+  t.currentComposableName
+  t.currentComposableName = ComposableName("")
+  currentComposableInvalidationTracer.currentComposableKeyName
 }
diff --git a/.../sungbin/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/normalLambda.kt b/.../sungbin/composeinvestigator/compiler/_source/frontend/traceTableApiUsage/normalLambda.kt
@@ -9,17 +9,20 @@
 
 package land.sungbin.composeinvestigator.compiler._source.frontend.traceTableApiUsage
 
+import land.sungbin.composeinvestigator.runtime.ComposableInvalidationTraceTable
 import land.sungbin.composeinvestigator.runtime.ComposableName
 import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
 
 private fun normalLambda() {
-  l {
-    with(currentComposableInvalidationTracer) {
-      currentComposableName
-      currentComposableName = ComposableName("")
-      currentComposableKeyName
-    }
+  val t = currentComposableInvalidationTracer
+  l(t) {
+    currentComposableName
+    it.currentComposableKeyName
+    currentComposableInvalidationTracer.currentComposableName = ComposableName("")
   }
 }
 
-private fun l(b: () -> Unit) = Unit
+private fun l(
+  t: ComposableInvalidationTraceTable,
+  b: ComposableInvalidationTraceTable.(ComposableInvalidationTraceTable) -> Unit,
+) = Unit
diff --git a/...r/_source/frontend/traceTableApiUsage/topLevelApiCallWithAllNoInvestigationComposables.kt b/...r/_source/frontend/traceTableApiUsage/topLevelApiCallWithAllNoInvestigationComposables.kt
@@ -0,0 +1,24 @@
+/*
+ * Developed by Ji Sungbin 2024.
+ *
+ * Licensed under the MIT.
+ * Please see full license: https://github.com/jisungbin/ComposeInvestigator/blob/main/LICENSE
+ */
+
+@file:Suppress("unused")
+
+package land.sungbin.composeinvestigator.compiler._source.frontend.traceTableApiUsage
+
+import androidx.compose.runtime.Composable
+import land.sungbin.composeinvestigator.runtime.NoInvestigation
+import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
+
+private val t = currentComposableInvalidationTracer
+private val a = Unit.run {
+  t.currentComposableName
+  t.currentComposableKeyName
+}
+
+@Composable @NoInvestigation private fun C() {}
+@Composable @NoInvestigation private fun C2() {}
+@Composable @NoInvestigation private fun C3() {}