Make StateInitializerTransform also handle StateObject

compiler-integration-test/build.gradle.kts

@@ -23,13 +23,14 @@ dependencies {
   implementation(projects.runtime)
   implementation(libs.compose.runtime)
 
+  implementation(kotlin("reflect")) // Used by assertk
+  implementation(libs.test.assertk)
+
   implementation("androidx.compose.runtime:runtime-test-utils:1.8.0-SNAPSHOT") {
     because("Why SNAPSHOT? See https://cs.android.com/androidx/platform/frameworks/support/+/androidx-main:compose/runtime/runtime-test-utils/build.gradle;l=71;drc=214c6abe4e624304956276717a0163fad3858be9")
   }
   kotlinCompilerPluginClasspath(projects.compiler)
 
   testImplementation(kotlin("test-junit5", version = libs.versions.kotlin.core.get()))
-  testImplementation(kotlin("reflect")) // Used by assertk
   testImplementation(libs.test.kotlin.coroutines)
-  testImplementation(libs.test.assertk)
 }

compiler-integration-test/src/main/kotlin/land/sungbin/composeinvestigator/compiler/test/StateObjects.kt

@@ -0,0 +1,59 @@
+// Copyright 2024 Ji Sungbin
+// SPDX-License-Identifier: Apache-2.0
+package land.sungbin.composeinvestigator.compiler.test
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableDoubleStateOf
+import androidx.compose.runtime.mutableFloatStateOf
+import androidx.compose.runtime.mutableIntStateOf
+import androidx.compose.runtime.mutableLongStateOf
+import androidx.compose.runtime.mutableStateListOf
+import androidx.compose.runtime.mutableStateMapOf
+import androidx.compose.runtime.mutableStateOf
+import assertk.assertAll
+import assertk.assertThat
+import assertk.assertions.isEqualTo
+import land.sungbin.composeinvestigator.runtime.currentComposableInvalidationTracer
+
+val stateObjectsTable by lazy { currentComposableInvalidationTracer }
+
+@Composable fun DirectStateObjects() {
+  val myUnitState = mutableStateOf(Unit)
+  val myIntState = mutableIntStateOf(0)
+  val myFloatState = mutableFloatStateOf(0f)
+  val myLongState = mutableLongStateOf(0L)
+  val myDoubleState = mutableDoubleStateOf(0.0)
+  val myMapState = mutableStateMapOf(0 to 0)
+  val myListState = mutableStateListOf(0)
+
+  assertAll {
+    assertThat(stateObjectsTable.findStateObjectName(myUnitState)).isEqualTo("myUnitState")
+    assertThat(stateObjectsTable.findStateObjectName(myIntState)).isEqualTo("myIntState")
+    assertThat(stateObjectsTable.findStateObjectName(myFloatState)).isEqualTo("myFloatState")
+    assertThat(stateObjectsTable.findStateObjectName(myLongState)).isEqualTo("myLongState")
+    assertThat(stateObjectsTable.findStateObjectName(myDoubleState)).isEqualTo("myDoubleState")
+    assertThat(stateObjectsTable.findStateObjectName(myMapState)).isEqualTo("myMapState")
+    assertThat(stateObjectsTable.findStateObjectName(myListState)).isEqualTo("myListState")
+  }
+}
+
+@Suppress("UnusedVariable", "unused")
+@Composable fun DelegateStateObjects() {
+  val objects = mutableListOf<Any>()
+
+  val myUnitState by mutableStateOf(Unit).also(objects::add)
+  val myIntState by mutableIntStateOf(0).also(objects::add)
+  val myFloatState by mutableFloatStateOf(0f).also(objects::add)
+  val myLongState by mutableLongStateOf(0L).also(objects::add)
+  val myDoubleState by mutableDoubleStateOf(0.0).also(objects::add)
+
+  @Suppress("KotlinConstantConditions") // false-positive (objects[N])
+  assertAll {
+    assertThat(stateObjectsTable.findStateObjectName(objects[0])).isEqualTo("myUnitState")
+    assertThat(stateObjectsTable.findStateObjectName(objects[1])).isEqualTo("myIntState")
+    assertThat(stateObjectsTable.findStateObjectName(objects[2])).isEqualTo("myFloatState")
+    assertThat(stateObjectsTable.findStateObjectName(objects[3])).isEqualTo("myLongState")
+    assertThat(stateObjectsTable.findStateObjectName(objects[4])).isEqualTo("myDoubleState")
+  }
+}

