Implement snapshot changes analysis by looking up the slot table (#147)

Author: skydoves

app/src/main/kotlin/com/skydoves/myapplication/ExampleRecompositionTracking.kt

@@ -46,6 +46,7 @@ import com.skydoves.myapplication.models.UnstableUser
  *    ComposeStabilityAnalyzer.setEnabled(BuildConfig.DEBUG)
  * 4. Run your app and check Logcat for "Recomposition" tag
  */
+@TraceRecomposition(traceStates = true)
 @Composable
 fun RecompositionTrackingExample() {
   var counter by remember { mutableIntStateOf(0) }

gradle.properties

@@ -51,7 +51,7 @@ kotlin.mpp.androidGradlePluginCompatibility.nowarn=true
 
 # Maven publishing
 GROUP=com.github.skydoves
-VERSION_NAME=0.7.2
+VERSION_NAME=0.7.3-SNAPSHOT
 
 POM_URL=https://github.com/skydoves/compose-stability-analyzer/
 POM_SCM_URL=https://github.com/skydoves/compose-stability-analyzer/

gradle/libs.versions.toml

@@ -13,7 +13,7 @@ androidGradlePlugin = "8.13.1"
 androidxActivity = "1.11.0"
 androidxComposeBom = "2025.12.00"
 jetbrains-compose = "1.9.3"
-compose-stability-analyzer = "0.7.0"
+compose-stability-analyzer = "0.7.2"
 runtimeAnnotation = "1.9.0"
 spotless = "6.21.0"
 shadow = "8.1.1"

stability-compiler/api/stability-compiler.api

@@ -156,7 +156,8 @@ public final class com/skydoves/compose/stability/compiler/StabilityReport$Compa
 public final class com/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder {
 	public fun <init> (Lorg/jetbrains/kotlin/backend/common/extensions/IrPluginContext;)V
 	public final fun initializeSymbols ()Z
-	public final fun injectTrackingCode (Lorg/jetbrains/kotlin/ir/declarations/IrFunction;Ljava/lang/String;Ljava/lang/String;ILjava/util/List;)Z
+	public final fun injectTrackingCode (Lorg/jetbrains/kotlin/ir/declarations/IrFunction;Ljava/lang/String;Ljava/lang/String;ILjava/util/List;Ljava/util/List;)Z
+	public static synthetic fun injectTrackingCode$default (Lcom/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder;Lorg/jetbrains/kotlin/ir/declarations/IrFunction;Ljava/lang/String;Ljava/lang/String;ILjava/util/List;Ljava/util/List;ILjava/lang/Object;)Z
 }
 
 public final class com/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder$ParameterStabilityData {
@@ -176,6 +177,26 @@ public final class com/skydoves/compose/stability/compiler/lower/RecompositionIr
 	public fun toString ()Ljava/lang/String;
 }
 
+public final class com/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder$StateVariableData {
+	public fun <init> (Ljava/lang/String;Ljava/lang/String;Lorg/jetbrains/kotlin/ir/declarations/IrVariable;ZLorg/jetbrains/kotlin/ir/declarations/IrSimpleFunction;)V
+	public synthetic fun <init> (Ljava/lang/String;Ljava/lang/String;Lorg/jetbrains/kotlin/ir/declarations/IrVariable;ZLorg/jetbrains/kotlin/ir/declarations/IrSimpleFunction;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun component1 ()Ljava/lang/String;
+	public final fun component2 ()Ljava/lang/String;
+	public final fun component3 ()Lorg/jetbrains/kotlin/ir/declarations/IrVariable;
+	public final fun component4 ()Z
+	public final fun component5 ()Lorg/jetbrains/kotlin/ir/declarations/IrSimpleFunction;
+	public final fun copy (Ljava/lang/String;Ljava/lang/String;Lorg/jetbrains/kotlin/ir/declarations/IrVariable;ZLorg/jetbrains/kotlin/ir/declarations/IrSimpleFunction;)Lcom/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder$StateVariableData;
+	public static synthetic fun copy$default (Lcom/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder$StateVariableData;Ljava/lang/String;Ljava/lang/String;Lorg/jetbrains/kotlin/ir/declarations/IrVariable;ZLorg/jetbrains/kotlin/ir/declarations/IrSimpleFunction;ILjava/lang/Object;)Lcom/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder$StateVariableData;
+	public fun equals (Ljava/lang/Object;)Z
+	public final fun getGetter ()Lorg/jetbrains/kotlin/ir/declarations/IrSimpleFunction;
+	public final fun getName ()Ljava/lang/String;
+	public final fun getTypeString ()Ljava/lang/String;
+	public final fun getVariable ()Lorg/jetbrains/kotlin/ir/declarations/IrVariable;
+	public fun hashCode ()I
+	public final fun isDelegated ()Z
+	public fun toString ()Ljava/lang/String;
+}
+
 public final class com/skydoves/compose/stability/compiler/lower/StabilityAnalyzerTransformer : org/jetbrains/kotlin/backend/common/IrElementTransformerVoidWithContext {
 	public static final field Companion Lcom/skydoves/compose/stability/compiler/lower/StabilityAnalyzerTransformer$Companion;
 	public fun <init> (Lorg/jetbrains/kotlin/backend/common/extensions/IrPluginContext;Lcom/skydoves/compose/stability/compiler/StabilityInfoCollector;Ljava/util/List;)V

stability-compiler/src/main/kotlin/com/skydoves/compose/stability/compiler/lower/RecompositionIrBuilder.kt

@@ -27,10 +27,14 @@ import org.jetbrains.kotlin.ir.builders.irGet
 import org.jetbrains.kotlin.ir.builders.irInt
 import org.jetbrains.kotlin.ir.builders.irString
 import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.declarations.IrLocalDelegatedProperty
+import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
 import org.jetbrains.kotlin.ir.declarations.IrValueParameter
 import org.jetbrains.kotlin.ir.declarations.IrVariable
+import org.jetbrains.kotlin.ir.expressions.IrBlock
 import org.jetbrains.kotlin.ir.expressions.IrBlockBody
 import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrWhen
 import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
@@ -64,6 +68,7 @@ public class RecompositionIrBuilder(
   private var trackerClassSymbol: IrClassSymbol? = null
   private var rememberTrackerFunctionSymbol: IrSimpleFunctionSymbol? = null
   private var trackParameterFunctionSymbol: IrSimpleFunctionSymbol? = null
+  private var trackStateFunctionSymbol: IrSimpleFunctionSymbol? = null
   private var logIfThresholdMetFunctionSymbol: IrSimpleFunctionSymbol? = null
 
   /**
@@ -104,6 +109,12 @@ public class RecompositionIrBuilder(
         return false
       }
 
+      trackStateFunctionSymbol = trackerClass?.functions?.firstOrNull {
+        it.name.asString() == "trackState"
+      }?.symbol
+
+      // trackState is optional - don't fail if missing (backward compat)
+
       logIfThresholdMetFunctionSymbol = trackerClass?.functions?.firstOrNull {
         it.name.asString() == "logIfThresholdMet"
       }?.symbol
@@ -136,52 +147,154 @@ public class RecompositionIrBuilder(
     tag: String,
     threshold: Int,
     parameterStabilities: List<ParameterStabilityData>,
+    stateVariables: List<StateVariableData> = emptyList(),
   ): Boolean {
     try {
       val body = function.body as? IrBlockBody ?: return false
 
       val builder = DeclarationIrBuilder(context, function.symbol)
 
-      // Create new body with tracking code prepended
-      val newStatements = mutableListOf<org.jetbrains.kotlin.ir.IrStatement>()
-
-      // 1. Create tracker variable: val _tracker = remember { ... }
+      // 1. Create tracker variable: val _tracker = rememberRecompositionTracker(...)
       val trackerVariable = createTrackerVariable(
         builder = builder,
         function = function,
         functionName = functionName,
         tag = tag,
         threshold = threshold,
       )
-      newStatements.add(trackerVariable)
 
-      // 2. Add trackParameter calls for each parameter
-      parameterStabilities.forEach { paramData ->
-        val trackParameterCall = createTrackParameterCall(
+      // 2. Create trackParameter calls for each parameter
+      val trackParameterCalls = parameterStabilities.map { paramData ->
+        createTrackParameterCall(
           builder = builder,
           trackerVariable = trackerVariable,
           paramData = paramData,
         )
-        newStatements.add(trackParameterCall)
       }
 
-      // 3. Add logIfThresholdMet call
+      // 3. Create logIfThresholdMet call
       val logCall = createLogIfThresholdMetCall(builder, trackerVariable)
-      newStatements.add(logCall)
 
-      // 4. Add original body statements
-      newStatements.addAll(body.statements)
+      if (stateVariables.isEmpty()) {
+        // Standard mode: prepend tracker + params + log, then body
+        val newStatements = mutableListOf<org.jetbrains.kotlin.ir.IrStatement>()
+        newStatements.add(trackerVariable)
+        newStatements.addAll(trackParameterCalls)
+        newStatements.add(logCall)
+        newStatements.addAll(body.statements)
+
+        body.statements.clear()
+        body.statements.addAll(newStatements)
+      } else {
+        // State tracking mode: interleave trackState calls after
+        // state variable declarations (may be nested in IrBlocks
+        // due to Compose compiler lowering), and move
+        // logIfThresholdMet to the end
+        val stateVarSet = stateVariables.associateBy {
+          it.variable
+        }
+
+        injectTrackStateRecursive(
+          body.statements,
+          builder,
+          trackerVariable,
+          stateVarSet,
+        )
 
-      // Replace body with new statements
-      body.statements.clear()
-      body.statements.addAll(newStatements)
+        val newStatements =
+          mutableListOf<org.jetbrains.kotlin.ir.IrStatement>()
+        newStatements.add(trackerVariable)
+        newStatements.addAll(trackParameterCalls)
+        newStatements.addAll(body.statements)
+        newStatements.add(logCall)
+
+        body.statements.clear()
+        body.statements.addAll(newStatements)
+      }
 
       return true
     } catch (e: Exception) {
       return false
     }
   }
 
+  /**
+   * Recursively walks IR statements and inserts trackState calls
+   * after detected state variable declarations. Handles IrBlock
+   * nesting from Compose compiler lowering.
+   */
+  private fun injectTrackStateRecursive(
+    statements: MutableList<org.jetbrains.kotlin.ir.IrStatement>,
+    builder: IrBuilderWithScope,
+    trackerVariable: IrVariable,
+    stateVarSet: Map<IrVariable, StateVariableData>,
+  ) {
+    var i = 0
+    while (i < statements.size) {
+      val stmt = statements[i]
+
+      // Recurse into blocks
+      if (stmt is IrBlock) {
+        injectTrackStateRecursive(
+          stmt.statements,
+          builder,
+          trackerVariable,
+          stateVarSet,
+        )
+      }
+
+      // Recurse into when branches
+      if (stmt is IrWhen) {
+        for (branch in stmt.branches) {
+          val branchResult = branch.result
+          if (branchResult is IrBlock) {
+            injectTrackStateRecursive(
+              branchResult.statements,
+              builder,
+              trackerVariable,
+              stateVarSet,
+            )
+          }
+        }
+      }
+
+      // Check if this is a detected state variable (IrVariable)
+      val variable = stmt as? IrVariable
+      if (variable != null && variable in stateVarSet) {
+        val stateData = stateVarSet[variable]!!
+        val trackStateCall = createTrackStateCall(
+          builder,
+          trackerVariable,
+          stateData,
+        )
+        if (trackStateCall != null) {
+          statements.add(i + 1, trackStateCall)
+          i++
+        }
+      }
+
+      // Check IrLocalDelegatedProperty (delegate is our key)
+      val delegatedProp = stmt as? IrLocalDelegatedProperty
+      if (delegatedProp != null) {
+        val delegate = delegatedProp.delegate
+        if (delegate in stateVarSet) {
+          val stateData = stateVarSet[delegate]!!
+          val trackStateCall = createTrackStateCall(
+            builder,
+            trackerVariable,
+            stateData,
+          )
+          if (trackStateCall != null) {
+            statements.add(i + 1, trackStateCall)
+            i++
+          }
+        }
+      }
+
+      i++
+    }
+  }
+
   /**
    * Creates an IrVariable for the tracker.
    */
@@ -275,6 +388,38 @@ public class RecompositionIrBuilder(
     return call
   }
 
+  /**
+   * Creates IR call to `tracker.trackState(name, type, value)`.
+   *
+   * For delegated state variables (var x by ...), reads the variable directly (already unwrapped).
+   * For non-delegated state variables (val s = mutableStateOf()),
+   * reads s.value via property getter.
+   */
+  private fun createTrackStateCall(
+    builder: IrBuilderWithScope,
+    trackerVariable: IrVariable,
+    stateData: StateVariableData,
+  ): IrExpression? {
+    val trackStateSymbol = trackStateFunctionSymbol ?: return null
+
+    val call = builder.irCall(trackStateSymbol)
+    call.arguments[0] = builder.irGet(trackerVariable) // dispatch receiver
+    call.arguments[1] = builder.irString(stateData.name) // name: String
+    call.arguments[2] = builder.irString(stateData.typeString) // type: String
+
+    if (stateData.getter != null) {
+      // Use the getter from IrLocalDelegatedProperty to read
+      // the unwrapped value (calls getValue on the delegate)
+      val getterCall = builder.irCall(stateData.getter.symbol)
+      call.arguments[3] = getterCall
+    } else {
+      // Fallback: read the variable directly
+      call.arguments[3] = builder.irGet(stateData.variable)
+    }
+
+    return call
+  }
+
   /**
    * Data class holding parameter information for tracking.
    */
@@ -284,4 +429,15 @@ public class RecompositionIrBuilder(
     val parameter: IrValueParameter,
     val stability: ParameterStability,
   )
+
+  /**
+   * Data class holding state variable information for tracking.
+   */
+  public data class StateVariableData(
+    val name: String,
+    val typeString: String,
+    val variable: IrVariable,
+    val isDelegated: Boolean,
+    val getter: IrSimpleFunction? = null,
+  )
 }