implement lazy caching for continuations

CLAUDE.md

@@ -121,7 +121,11 @@ Custom `_plugins/bytecode-processor` Gradle plugin applies compile-time transfor
 
 **`SpecCache.specMethod` benign race**: Written without synchronization (two threads may both compute and write the same handle), which is benign since MethodHandle instances are immutable and the computation is pure.
 
-**`ManualContinuation.$decoroutinator$cacheField` is nullable at runtime**: The bytecode transformer (`tryAddManualContinuation`) only runs the `PUTSTATIC` for `$decoroutinator$cacheField` in `<clinit>` when `fillUnknownElementsWithClassName()` is true at class-load time. When false, the JVM reference field stays null (default). The interface property is declared `SpecCache?` and the getter uses `?: return null` to handle this. Any code path that reads `$decoroutinator$cacheField` must handle null even though the Kotlin type looks non-nullable from the source.
+**`ManualContinuation.$decoroutinator$cacheField` can be JVM-null despite its non-null Kotlin type**: The bytecode transformer (`tryAddManualContinuation`) only runs the `PUTSTATIC` for the (`static`) `$decoroutinator$cacheField` in `<clinit>` when `fillUnknownElementsWithClassName` is `true` at class-load time; when false, the raw JVM field stays at its default `null`, even though the interface declares it `val $decoroutinator$cacheField: SpecCache` (non-null). This is only sound because `ManualContinuation.$decoroutinator$cache`'s getter checks `provider.fillUnknownElementsWithClassName` itself and returns `provider.nullElementSpecCache` *before* ever touching `$decoroutinator$cacheField` — since that flag is a process-lifetime-constant `val`, the `<clinit>`-time gate and the getter's runtime gate always agree. Never read `$decoroutinator$cacheField` directly without first replicating that same flag check.
+
+**`LazilyCachedContinuation` vs `ManualContinuation`**: Both bytecode-inject a `ContinuationCached` implementation into external (non-project) classes via `class-transformer/.../classTransformer.kt`, but for different situations. `ManualContinuation` (`tryAddManualContinuation`, whitelist `manualContinuationsInternalClassNames`, e.g. `kotlinx.coroutines.internal.ScopeCoroutine`) is for classes whose real `getStackTraceElement()` is class-wide constant (e.g. always `null`) — it eagerly bakes one synthetic `SpecCache` into a *static* field at `<clinit>`. `LazilyCachedContinuation` (`tryAddLazilyCachedContinuation`, whitelist `lazilyCachedContinuationsInternalClassNames`, e.g. `kotlinx.coroutines.debug.internal.DebugProbesImpl$CoroutineOwner`) is for classes whose real `getStackTraceElement()` carries genuine, varying per-*instance* data — it lazily computes and caches a `SpecCache` into an *instance* field the first time `$decoroutinator$cache` is read, which is always correct because a real per-instance value never changes after construction. Both mechanisms also patch the class's own `getStackTraceElement()` (via the shared `updateGetStackTraceElementMethod`) to serve the cache once populated, so external callers of that method benefit too — but the two patches call *different* things to fetch the cache, and this difference is load-bearing: `ManualContinuation`'s patch calls the smart `$decoroutinator$cache` getter (safe — its computation never calls back into `getStackTraceElement()`), while `LazilyCachedContinuation`'s patch calls the *raw* `$decoroutinator$cacheField` getter, not the smart `$decoroutinator$cache` getter. Calling the smart getter there would recurse infinitely: `$decoroutinator$cache`'s own computation calls `getStackTraceElement()` to get the real value, which — if patched to call the smart getter — would re-enter `$decoroutinator$cache` before the field is ever populated.
+
+**`SpecCache.element` is nullable; `provider.nullElementSpecCache`** is the shared singleton (`element = null`, `specMethod` pre-set to `methodHandleInvoker.unknownSpecMethodHandle`) used whenever a `ManualContinuation`/`LazilyCachedContinuation` frame has no usable element and `fillUnknownElementsWithClassName` is off. Reusing one singleton (rather than returning plain `null`) means the caller — the injected `$decoroutinator$cache` field-based caching — never has to distinguish "not yet computed" from "computed, nothing useful," so it never recomputes.
 
 **`Method.invoke` wraps checked exceptions in `InvocationTargetException`**: When calling a protected/private method via reflection (e.g., `ClassLoader.findClass`), any checked exception thrown by the method is wrapped in `InvocationTargetException`. Direct `catch (e: ClassNotFoundException)` at the call site of `Method.invoke(...)` will NOT catch it — you must catch `InvocationTargetException` and rethrow `e.cause`. See `generator-android/specMethodsFactory-generator-android.kt`'s `ClassLoader.findClass` extension for the canonical pattern.
 