compiler-integration-test/src/test/kotlin/land/sungbin/composeinvestigator/compiler/test/StateObjectsTest.kt

@@ -0,0 +1,22 @@
+// Copyright 2024 Ji Sungbin
+// SPDX-License-Identifier: Apache-2.0
+package land.sungbin.composeinvestigator.compiler.test
+
+import androidx.compose.runtime.mock.compositionTest
+import kotlin.test.BeforeTest
+import org.junit.jupiter.api.Test
+
+class StateObjectsTest {
+  @BeforeTest fun prepare() {
+    TestConfiguration.reset()
+    stateObjectsTable.reset()
+  }
+
+  @Test fun directStateObjects() = compositionTest {
+    compose { DirectStateObjects() }
+  }
+
+  @Test fun delegateStateObjects() = compositionTest {
+    compose { DelegateStateObjects() }
+  }
+}

compiler/src/main/kotlin/land/sungbin/composeinvestigator/compiler/lower/ComposeInvestigatorBaseLower.kt

@@ -13,6 +13,7 @@ import land.sungbin.composeinvestigator.compiler.MUTABLE_LIST_OF_FQN
 import land.sungbin.composeinvestigator.compiler.NO_INVESTIGATION_FQN
 import land.sungbin.composeinvestigator.compiler.STABILITY_FQN
 import land.sungbin.composeinvestigator.compiler.STATE_FQN
+import land.sungbin.composeinvestigator.compiler.STATE_OBJECT_FQN
 import land.sungbin.composeinvestigator.compiler.Stability_CERTAIN
 import land.sungbin.composeinvestigator.compiler.Stability_COMBINED
 import land.sungbin.composeinvestigator.compiler.Stability_PARAMETER
@@ -87,6 +88,7 @@ public open class ComposeInvestigatorBaseLower(
   protected val composerCompoundKeyHashSymbol: IrSimpleFunctionSymbol = composerSymbol.getPropertyGetter(Composer_COMPOUND_KEY_HASH.toString())!!
 
   private val stateSymbol = context.referenceClass(ClassId.topLevel(STATE_FQN))!!
+  private val stateObjectSymbol = context.referenceClass(ClassId.topLevel(STATE_OBJECT_FQN))!!
 
   protected val mutableListOfSymbol: IrSimpleFunctionSymbol by unsafeLazy {
     context
@@ -211,10 +213,13 @@ public open class ComposeInvestigatorBaseLower(
   ): IrExpression = expression
 
   private fun IrVariable.isValidStateDeclaration(): Boolean {
-    val isState = type.classOrNull?.isSubtypeOfClass(stateSymbol.defaultType.classOrFail) == true
+    val isState = type.classOrNull?.let { clazz ->
+      clazz.isSubtypeOfClass(stateSymbol.defaultType.classOrFail) == true ||
+        clazz.isSubtypeOfClass(stateObjectSymbol.defaultType.classOrFail) == true
+    }
     val isTempVariable = origin == IrDeclarationOrigin.IR_TEMPORARY_VARIABLE
     val hasInitializer = initializer != null
-    return !type.isNullable() && isState && !isTempVariable && hasInitializer
+    return !type.isNullable() && isState == true && !isTempVariable && hasInitializer
   }
 
   protected fun Stability.asOwnStability(): IrConstructorCall {

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

@@ -34,6 +34,7 @@ public val Composer_SKIP_TO_GROUP_END: Name = identifier("skipToGroupEnd")
 public val Composer_COMPOUND_KEY_HASH: Name = identifier("compoundKeyHash")
 
 public val STATE_FQN: FqName = FqName("$AndroidxComposeRuntime.State")
+public val STATE_OBJECT_FQN: FqName = FqName("$AndroidxComposeRuntime.snapshots.StateObject")
 // END Compose Runtime
 
 // START ComposableScope