Repository: yandex/yatagan (kotlin)

Issue:
Generate cascading switch statements when number of bindings is large
Yatagan, as Dagger in `fastInit` uses _"switching providers"_ mechanism to reduce classloading. This means that for every binding, that can't be created inline (has `Lazy`/`Provider` usages) there's an index assigned. And every provider/lazy object instantiates the binding using the index and a lookup (switch) table. 

Now, the hypothesis is, that when the lookup table becomes very large, the performance is going to be penalized. Yatagan currently generates plain switch statement, while Dagger did so with only up to 100 objects, after that it generated another switch level with `index % 100` labels.

Let's try to investigate and possibly introduce the optimization.

Might be reasonable to do when the performance testing is implemented in #7

Pull request:
Introduce an `yatagan.experimental.maxSlotsPerSwitch`
 option to enable/customize nested switches behavior.
 It's set to infinity by default as before the patch.

Add a "heavy" test-case to trigger this mechanic.

Resolves #8

Patch:
diff --git a/codegen/impl/src/main/kotlin/AccessStrategyManager.kt b/codegen/impl/src/main/kotlin/AccessStrategyManager.kt
index b3e2ceda..39ae2a52 100644
--- a/codegen/impl/src/main/kotlin/AccessStrategyManager.kt
+++ b/codegen/impl/src/main/kotlin/AccessStrategyManager.kt
@@ -113,7 +113,7 @@ internal class AccessStrategyManager(
                             lazyStrategy = if (usage.lazy + usage.provider > 0) {
                                 SlotProviderStrategy(
                                     binding = binding,
-                                    multiFactory = multiFactory.get(),
+                                    slot = multiFactory.get().requestSlot(binding),
                                     cachingProvider = unscopedProviderGenerator.get(),
                                 )
                             } else null
diff --git a/codegen/impl/src/main/kotlin/ComponentGenerator.kt b/codegen/impl/src/main/kotlin/ComponentGenerator.kt
index 52e1c732..a41005e4 100644
--- a/codegen/impl/src/main/kotlin/ComponentGenerator.kt
+++ b/codegen/impl/src/main/kotlin/ComponentGenerator.kt
@@ -35,9 +35,14 @@ import javax.lang.model.element.Modifier.STATIC
 internal class ComponentGenerator private constructor(
     private val graph: BindingGraph,
     val generatedClassName: ClassName,
+    maxSlotsPerSwitch: Int,
 ) {
-    constructor(graph: BindingGraph): this(
+    constructor(
+        graph: BindingGraph,
+        maxSlotsPerSwitch: Int,
+    ): this(
         graph = graph,
+        maxSlotsPerSwitch = maxSlotsPerSwitch,
         generatedClassName = graph.model.name.let { name ->
             check(name is ClassNameModel)
             // Keep name mangling in sync with loader!
@@ -57,6 +62,7 @@ internal class ComponentGenerator private constructor(
     private val slotSwitchingGenerator = lazyProvider {
         SlotSwitchingGenerator(
             thisGraph = graph,
+            maxSlotsPerSwitch = maxSlotsPerSwitch,
         ).also(::registerContributor)
     }
     private val unscopedProviderGenerator = lazyProvider {
@@ -120,6 +126,7 @@ internal class ComponentGenerator private constructor(
     private val childGenerators: Collection<ComponentGenerator> = graph.children.map { childGraph ->
         ComponentGenerator(
             graph = childGraph,
+            maxSlotsPerSwitch = maxSlotsPerSwitch,
             generatedClassName = generatedClassName.nestedClass(
                 subcomponentNs.name(childGraph.model.name, suffix = "Impl", firstCapital = true)
             ),
diff --git a/codegen/impl/src/main/kotlin/ComponentGeneratorFacade.kt b/codegen/impl/src/main/kotlin/ComponentGeneratorFacade.kt
index dce872ca..b8f96589 100644
--- a/codegen/impl/src/main/kotlin/ComponentGeneratorFacade.kt
+++ b/codegen/impl/src/main/kotlin/ComponentGeneratorFacade.kt
@@ -21,9 +21,11 @@ import com.yandex.yatagan.core.graph.BindingGraph
 
 class ComponentGeneratorFacade(
     graph: BindingGraph,
+    maxSlotsPerSwitch: Int,
 ) {
     private val generator = ComponentGenerator(
         graph = graph,
+        maxSlotsPerSwitch = maxSlotsPerSwitch,
     )
 
     val targetPackageName: String
diff --git a/codegen/impl/src/main/kotlin/SlotSwitchingGenerator.kt b/codegen/impl/src/main/kotlin/SlotSwitchingGenerator.kt
index e846056f..75a2b0d9 100644
--- a/codegen/impl/src/main/kotlin/SlotSwitchingGenerator.kt
+++ b/codegen/impl/src/main/kotlin/SlotSwitchingGenerator.kt
@@ -21,16 +21,23 @@ import com.yandex.yatagan.codegen.poetry.TypeSpecBuilder
 import com.yandex.yatagan.codegen.poetry.buildExpression
 import com.yandex.yatagan.core.graph.BindingGraph
 import com.yandex.yatagan.core.graph.bindings.Binding
+import javax.lang.model.element.Modifier
 
 internal class SlotSwitchingGenerator(
     private val thisGraph: BindingGraph,
+    private val maxSlotsPerSwitch: Int,
 ) : ComponentGenerator.Contributor {
+    init {
+        require(maxSlotsPerSwitch > 1 || maxSlotsPerSwitch == -1) {
+            "maxSlotsPerSwitch should be at least 2 or -1 (infinity)"
+        }
+    }
 
-    private val boundSlots = mutableMapOf<Binding, Int>()
-    private var nextFreeSlot = 0
+    private val bindings = arrayListOf<Binding>()
 
     fun requestSlot(forBinding: Binding): Int {
-        return boundSlots.getOrPut(forBinding) { nextFreeSlot++ }
+        bindings += forBinding
+        return bindings.lastIndex
     }
 
     override fun generate(builder: TypeSpecBuilder) {
@@ -38,18 +45,57 @@ internal class SlotSwitchingGenerator(
             modifiers(/*package-private*/)
             returnType(ClassName.OBJECT)
             parameter(ClassName.INT, "slot")
-            controlFlow("switch(slot)") {
-                boundSlots.forEach { (binding, slot) ->
-                    +buildExpression {
-                        +"case $slot: return "
-                        binding.generateAccess(
-                            builder = this,
-                            inside = thisGraph,
-                            isInsideInnerClass = false,
-                        )
+            if (maxSlotsPerSwitch > 1 && bindings.size > maxSlotsPerSwitch) {
+                // Strategy with two-level-nested switches
+
+                val outerSlotsCount = bindings.size / maxSlotsPerSwitch
+                controlFlow("switch(slot / $maxSlotsPerSwitch)") {
+                    for(outerSlot in 0 .. outerSlotsCount) {
+                        val nestedFactoryName = "$FactoryMethodName\$$outerSlot"
+                        builder.method(nestedFactoryName) {
+                            modifiers(Modifier.PRIVATE)
+                            returnType(ClassName.OBJECT)
+                            parameter(ClassName.INT, "slot")
+                            controlFlow("switch(slot)") {
+                                val startIndex = outerSlot * maxSlotsPerSwitch
+                                var nestedSlot = 0
+                                while((startIndex + nestedSlot) < bindings.size && nestedSlot < maxSlotsPerSwitch) {
+                                    val binding = bindings[startIndex + nestedSlot]
+                                    +buildExpression {
+                                        +"case $nestedSlot: return "
+                                        binding.generateAccess(
+                                            builder = this,
+                                            inside = thisGraph,
+                                            isInsideInnerClass = false,
+                                        )
+                                    }
+                                    ++nestedSlot
+                                }
+                                +"default: throw new %T()".formatCode(Names.AssertionError)
+                            }
+                        }
+                        +buildExpression {
+                            +"case $outerSlot: return %N(slot %L 100)".formatCode(nestedFactoryName, "%")
+                        }
+                    }
+                    +"default: throw new %T()".formatCode(Names.AssertionError)
+                }
+            } else {
+                // Single switch statement
+
+                controlFlow("switch(slot)") {
+                    bindings.forEachIndexed { slot, binding ->
+                        +buildExpression {
+                            +"case $slot: return "
+                            binding.generateAccess(
+                                builder = this,
+                                inside = thisGraph,
+                                isInsideInnerClass = false,
+                            )
+                        }
                     }
+                    +"default: throw new %T()".formatCode(Names.AssertionError)
                 }
-                +"default: throw new %T()".formatCode(Names.AssertionError)
             }
         }
     }
diff --git a/codegen/impl/src/main/kotlin/accessStrategies.kt b/codegen/impl/src/main/kotlin/accessStrategies.kt
index 0e5490f6..33b96281 100644
--- a/codegen/impl/src/main/kotlin/accessStrategies.kt
+++ b/codegen/impl/src/main/kotlin/accessStrategies.kt
@@ -130,12 +130,11 @@ internal class CachingStrategyMultiThread(
 
 internal class SlotProviderStrategy(
     private val binding: Binding,
-    multiFactory: SlotSwitchingGenerator,
+    private val slot: Int,
     cachingProvider: UnscopedProviderGenerator,
 ) : AccessStrategy {
 
     private val providerName = cachingProvider.name
-    private val slot = multiFactory.requestSlot(binding)
 
     override fun generateAccess(
         builder: ExpressionBuilder,
diff --git a/processor/common/src/main/kotlin/Options.kt b/processor/common/src/main/kotlin/Options.kt
index d17ebcdb..f8940b46 100644
--- a/processor/common/src/main/kotlin/Options.kt
+++ b/processor/common/src/main/kotlin/Options.kt
@@ -63,5 +63,7 @@ class Options(
         val MaxIssueEncounterPaths = IntOption("yatagan.maxIssueEncounterPaths", default = 5)
 
         val UsePlainOutput = BooleanOption("yatagan.usePlainOutput", default = false)
+
+        val MaxSlotsPerSwitch = IntOption("yatagan.experimental.maxSlotsPerSwitch", default = -1)
     }
 }
\ No newline at end of file
diff --git a/processor/common/src/main/kotlin/process.kt b/processor/common/src/main/kotlin/process.kt
index 4906f8df..509e6889 100644
--- a/processor/common/src/main/kotlin/process.kt
+++ b/processor/common/src/main/kotlin/process.kt
@@ -85,6 +85,7 @@ fun <Source> process(
             try {
                 val generator = ComponentGeneratorFacade(
                     graph = graphRoot,
+                    maxSlotsPerSwitch = delegate.options[Options.MaxSlotsPerSwitch],
                 )
                 delegate.openFileForGenerating(
                     sources = allSourcesSequence(delegate, graphRoot),
diff --git a/rt/engine/src/main/kotlin/RuntimeComponent.kt b/rt/engine/src/main/kotlin/RuntimeComponent.kt
index 0f7142a1..b78cbd9e 100644
--- a/rt/engine/src/main/kotlin/RuntimeComponent.kt
+++ b/rt/engine/src/main/kotlin/RuntimeComponent.kt
@@ -339,7 +339,8 @@ internal class RuntimeComponent(
     }
 
     override fun visitEmpty(binding: EmptyBinding): Any {
-        throw IllegalStateException("Missing binding encountered in `$graph`: $binding")
+        throw IllegalStateException(
+            "Missing binding encountered in `${graph.toString(null)}`: ${binding.toString(null)}")
     }
 
     private inner class EntryPointHandler(val dependency: NodeDependency) : MethodHandler {


Test patch:
diff --git a/testing/tests/src/main/kotlin/CompileTestDriverBase.kt b/testing/tests/src/main/kotlin/CompileTestDriverBase.kt
index 61cec2e4..7c7f009c 100644
--- a/testing/tests/src/main/kotlin/CompileTestDriverBase.kt
+++ b/testing/tests/src/main/kotlin/CompileTestDriverBase.kt
@@ -163,6 +163,7 @@ abstract class CompileTestDriverBase private constructor(
         ),
         processorOptions = mapOf(
             Options.MaxIssueEncounterPaths.key to "100",
+            Options.MaxSlotsPerSwitch.key to "100",
         ),
     )
 
diff --git a/testing/tests/src/test/kotlin/HeavyTest.kt b/testing/tests/src/test/kotlin/HeavyTest.kt
index 015b05d4..f5983571 100644
--- a/testing/tests/src/test/kotlin/HeavyTest.kt
+++ b/testing/tests/src/test/kotlin/HeavyTest.kt
@@ -28,17 +28,12 @@ import javax.inject.Provider
 @RunWith(Parameterized::class)
 @OptIn(ExperimentalGenerationApi::class)
 class HeavyTest(
-    driverProvider: Provider<CompileTestDriverBase>
+    driverProvider: Provider<CompileTestDriverBase>,
+    private val params: GenerationParams,
+    private val name: String,
 ) : CompileTestDriver by driverProvider.get() {
     companion object {
-        @JvmStatic
-        @Parameterized.Parameters(name = "{0}")
-        fun parameters() = compileTestDrivers()
-    }
-
-    @Test
-    fun `base case`() {
-        val params = GenerationParams(
+        private val baseParams = GenerationParams(
             componentTreeMaxDepth = 6,
             totalGraphNodesCount = 300,
             bindings = Distribution.build {
@@ -63,6 +58,24 @@ class HeavyTest(
             seed = 1236473289L,
         )
 
+        @JvmStatic
+        @Parameterized.Parameters(name = "{0} - {2}")
+        fun parameters() = compileTestDrivers().flatMap {
+            listOf(
+                arrayOf(it, baseParams, "Tall & Thin"),
+                arrayOf(it, baseParams.copy(
+                    componentTreeMaxDepth = 1,
+                    totalGraphNodesCount = 600,
+                    totalRootCount = 2,
+                    maxEntryPointsPerComponent = 60,
+                ), "Fat & Flat"),
+            )
+        }
+    }
+
+    @Test
+    fun procedural() {
+        println("Running '$name' heavy test")
         generate(
             params = params,
             testMethodName = "test",