From 0dd5ad08f3ea07c5fc9074d10cc5ef85dd577cec Mon Sep 17 00:00:00 2001
From: Kristoffer Andersen <kjaa@google.com>
Date: Mon, 8 Nov 2021 08:50:18 +0100
Subject: [PATCH] [EE-IR] Support inline functions in IR evaluator

Inline functions called by the fragment is detected by the fragment
frontend and the PSI for them is supplied to the compiler backend.

For the purpose of compilation, the inline function is considered as
part of the same source module as the fragment, and then discarded
upon loading the compiled fragment for evaluation.
---
 .../psi2ir/generators/ModuleGenerator.kt      |  4 +--
 .../fragments/FragmentModuleGenerator.kt      | 31 ++++++++++++-------
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ModuleGenerator.kt b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ModuleGenerator.kt
index e2c39f64687..0c5eef87824 100644
--- a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ModuleGenerator.kt
+++ b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/ModuleGenerator.kt
@@ -52,7 +52,7 @@ open class ModuleGenerator(
         ExternalDependenciesGenerator(context.symbolTable, irProviders).generateUnboundSymbolsAsDependencies()
     }
 
-    private fun generateSingleFile(irDeclarationGenerator: DeclarationGenerator, ktFile: KtFile, module: IrModuleFragment): IrFileImpl {
+    fun generateSingleFile(irDeclarationGenerator: DeclarationGenerator, ktFile: KtFile, module: IrModuleFragment): IrFileImpl {
         val irFile = createEmptyIrFile(ktFile, module)
 
         val constantValueGenerator = irDeclarationGenerator.context.constantValueGenerator
@@ -85,7 +85,7 @@ open class ModuleGenerator(
         return irFile
     }
 
-    private fun createEmptyIrFile(ktFile: KtFile, module: IrModuleFragment): IrFileImpl {
+     fun createEmptyIrFile(ktFile: KtFile, module: IrModuleFragment): IrFileImpl {
         val fileEntry = PsiIrFileEntry(ktFile)
         val packageFragmentDescriptor = context.moduleDescriptor.findPackageFragmentForFile(ktFile)!!
         return IrFileImpl(fileEntry, packageFragmentDescriptor, module).apply {
diff --git a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/fragments/FragmentModuleGenerator.kt b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/fragments/FragmentModuleGenerator.kt
index dd0ccd67287..08c6a80b32e 100644
--- a/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/fragments/FragmentModuleGenerator.kt
+++ b/compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/fragments/FragmentModuleGenerator.kt
@@ -12,8 +12,7 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrModuleFragmentImpl
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
 import org.jetbrains.kotlin.psi.KtBlockCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi2ir.generators.GeneratorContext
-import org.jetbrains.kotlin.psi2ir.generators.ModuleGenerator
+import org.jetbrains.kotlin.psi2ir.generators.*
 import org.jetbrains.kotlin.resolve.lazy.descriptors.findPackageFragmentForFile
 
 class FragmentModuleGenerator(
@@ -24,18 +23,26 @@ class FragmentModuleGenerator(
     override fun generateModuleFragment(
         ktFiles: Collection<KtFile>,
     ): IrModuleFragment {
-        val ktBlockCodeFragment = ktFiles.singleOrNull() as? KtBlockCodeFragment
-            ?: TODO("Multiple fragments in one compilation not understood and implemented yet")
+        assert(ktFiles.singleOrNull { it is KtBlockCodeFragment} != null) {
+            "Amongst all files passed to the FragmentModuleGenerator should be exactly one KtBlockCodeFragment"
+        }
         return IrModuleFragmentImpl(context.moduleDescriptor, context.irBuiltIns).also { irModule ->
             val irDeclarationGenerator = FragmentDeclarationGenerator(context, fragmentInfo)
-            irModule.files.add(
-                createEmptyIrFile(ktBlockCodeFragment).apply {
-                    declarations.add(
-                        irDeclarationGenerator.generateClassForCodeFragment(ktBlockCodeFragment)
-                    )
-                    patchDeclarationParents()
-                }
-            )
+            ktFiles.forEach { ktFile ->
+                irModule.files.add(
+                    if (ktFile is KtBlockCodeFragment) {
+                        createEmptyIrFile(ktFile).apply {
+                            declarations.add(
+                                irDeclarationGenerator.generateClassForCodeFragment(ktFile)
+                            )
+                            patchDeclarationParents()
+                        }
+                    } else {
+                        val fileContext = context.createFileScopeContext(ktFile)
+                        generateSingleFile(DeclarationGenerator(fileContext), ktFile, irModule)
+                    }
+                )
+            }
         }
     }