fix: synthesize minimal unresolved class stubs for main sources

Author: edengilbertus

amber-pipeline/src/main/kotlin/dev/amber/pipeline/codegen/JavaEmitter.kt

@@ -12,6 +12,7 @@ import dev.amber.core.model.SourcePartitionKind
 import dev.amber.core.model.StmtIr
 import dev.amber.core.model.TermIr
 import dev.amber.core.model.placeholderReasons
+import dev.amber.pipeline.dex.SourcePartitioner
 import dev.amber.pipeline.dex.TypeAnnotatedMethod
 import dev.amber.pipeline.mapping.MappingRenameIndex
 import java.nio.file.Path
@@ -58,6 +59,11 @@ class JavaEmitter {
         val statementCount: Int,
     )
 
+    private data class StubClassNode(
+        val simpleName: String,
+        val children: MutableMap<String, StubClassNode> = sortedMapOf(),
+    )
+
     fun emit(methods: List<TypeAnnotatedMethod>, classes: List<ClassIr> = emptyList()): EmitResult {
         val diagnostics = mutableListOf<String>()
         val fileBuilders = mutableMapOf<String, TypeSpec.Builder>()
@@ -143,6 +149,9 @@ class JavaEmitter {
             .forEach { (sourceSetName, partitionMethods) ->
                 val partitionOutputDir = outputDir.resolve(sourceSetName).resolve("java")
                 partitionOutputDir.createDirectories()
+                val partitionKind = partitionMethods.firstOrNull()?.sourcePartition?.kind
+                    ?: SourcePartitionKind.PRIMARY_APP
+                val includeInMainCompile = partitionKind != SourcePartitionKind.LIKELY_LIBRARY
                 val classMethodCounts = partitionMethods.groupingBy { annotated ->
                     parseClassName(annotated.emittedMethodId)
                 }.eachCount()
@@ -167,8 +176,15 @@ class JavaEmitter {
                     )
                 }
 
-                val partitionKind = partitionMethods.firstOrNull()?.sourcePartition?.kind
-                    ?: SourcePartitionKind.PRIMARY_APP
+                if (includeInMainCompile) {
+                    emitUnresolvedReferenceStubs(
+                        methods = partitionMethods,
+                        outputDir = partitionOutputDir,
+                        diagnostics = diagnostics,
+                        packageFilter = packageFilter,
+                    )
+                }
+
                 val partitionSummary = summarizeSourcePartitions(
                     partitionMethods.map { annotated -> annotated.sourcePartition },
                 )
@@ -183,7 +199,7 @@ class JavaEmitter {
                     language = EmittedSourceLanguage.JAVA,
                     partitionKind = partitionKind,
                     classCount = emittedClassCount,
-                    includeInMainCompile = partitionKind != SourcePartitionKind.LIKELY_LIBRARY,
+                    includeInMainCompile = includeInMainCompile,
                     partitionHeuristics = partitionSummary.heuristics,
                     heuristicCounts = partitionSummary.heuristicCounts,
                     tplFingerprintDigests = partitionSummary.fingerprintDigests,
@@ -400,6 +416,56 @@ class JavaEmitter {
         }
     }
 
+    private fun emitUnresolvedReferenceStubs(
+        methods: List<TypeAnnotatedMethod>,
+        outputDir: Path,
+        diagnostics: MutableList<String>,
+        packageFilter: String?,
+    ) {
+        val emittedClassNames = methods.asSequence()
+            .map { annotated -> parseClassName(annotated.emittedMethodId) }
+            .filter { className -> packageFilter == null || className.startsWith(packageFilter) }
+            .toSet()
+        val emittedPackages = emittedClassNames.map(::internalPackageName).toSet()
+        val unresolvedClassNames = methods.asSequence()
+            .flatMap { annotated -> collectReferencedClassNames(annotated).asSequence() }
+            .distinct()
+            .filter { className ->
+                shouldStubReferencedClass(
+                    className = className,
+                    emittedClassNames = emittedClassNames,
+                    emittedPackages = emittedPackages,
+                )
+            }
+            .sorted()
+            .toList()
+
+        if (unresolvedClassNames.isEmpty()) {
+            return
+        }
+
+        val stubTrees = linkedMapOf<String, StubClassNode>()
+        unresolvedClassNames.forEach { className ->
+            addStubClassNode(stubTrees, className)
+        }
+
+        stubTrees.forEach { (topLevelClassName, rootNode) ->
+            try {
+                JavaFile.builder(packageName(topLevelClassName), stubTypeSpec(rootNode, isNested = false).build())
+                    .skipJavaLangImports(true)
+                    .build()
+                    .writeTo(outputDir)
+            } catch (e: Exception) {
+                diagnostics.add("Failed to write unresolved stub L$topLevelClassName;: ${e.message}")
+            }
+        }
+
+        diagnostics.add(
+            "Stubbed ${unresolvedClassNames.size} unresolved referenced classes: ${unresolvedClassNames.take(10).joinToString()}" +
+                if (unresolvedClassNames.size > 10) " ..." else "",
+        )
+    }
+
     private fun emitClassBuilder(
         className: String,
         classMethods: List<TypeAnnotatedMethod>,
@@ -440,6 +506,141 @@ class JavaEmitter {
             )
     }
 
+    private fun stubTypeSpec(node: StubClassNode, isNested: Boolean): TypeSpec.Builder {
+        val modifiers = if (isNested) {
+            arrayOf(Modifier.PUBLIC, Modifier.STATIC)
+        } else {
+            arrayOf(Modifier.PUBLIC)
+        }
+        return TypeSpec.classBuilder(sanitizeIdentifier(node.simpleName, isClass = true))
+            .addModifiers(*modifiers)
+            .apply {
+                node.children.values.forEach { child ->
+                    addType(stubTypeSpec(child, isNested = true).build())
+                }
+            }
+    }
+
+    private fun addStubClassNode(
+        stubTrees: MutableMap<String, StubClassNode>,
+        className: String,
+    ) {
+        val packageName = internalPackageName(className)
+        val classPart = className.substringAfterLast('/')
+        val segments = classPart.split('$').filter { segment -> segment.isNotEmpty() }
+        if (segments.isEmpty()) {
+            return
+        }
+
+        val topLevelClassName = if (packageName.isEmpty()) segments.first() else "$packageName/${segments.first()}"
+        val rootNode = stubTrees.getOrPut(topLevelClassName) {
+            StubClassNode(simpleName = segments.first())
+        }
+        var current = rootNode
+        segments.drop(1).forEach { segment ->
+            current = current.children.getOrPut(segment) {
+                StubClassNode(simpleName = segment)
+            }
+        }
+    }
+
+    private fun shouldStubReferencedClass(
+        className: String,
+        emittedClassNames: Set<String>,
+        emittedPackages: Set<String>,
+    ): Boolean {
+        if (className.isBlank()) {
+            return false
+        }
+        if (className in emittedClassNames) {
+            return false
+        }
+        val topLevelClassName = topLevelInternalClassName(className)
+        if (topLevelClassName in emittedClassNames) {
+            return false
+        }
+        if (internalPackageName(className) !in emittedPackages) {
+            return false
+        }
+        if (SourcePartitioner.partitionClassId("L$className;").kind == SourcePartitionKind.LIKELY_LIBRARY) {
+            return false
+        }
+        return true
+    }
+
+    private fun collectReferencedClassNames(annotated: TypeAnnotatedMethod): Set<String> {
+        val referencedClassNames = linkedSetOf<String>()
+        collectClassNamesFromMethodRef(annotated.emittedMethodId, referencedClassNames)
+        annotated.typeFacts.values.forEach { fact ->
+            collectClassNamesFromDescriptor(annotated.emittedDescriptor(fact.descriptor), referencedClassNames)
+        }
+        annotated.syntheticAccessor?.targetMethodRef?.let { methodRef ->
+            collectClassNamesFromMethodRef(annotated.emittedMethodRef(methodRef), referencedClassNames)
+        }
+        annotated.syntheticAccessor?.targetFieldRef?.let { fieldRef ->
+            val emittedFieldRef = annotated.emittedFieldRef(fieldRef)
+            collectClassNamesFromDescriptor(emittedFieldRef.definingClass, referencedClassNames)
+            collectClassNamesFromDescriptor(emittedFieldRef.type, referencedClassNames)
+        }
+        annotated.methodIr.tryRegions.forEach { tryRegion ->
+            tryRegion.handlers.forEach { handler ->
+                handler.catchType?.let { catchType ->
+                    collectClassNamesFromDescriptor(annotated.emittedDescriptor(catchType), referencedClassNames)
+                }
+            }
+        }
+        annotated.methodIr.handlers.forEach { handler ->
+            handler.caughtTypes.forEach { caughtType ->
+                collectClassNamesFromDescriptor(annotated.emittedDescriptor(caughtType), referencedClassNames)
+            }
+        }
+        annotated.renameResult.blocks.forEach { block ->
+            block.stmts.forEach { stmt ->
+                stmt.invoke?.let { invoke ->
+                    collectClassNamesFromMethodRef(annotated.emittedMethodRef(invoke.methodRef), referencedClassNames)
+                }
+                stmt.fieldRef?.let { fieldRef ->
+                    val emittedFieldRef = annotated.emittedFieldRef(fieldRef)
+                    collectClassNamesFromDescriptor(emittedFieldRef.definingClass, referencedClassNames)
+                    collectClassNamesFromDescriptor(emittedFieldRef.type, referencedClassNames)
+                }
+                stmt.imm?.typeValue?.let { typeValue ->
+                    collectClassNamesFromDescriptor(annotated.emittedDescriptor(typeValue), referencedClassNames)
+                }
+            }
+        }
+        return referencedClassNames
+    }
+
+    private fun collectClassNamesFromMethodRef(methodRef: String, sink: MutableSet<String>) {
+        collectClassNamesFromDescriptor(parseOwnerClass(methodRef), sink)
+        parseParams(methodRef).forEach { descriptor ->
+            collectClassNamesFromDescriptor(descriptor, sink)
+        }
+        collectClassNamesFromDescriptor(parseReturnType(methodRef), sink)
+    }
+
+    private fun collectClassNamesFromDescriptor(descriptor: String, sink: MutableSet<String>) {
+        val normalized = descriptor.trim().trimEnd(';')
+        when {
+            normalized.isEmpty() -> Unit
+            normalized.startsWith("[") -> collectClassNamesFromDescriptor(normalized.substring(1), sink)
+            normalized.length == 1 && normalized[0] in "VZBCSIJFD" -> Unit
+            normalized == "java/lang/Object" || normalized == "java/lang/String" -> Unit
+            normalized.startsWith("L") -> sink += normalized.substring(1)
+            descriptor.startsWith("L") && descriptor.endsWith(";") -> sink += descriptor.substring(1, descriptor.length - 1)
+            '/' in normalized -> sink += normalized.removePrefix("L")
+        }
+    }
+
+    private fun internalPackageName(className: String): String = className.substringBeforeLast('/', "")
+
+    private fun topLevelInternalClassName(className: String): String {
+        val packageName = internalPackageName(className)
+        val topLevelSimpleName = className.substringAfterLast('/').substringBefore('$')
+        return if (packageName.isEmpty()) topLevelSimpleName else "$packageName/$topLevelSimpleName"
+    }
+
     private fun methodsAreContiguousByClass(methods: List<TypeAnnotatedMethod>): Boolean {
         val closedClasses = mutableSetOf<String>()
         var currentClassName: String? = null