JS: review fixes (renamings, etc.)

js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineTransformer.kt

@@ -20,7 +20,7 @@ import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
 import org.jetbrains.kotlin.js.backend.ast.metadata.isInlineableCoroutineBody
 import org.jetbrains.kotlin.js.inline.InlineSuspendFunctionSplitter
-import org.jetbrains.kotlin.js.inline.ProgramFragmentInliningScope
+import org.jetbrains.kotlin.js.inline.ImportInfoFragmentInliningScope
 import org.jetbrains.kotlin.js.translate.declaration.transformCoroutineMetadataToSpecialFunctions
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
@@ -72,9 +72,9 @@ class CoroutineTransformer : JsVisitorWithContextImpl() {
 fun transformCoroutines(fragments: Iterable<JsProgramFragment>) {
     val coroutineTransformer = CoroutineTransformer()
     for (fragment in fragments) {
-        val scope = ProgramFragmentInliningScope(fragment)
+        ImportInfoFragmentInliningScope.process(fragment) { scope ->
-        InlineSuspendFunctionSplitter(scope).accept(scope.allCode)
+            InlineSuspendFunctionSplitter(scope).accept(scope.allCode)
-        coroutineTransformer.accept(scope.allCode)
+            coroutineTransformer.accept(scope.allCode)
-        scope.update()
+        }
     }
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/FunctionReader.kt

@@ -171,11 +171,11 @@ class FunctionReader(
             readFunction(key) ?: NotFoundMarker
     }
 
-    operator fun get(descriptor: CallableDescriptor, fragment: JsProgramFragment): FunctionWithWrapper? {
+    operator fun get(descriptor: CallableDescriptor, callsiteFragment: JsProgramFragment): FunctionWithWrapper? {
         return functionCache.get(descriptor).let {
             if (it === NotFoundMarker) null else {
                 val (fn, info) = it as Pair<*, *>
-                renameModules(descriptor, (fn as FunctionWithWrapper).deepCopy(), info as ModuleInfo, fragment)
+                renameModules(descriptor, (fn as FunctionWithWrapper).deepCopy(), info as ModuleInfo, callsiteFragment)
             }
         }
     }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/InlineAstVisitor.kt

@@ -19,9 +19,10 @@ import org.jetbrains.kotlin.js.inline.util.refreshLabelNames
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 
-class InlinerImpl(
+// Visits AST, detects inline function declarations and invocations.
-    val jsInliner: JsInliner,
+class InlineAstVisitor(
-    val scope: InliningScope
+    private val jsInliner: JsInliner,
+    private val scope: InliningScope
 ) : JsVisitorWithContextImpl() {
 
     override fun visit(x: JsBinaryOperation, ctx: JsContext<*>): Boolean {
@@ -32,7 +33,7 @@ class InlinerImpl(
                 val name = left.name
                 if (name != null) {
                     extractFunction(right)?.let { function ->
-                        jsInliner.process(InlineFunctionDefinition(function, null), null, scope)
+                        jsInliner.process(InlineFunctionDefinition(function, null), null, scope.fragment)
                     }
                 }
             }
@@ -46,7 +47,7 @@ class InlinerImpl(
         val name = x.name
         if (initializer != null && name != null) {
             extractFunction(initializer)?.let { function ->
-                jsInliner.process(InlineFunctionDefinition(function, null), null, scope)
+                jsInliner.process(InlineFunctionDefinition(function, null), null, scope.fragment)
             }
         }
 
@@ -55,7 +56,7 @@ class InlinerImpl(
 
     override fun visit(x: JsInvocation, ctx: JsContext<*>): Boolean {
         InlineMetadata.decompose(x)?.let {
-            jsInliner.process(InlineFunctionDefinition(it.function, it.tag.value), x, scope)
+            jsInliner.process(InlineFunctionDefinition(it.function, it.tag.value), x, scope.fragment)
         }
 
         return super.visit(x, ctx)
@@ -129,7 +130,7 @@ class InlinerImpl(
 
     private fun hasToBeInlined(call: JsInvocation): Boolean {
         val strategy = call.inlineStrategy
-        return if (strategy == null || !strategy.isInline) false else jsInliner.functionContext.hasFunctionDefinition(call, scope)
+        return if (strategy == null || !strategy.isInline) false else jsInliner.functionDefinitionLoader.hasFunctionDefinition(call, scope.fragment)
     }
 
     private fun patchReturnsFromSecondaryConstructor(function: JsFunction) {

js/js.inliner/src/org/jetbrains/kotlin/js/inline/InlineSuspendFunctionSplitter.kt

@@ -10,8 +10,12 @@ import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
 import org.jetbrains.kotlin.js.backend.ast.metadata.isInlineableCoroutineBody
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 
+// Goes through `suspend inline` function definitions and transforms into two entites:
+// 1. A `suspend` function definition, which could be invoked at run time
+// 2. A `suspend inline` function definition, which doesn't work at run time, but is understood by the inliner.
+// TODO remove the wrapped version for `private inline suspend fun`'s
 class InlineSuspendFunctionSplitter(
-    val scope: ProgramFragmentInliningScope
+    val scope: ImportInfoFragmentInliningScope
 ) : JsVisitorWithContextImpl() {
 
     override fun endVisit(x: JsExpressionStatement, ctx: JsContext<*>) {

js/js.inliner/src/org/jetbrains/kotlin/js/inline/InlinerCycleReporter.kt

@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.js.backend.ast.JsInvocation
 import org.jetbrains.kotlin.js.backend.ast.metadata.descriptor
 import org.jetbrains.kotlin.js.backend.ast.metadata.inlineStrategy
 import org.jetbrains.kotlin.js.backend.ast.metadata.psiElement
-import org.jetbrains.kotlin.js.inline.context.FunctionContext
+import org.jetbrains.kotlin.js.inline.context.FunctionDefinitionLoader
 import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper
 import org.jetbrains.kotlin.resolve.inline.InlineStrategy
 import java.util.*
@@ -26,7 +26,7 @@ import java.util.*
  */
 class InlinerCycleReporter(
     val trace: DiagnosticSink,
-    private val functionContext: FunctionContext
+    private val functionDefinitionLoader: FunctionDefinitionLoader
 ) {
 
     private enum class VisitedState { IN_PROCESS, PROCESSED }
@@ -43,7 +43,7 @@ class InlinerCycleReporter(
 
     // Puts `function` on the `namedFunctionsStack` for inline call cycles reporting
     fun <T> withFunction(function: JsFunction, body: () -> T): T {
-        if (function in functionContext.functionsByFunctionNodes.keys) {
+        if (function in functionDefinitionLoader.functionsByFunctionNodes.keys) {
             namedFunctionsStack.push(function)
         }
 

js/js.inliner/src/org/jetbrains/kotlin/js/inline/InliningScope.kt

@@ -23,15 +23,15 @@ sealed class InliningScope {
 
     abstract val fragment: JsProgramFragment
 
-    abstract fun addInlinedDeclaration(tag: String?, declaration: JsStatement)
+    protected abstract fun addInlinedDeclaration(tag: String?, declaration: JsStatement)
 
-    abstract fun hasImport(name: JsName, tag: String): JsName?
+    protected abstract fun hasImport(name: JsName, tag: String): JsName?
 
-    abstract fun addImport(tag: String, vars: JsVars)
+    protected abstract fun addImport(tag: String, vars: JsVars)
 
-    open fun addLocalDeclarationBinding(inlineFunctionTag: String, name: JsName, index: Int) {}
+    protected open fun addLocalDeclarationBinding(inlineFunctionTag: String, name: JsName, index: Int) {}
 
-    open fun preprocess(statement: JsStatement) {}
+    protected open fun preprocess(statement: JsStatement) {}
 
     abstract fun update()
 
@@ -97,7 +97,7 @@ sealed class InliningScope {
                     newReplacements[name] = replacement
                 }
 
-            // Add local declarations to the nae bindings in order to correctly rename usages of the same imported local declaraions
+            // Add local declarations to the name bindings in order to correctly rename usages of the same imported local declaraions
             definition.tag?.let { tag ->
                 localDeclarations.forEachIndexed { index, name ->
                     addLocalDeclarationBinding(tag, name, index)
@@ -146,7 +146,7 @@ sealed class InliningScope {
     }
 }
 
-class ProgramFragmentInliningScope(
+class ImportInfoFragmentInliningScope private constructor(
     override val fragment: JsProgramFragment
 ) : InliningScope() {
 
@@ -242,11 +242,18 @@ class ProgramFragmentInliningScope(
             }
         }.internalName
     }
+
+    companion object {
+        fun process(fragment: JsProgramFragment, fn: (ImportInfoFragmentInliningScope) -> Unit) {
+            val scope = ImportInfoFragmentInliningScope(fragment)
+            fn(scope)
+            scope.update()
+        }
+    }
 }
 
-class PublicInlineFunctionInliningScope(
+class ImportIntoWrapperInliningScope private constructor(
-    val function: JsFunction,
+    private val wrapperBody: JsBlock,
-    val wrapperBody: JsBlock,
     override val fragment: JsProgramFragment
 ) : InliningScope() {
     private val existingImports = mutableMapOf<String, JsName>().also { map ->
@@ -275,4 +282,12 @@ class PublicInlineFunctionInliningScope(
     override fun update() {
         wrapperBody.statements.addAll(0, additionalStatements)
     }
+
+    companion object {
+        fun process(wrapperBody: JsBlock, fragment: JsProgramFragment, fn: (ImportIntoWrapperInliningScope) -> Unit) {
+            val scope = ImportIntoWrapperInliningScope(wrapperBody, fragment)
+            fn(scope)
+            scope.update()
+        }
+    }
 }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/JsInliner.kt

@@ -9,7 +9,7 @@ import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.synthetic
 import org.jetbrains.kotlin.js.config.JsConfig
-import org.jetbrains.kotlin.js.inline.context.FunctionContext
+import org.jetbrains.kotlin.js.inline.context.FunctionDefinitionLoader
 import org.jetbrains.kotlin.js.inline.context.InliningContext
 
 import org.jetbrains.kotlin.js.translate.general.AstGenerationResult
@@ -21,44 +21,35 @@ class JsInliner(
     val translationResult: AstGenerationResult
 ) {
 
-    val functionContext = FunctionContext(this)
+    val functionDefinitionLoader = FunctionDefinitionLoader(this)
 
-    val cycleReporter = InlinerCycleReporter(trace, functionContext)
+    val cycleReporter = InlinerCycleReporter(trace, functionDefinitionLoader)
 
     fun process() {
         for (fragment in translationResult.newFragments) {
-            process(fragment)
+            ImportInfoFragmentInliningScope.process(fragment) { fragmentScope ->
-        }
+                InlineAstVisitor(this, fragmentScope).accept(fragmentScope.allCode)
-    }
-
-    fun process(fragment: JsProgramFragment) {
-        val fragmentScope = functionContext.scopeForFragment(fragment) ?: return
-
-        fragmentScope.process(fragmentScope.allCode)
-
-        fragmentScope.update()
-    }
-
-    fun process(inlineFn: InlineFunctionDefinition, call: JsInvocation?, containingScope: InliningScope) {
-        cycleReporter.processInlineFunction(inlineFn.fn, call) {
-            val (function, wrapperBody) = inlineFn.fn
-
-            if (wrapperBody != null) {
-                val scope = PublicInlineFunctionInliningScope(function, wrapperBody, containingScope.fragment)
-                scope.process(wrapperBody)
-                scope.update()
-            } else {
-                containingScope.process(function)
             }
         }
     }
 
-    private fun InliningScope.process(node: JsNode) {
+    fun process(inlineFn: InlineFunctionDefinition, call: JsInvocation?, definitionFragment: JsProgramFragment) {
-        InlinerImpl(this@JsInliner, this).accept(node)
+        // Old fragments are already processed
+        if (definitionFragment !in translationResult.newFragments) return
+
+        cycleReporter.processInlineFunction(inlineFn.fn, call) {
+            val wrapperBody = inlineFn.fn.wrapperBody
+
+            checkNotNull(wrapperBody) { "All unprocessed inline functions should be generated with wrappers" }
+
+            ImportIntoWrapperInliningScope.process(wrapperBody, definitionFragment) { scope ->
+                InlineAstVisitor(this, scope).accept(wrapperBody)
+            }
+        }
     }
 
     fun inline(scope: InliningScope, call: JsInvocation, currentStatement: JsStatement?): InlineableResult {
-        val definition = functionContext.getFunctionDefinition(call, scope)
+        val definition = functionDefinitionLoader.getFunctionDefinition(call, scope.fragment)
 
         val function = scope.importFunctionDefinition(definition)
 
@@ -67,7 +58,7 @@ class JsInliner(
         val (inlineableBody, resultExpression) = FunctionInlineMutator.getInlineableCallReplacement(call, function, inliningContext)
 
         // body of inline function can contain call to lambdas that need to be inlined
-        scope.process(inlineableBody)
+        InlineAstVisitor(this, scope).accept<JsNode?>(inlineableBody)
 
         // TODO shouldn't we process the resultExpression qualifier along with the lambda inlining?
         resultExpression?.synthetic = true

js/js.inliner/src/org/jetbrains/kotlin/js/inline/context/FunctionDefinitionLoader.kt

@@ -25,26 +25,19 @@ import org.jetbrains.kotlin.js.inline.util.*
 import org.jetbrains.kotlin.js.translate.context.Namer
 import java.util.HashMap
 
-class FunctionContext(
+class FunctionDefinitionLoader(
-    val inliner: JsInliner
+    private val inliner: JsInliner
 ) {
-    fun getFunctionDefinition(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition {
+    fun getFunctionDefinition(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition {
-        return getFunctionDefinitionImpl(call, scope)!!
+        return getFunctionDefinitionImpl(call, fragment)!!
     }
 
-    fun hasFunctionDefinition(call: JsInvocation, scope: InliningScope): Boolean {
+    fun hasFunctionDefinition(call: JsInvocation, fragment: JsProgramFragment): Boolean {
-        return getFunctionDefinitionImpl(call, scope) != null
+        return getFunctionDefinitionImpl(call, fragment) != null
     }
 
     val functionsByFunctionNodes = HashMap<JsFunction, FunctionWithWrapper>()
 
-    fun scopeForFragment(fragment: JsProgramFragment) = if (fragment in inliner.translationResult.newFragments) {
-        inliningScopeCache.computeIfAbsent(fragment) {
-            loadFragment(fragment)
-            ProgramFragmentInliningScope(fragment)
-        }
-    } else null
-
     /**
      * Gets function definition by invocation.
      *
@@ -75,11 +68,11 @@ class FunctionContext(
      * 5. Qualifier can be JsNameRef with ref to case [3]
      *    in case of local function with closure.
      */
-    private fun getFunctionDefinitionImpl(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
+    private fun getFunctionDefinitionImpl(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition? {
         // Ensure we have the local function information
-        assert(scope.fragment in inliningScopeCache)
+        loadFragment(fragment)
 
-        return lookUpFunctionDirect(call) ?: lookUpFunctionIndirect(call, scope) ?: lookUpFunctionExternal(call, scope.fragment)
+        return lookUpFunctionDirect(call) ?: lookUpFunctionIndirect(call, fragment) ?: lookUpFunctionExternal(call, fragment)
     }
 
     private val functionReader = FunctionReader(inliner.reporter, inliner.config)
@@ -98,7 +91,7 @@ class FunctionContext(
         fragmentInfo[fragment]?.run { accessors[functionTag] }
 
 
-    private fun lookUpFunctionIndirect(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
+    private fun lookUpFunctionIndirect(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition? {
         /** remove ending `()` */
         val callQualifier: JsExpression = if (isCallInvocation(call)) {
             (call.qualifier as JsNameRef).qualifier!!
@@ -107,10 +100,10 @@ class FunctionContext(
         }
 
         call.descriptor?.let { descriptor ->
-            fragmentInfo[scope.fragment]?.let { info ->
+            fragmentInfo[fragment]?.let { info ->
                 info.localAccessors[descriptor]?.let { fn ->
                     return InlineFunctionDefinition(fn, null).also { def ->
-                        inliner.process(def, call, scope)
+                        inliner.process(def, call, fragment)
                     }
                 }
             }
@@ -121,10 +114,10 @@ class FunctionContext(
         return when (qualifier) {
             is JsInvocation -> {
                 tryExtractCallableReference(qualifier) ?: getSimpleName(qualifier)?.let { simpleName ->
-                    lookUpStaticFunction(simpleName, scope.fragment)?.let { if (isFunctionCreator(it.function)) it else null }
+                    lookUpStaticFunction(simpleName, fragment)?.let { if (isFunctionCreator(it.function)) it else null }
                 }
             }
-            is JsNameRef -> lookUpStaticFunction(qualifier.name, scope.fragment)
+            is JsNameRef -> lookUpStaticFunction(qualifier.name, fragment)
 
             // Since we could get functionWithWrapper as a simple function directly from staticRef (which always points on implementation)
             // we should check if we have a known wrapper for it
@@ -132,15 +125,13 @@ class FunctionContext(
             else -> null
         }?.let {
             InlineFunctionDefinition(it, null).also { definition ->
-                if (scope.fragment in inliner.translationResult.newFragments) {
+                if (fragment in inliner.translationResult.newFragments) {
-                    inliner.process(definition, call, scope)
+                    inliner.process(definition, call, fragment)
                 }
             }
         }
     }
 
-    private val inliningScopeCache = mutableMapOf<JsProgramFragment, ProgramFragmentInliningScope>()
-
     private fun loadFragment(fragment: JsProgramFragment) {
         fragmentInfo.computeIfAbsent(fragment) {
             FragmentInfo(
@@ -172,10 +163,8 @@ class FunctionContext(
 
         val definition = InlineFunctionDefinition(fn, tag)
 
-        // Process definition if it is in the new fragments
+        // Make sure definition has it's own inline calls inlined.
-        scopeForFragment(definitionFragment)?.let { definitionScope ->
+        inliner.process(definition, call, definitionFragment)
-            inliner.process(definition, call, definitionScope)
-        }
 
         return definition
     }