class-transformer/src/main/kotlin/internal/classTransformer.kt

@@ -16,6 +16,7 @@ import dev.reformator.stacktracedecoroutinator.intrinsics.LABEL_FIELD_NAME
 import dev.reformator.stacktracedecoroutinator.intrinsics.UNKNOWN_LINE_NUMBER
 import dev.reformator.stacktracedecoroutinator.provider.BaseContinuationExtractor
 import dev.reformator.stacktracedecoroutinator.provider.DecoroutinatorTransformed
+import dev.reformator.stacktracedecoroutinator.provider.LazilyCachedContinuation
 import dev.reformator.stacktracedecoroutinator.provider.ManualContinuation
 import dev.reformator.stacktracedecoroutinator.provider.SpecCache
 import dev.reformator.stacktracedecoroutinator.provider.internal.BaseContinuationAccessor
@@ -35,6 +36,7 @@ import java.io.InputStream
 import java.lang.invoke.MethodHandles
 import kotlin.coroutines.Continuation
 import kotlin.coroutines.jvm.internal.CoroutineStackFrame
+import kotlin.jvm.java
 
 class ClassBodyTransformationStatus(
     val updatedBody: ByteArray?,
@@ -68,7 +70,11 @@ fun transformClassBody(
             node.transformBaseContinuation()
             doTransformation = true
         } else {
-            if (node.tryAddBaseContinuationExtractor() || node.tryAddManualContinuation(lineNumbersBySpecMethodName)) {
+            if (
+                node.tryAddBaseContinuationExtractor() ||
+                node.tryAddManualContinuation(lineNumbersBySpecMethodName) ||
+                node.tryAddLazilyCachedContinuation()
+            ) {
                 doTransformation = true
             }
 
@@ -106,8 +112,14 @@ private val manualContinuationsInternalClassNames =
         "kotlinx.coroutines.internal.ScopeCoroutine"
     ).map { it.internalName }.toHashSet()
 
+private val lazilyCachedContinuationsInternalClassNames =
+    sequenceOf(
+        "kotlinx.coroutines.debug.internal.DebugProbesImpl\$CoroutineOwner"
+    ).map { it.internalName }.toHashSet()
+
 private const val baseContinuationCachesFieldName = "\$decoroutinator\$caches"
 private const val manualContinuationCacheFieldName = "\$decoroutinator\$cache"
+private const val lazilyCachedContinuationCacheFieldName = "\$decoroutinator\$cache"
 
 @Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN")
 private fun Metadata.getNonSuspendFunctionSignatures(): List<JvmMethodSignature> {
@@ -305,6 +317,24 @@ private fun ClassNode.tryAddManualContinuation(
         }
     }
 
+    updateGetStackTraceElementMethod(
+        isUsingElementCacheForGetElementMethodEnabledMethodName = isUsingElementCacheForManualContinuationGetElementMethodEnabledMethodName,
+        getSpecCacheMethodOwnerInterfaceClass = ManualContinuation::class.java,
+        getSpecCacheMethodName = manualContinuationGetCacheMethodName
+    )
+
+    lineNumbersBySpecMethodName.computeIfAbsent(Continuation<*>::resumeWith.name) {
+        hashSetOf(UNKNOWN_LINE_NUMBER)
+    }.add(UNKNOWN_LINE_NUMBER)
+
+    return true
+}
+
+private fun ClassNode.updateGetStackTraceElementMethod(
+    isUsingElementCacheForGetElementMethodEnabledMethodName: String,
+    getSpecCacheMethodOwnerInterfaceClass: Class<*>,
+    getSpecCacheMethodName: String,
+) {
     val getStackTraceElementMethod = methods.find { method ->
         method.name == CoroutineStackFrame::getStackTraceElement.name && !method.isStatic &&
         method.desc == "()${Type.getDescriptor(StackTraceElement::class.java)}" &&
@@ -313,31 +343,103 @@ private fun ClassNode.tryAddManualContinuation(
 
     @Suppress("IfThenToSafeAccess")
     if (getStackTraceElementMethod != null) {
-        getStackTraceElementMethod.instructions.insertBefore(getStackTraceElementMethod.instructions.first, InsnList().apply {
-            add(MethodInsnNode(
-                Opcodes.INVOKESTATIC,
-                Type.getInternalName(providerApiClass),
-                isUsingElementCacheForManualContinuationGetElementMethodEnabledMethodName,
-                "()${Type.BOOLEAN_TYPE.descriptor}"
-            ))
-            val disabledLabel = LabelNode()
-            add(JumpInsnNode(Opcodes.IFEQ, disabledLabel))
+        getStackTraceElementMethod.instructions.insertBefore(
+            getStackTraceElementMethod.instructions.first,
+            InsnList().apply {
+                add(MethodInsnNode(
+                    Opcodes.INVOKESTATIC,
+                    Type.getInternalName(providerApiClass),
+                    isUsingElementCacheForGetElementMethodEnabledMethodName,
+                    "()${Type.BOOLEAN_TYPE.descriptor}"
+                ))
+                val disabledLabel = LabelNode()
+                add(JumpInsnNode(Opcodes.IFEQ, disabledLabel))
+
+                add(VarInsnNode(Opcodes.ALOAD, 0))
+                add(MethodInsnNode(
+                    Opcodes.INVOKEINTERFACE,
+                    Type.getInternalName(getSpecCacheMethodOwnerInterfaceClass),
+                    getSpecCacheMethodName,
+                    "()${Type.getDescriptor(SpecCache::class.java)}"
+                ))
+                add(InsnNode(Opcodes.DUP))
+                val cacheIsNullLabel = LabelNode()
+                add(JumpInsnNode(Opcodes.IFNULL, cacheIsNullLabel))
+
+                add(MethodInsnNode(
+                    Opcodes.INVOKEVIRTUAL,
+                    Type.getInternalName(SpecCache::class.java),
+                    specCacheGetElementMethodName,
+                    "()${Type.getDescriptor(StackTraceElement::class.java)}"
+                ))
+                add(InsnNode(Opcodes.ARETURN))
+
+                add(cacheIsNullLabel)
+                add(FrameNode(
+                    Opcodes.F_SAME1,
+                    0,
+                    null,
+                    1,
+                    arrayOf(Type.getInternalName(SpecCache::class.java))
+                ))
+                add(InsnNode(Opcodes.POP))
+                add(disabledLabel)
+                add(FrameNode(Opcodes.F_SAME, 0, null, 0, null))
+            }
+        )
+    }
+}
+
+private fun ClassNode.tryAddLazilyCachedContinuation(): Boolean {
+    if (isInterface || name !in lazilyCachedContinuationsInternalClassNames) return false
+
+    interfaces = interfaces.orEmpty() + Type.getInternalName(LazilyCachedContinuation::class.java)
+
+    fields = fields.orEmpty() + FieldNode(
+        Opcodes.ASM9,
+        Opcodes.ACC_PRIVATE or Opcodes.ACC_SYNTHETIC,
+        lazilyCachedContinuationCacheFieldName,
+        Type.getDescriptor(SpecCache::class.java),
+        null,
+        null
+    )
+
+    methods = methods.orEmpty() + MethodNode(Opcodes.ASM9).apply {
+        access = Opcodes.ACC_PUBLIC or Opcodes.ACC_SYNTHETIC
+        name = lazilyCachedContinuationGetCacheFieldMethodName
+        desc = "()${Type.getDescriptor(SpecCache::class.java)}"
+        instructions = InsnList().apply {
             add(VarInsnNode(Opcodes.ALOAD, 0))
-            add(MethodInsnNode(
-                Opcodes.INVOKEVIRTUAL,
-                this@tryAddManualContinuation.name,
-                continuationCachedGetCacheElementMethodName,
-                "()${Type.getDescriptor(StackTraceElement::class.java)}"
+            add(FieldInsnNode(
+                Opcodes.GETFIELD,
+                this@tryAddLazilyCachedContinuation.name,
+                lazilyCachedContinuationCacheFieldName,
+                Type.getDescriptor(SpecCache::class.java)
             ))
             add(InsnNode(Opcodes.ARETURN))
-            add(disabledLabel)
-            add(FrameNode(Opcodes.F_SAME, 0, null, 0, null))
-        })
+        }
+    } + MethodNode(Opcodes.ASM9).apply {
+        access = Opcodes.ACC_PUBLIC or Opcodes.ACC_SYNTHETIC
+        name = lazilyCachedContinuationSetCacheFieldMethodName
+        desc = "(${Type.getDescriptor(SpecCache::class.java)})V"
+        instructions = InsnList().apply {
+            add(VarInsnNode(Opcodes.ALOAD, 0))
+            add(VarInsnNode(Opcodes.ALOAD, 1))
+            add(FieldInsnNode(
+                Opcodes.PUTFIELD,
+                this@tryAddLazilyCachedContinuation.name,
+                lazilyCachedContinuationCacheFieldName,
+                Type.getDescriptor(SpecCache::class.java)
+            ))
+            add(InsnNode(Opcodes.RETURN))
+        }
     }
 
-    lineNumbersBySpecMethodName.computeIfAbsent(Continuation<*>::resumeWith.name) {
-        hashSetOf(UNKNOWN_LINE_NUMBER)
-    }.add(UNKNOWN_LINE_NUMBER)
+    updateGetStackTraceElementMethod(
+        isUsingElementCacheForGetElementMethodEnabledMethodName = isUsingElementCacheForLazilyCachedContinuationGetElementMethodEnabledMethodName,
+        getSpecCacheMethodOwnerInterfaceClass = LazilyCachedContinuation::class.java,
+        getSpecCacheMethodName = lazilyCachedContinuationGetCacheFieldMethodName
+    )
 
     return true
 }
@@ -832,9 +934,6 @@ private val debugMetadataMethodNameMethodName: String
 private val debugMetadataClassNameMethodName: String
     @LoadConstant("debugMetadataClassNameMethodName") get() = fail()
 
-private val continuationCachedGetCacheElementMethodName: String
-    @LoadConstant("continuationCachedGetCacheElementMethodName") get() = fail()
-
 private val isDecoroutinatorEnabledMethodName: String
     @LoadConstant("isDecoroutinatorEnabledMethodName") get() = fail()
 
@@ -891,3 +990,18 @@ private val fillUnknownElementsWithClassNameMethodName: String
 
 private val isUsingElementCacheForManualContinuationGetElementMethodEnabledMethodName: String
     @LoadConstant("isUsingElementCacheForManualContinuationGetElementMethodEnabledMethodName") get() = fail()
+
+private val manualContinuationGetCacheMethodName: String
+    @LoadConstant("manualContinuationGetCacheMethodName") get() = fail()
+
+private val lazilyCachedContinuationGetCacheFieldMethodName: String
+    @LoadConstant("lazilyCachedContinuationGetCacheFieldMethodName") get() = fail()
+
+private val lazilyCachedContinuationSetCacheFieldMethodName: String
+    @LoadConstant("lazilyCachedContinuationSetCacheFieldMethodName") get() = fail()
+
+private val specCacheGetElementMethodName: String
+    @LoadConstant("specCacheGetElementMethodName") get() = fail()
+
+private val isUsingElementCacheForLazilyCachedContinuationGetElementMethodEnabledMethodName: String
+    @LoadConstant("isUsingElementCacheForLazilyCachedContinuationGetElementMethodEnabledMethodName") get() = fail()

common/src/main/kotlin/internal/awakener.kt

@@ -59,8 +59,9 @@ private inline fun CoroutineStackFrame.getElementAndSpecMethod(
     val cache = (this as? ContinuationCached)?.`$decoroutinator$cache`
     if (cache != null) {
         val specMethod = cache.specMethod ?: run {
-            val specMethod = specMethodsFactory.getSpecMethodHandle(cache.element) ?:
-                methodHandleInvoker.unknownSpecMethodHandle
+            val specMethod = cache.element
+                ?.let { specMethodsFactory.getSpecMethodHandle(it) }
+                ?: methodHandleInvoker.unknownSpecMethodHandle
             cache.specMethod = specMethod
             specMethod
         }
@@ -93,20 +94,6 @@ private fun BaseContinuation.getElementsAndSpecMethods(): List<ElementAndSpecMet
         }
     }
 
-private fun CoroutineStackFrame.getNormalizedStackTraceElement(): StackTraceElement? {
-    val element = getStackTraceElement()
-    return when {
-        element != null -> element
-        fillUnknownElementsWithClassName -> StackTraceElement(
-            javaClass.name,
-            Continuation<*>::resumeWith.name,
-            null,
-            UNKNOWN_LINE_NUMBER
-        )
-        else -> null
-    }
-}
-
 private fun BaseContinuation.stdlibAwake(accessor: BaseContinuationAccessor, result: Any?) {
     var newResult = result
     var baseContinuation = this

common/src/main/kotlin/internal/di-common.kt

@@ -78,6 +78,15 @@ internal val isUsingElementFactoryForBaseContinuationEnabled: Boolean =
         ).toBoolean()
     }
 
+internal val isUsingElementCacheForLazilyCachedContinuationGetElementMethodEnabled =
+    fillUnknownElementsWithClassName &&
+    getRuntimeSettingsValue({ isUsingElementCacheForLazilyCachedContinuationGetElementMethodEnabled }) {
+        System.getProperty(
+            "dev.reformator.stacktracedecoroutinator.isUsingElementCacheForLazilyCachedContinuationGetElementMethodEnabled",
+            "true"
+        ).toBoolean()
+    }
+
 @Suppress("ObjectPropertyName")
 private val _transformedClassesRegistry: TransformedClassesRegistry? =
     if (enabled) TransformedClassesRegistryImpl() else null