Add a test for the caching behavior of StabilityInferencer

Author: derekxu16

plugins/compose/compiler-hosted/integration-tests/src/jvmTest/kotlin/androidx/compose/compiler/plugins/kotlin/ClassStabilityTransformTests.kt

@@ -32,14 +32,17 @@ import org.jetbrains.kotlin.ir.builders.declarations.buildClass
 import org.jetbrains.kotlin.ir.builders.declarations.buildReceiverParameter
 import org.jetbrains.kotlin.ir.declarations.IrClass
 import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
 import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
 import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction
+import org.jetbrains.kotlin.ir.declarations.impl.IrFileImpl
 import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrReturn
 import org.jetbrains.kotlin.ir.symbols.UnsafeDuringIrConstructionAPI
 import org.jetbrains.kotlin.ir.types.IrType
 import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl
 import org.jetbrains.kotlin.ir.util.addChild
+import org.jetbrains.kotlin.ir.util.deepCopyWithoutPatchingParents
 import org.jetbrains.kotlin.ir.util.defaultType
 import org.jetbrains.kotlin.ir.util.file
 import org.jetbrains.kotlin.ir.util.statements
@@ -1872,6 +1875,54 @@ class ClassStabilityTransformTests(useFir: Boolean) : AbstractIrTransformTest(us
         )
     }
 
+    @Test
+    fun testStabilityCaching() {
+        val testFileSource = """
+            class A
+            class B(val a: A)
+            class C(val b: B)
+            class D(val c: C)
+        """.trimIndent()
+
+        val files = listOf(SourceFile("Test.kt", testFileSource))
+        val irModule = compileToIr(files, additionalPaths)
+        val stabilityInferencer = StabilityInferencer(isTargetJvm = true, irModule.descriptor, setOf())
+
+        val testIrFile = irModule.files.last()
+        val (aClass, bClass, cClass, dClass) = testIrFile.declarations.filterIsInstance<IrClass>()
+
+        // When a dependent in a different file requests the stability of a class defined in
+        // `testIrFile`, we expect the computed stability not to be cached. Note that
+        // `deepCopyWithoutPatchingParents` is used on the following line to obtain an object that
+        // is not referentially equal to `testIrFile`.
+        stabilityInferencer.stabilityOf(dClass.defaultType as IrType, testIrFile.deepCopyWithoutPatchingParents())
+        assertEquals(0, stabilityInferencer.cache.size)
+
+        // When `StabilityInferencer`'s cache is initially empty and `D`'s stability is requested,
+        // we expect the stability values of `A`, `B`, `C`, and `D` to each be computed and added to
+        // the cache.
+        val stabilityOfD = stabilityInferencer.stabilityOf(dClass.defaultType as IrType, testIrFile)
+        assertEquals(4, stabilityInferencer.cache.size)
+
+        // When the stability of a class is cached, we expect it to be returned from the cache, and
+        // we expect no cache modifications to happen.
+
+        assertEquals(
+            stabilityOfD,
+            stabilityInferencer.stabilityOf(dClass.defaultType as IrType, testIrFile)
+        )
+        assertEquals(4, stabilityInferencer.cache.size)
+
+        stabilityInferencer.stabilityOf(cClass.defaultType as IrType, testIrFile)
+        assertEquals(4, stabilityInferencer.cache.size)
+
+        stabilityInferencer.stabilityOf(bClass.defaultType as IrType, testIrFile)
+        assertEquals(4, stabilityInferencer.cache.size)
+
+        stabilityInferencer.stabilityOf(aClass.defaultType as IrType, testIrFile)
+        assertEquals(4, stabilityInferencer.cache.size)
+    }
+
     private fun assertStability(
         externalSrc: String,
         localSrc: String,

plugins/compose/compiler-hosted/src/main/java/androidx/compose/compiler/plugins/kotlin/analysis/Stability.kt

@@ -21,6 +21,7 @@ package androidx.compose.compiler.plugins.kotlin.analysis
 import androidx.compose.compiler.plugins.kotlin.ComposeFqNames
 import androidx.compose.compiler.plugins.kotlin.lower.annotationClass
 import androidx.compose.compiler.plugins.kotlin.lower.isSyntheticComposableFunction
+import com.google.common.annotations.VisibleForTesting
 import org.jetbrains.kotlin.backend.jvm.ir.isInlineClassType
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
@@ -192,7 +193,8 @@ private fun IrAnnotationContainer.stabilityParamBitmask(): Int? =
     (annotations.findAnnotation(ComposeFqNames.StabilityInferred)?.arguments[0] as? IrConst)
         ?.value as? Int
 
-private data class SymbolForAnalysis(
+@VisibleForTesting
+data class SymbolForAnalysis(
     val symbol: IrClassifierSymbol,
     val typeParameters: List<IrTypeArgument?>,
     /**
@@ -214,7 +216,9 @@ class StabilityInferencer(
 ) {
     private val externalTypeMatcherCollection = FqNameMatcherCollection(externalStableTypeMatchers)
 
-    private val cache = mutableMapOf<SymbolForAnalysis, Stability>()
+    @VisibleForTesting
+    var cache = mutableMapOf<SymbolForAnalysis, Stability>()
+        private set
 
     /**
      * Returns the stability of [irType].