JS: fix bugs

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

@@ -70,6 +70,7 @@ class CoroutineTransformer : JsVisitorWithContextImpl() {
 fun transformCoroutines(fragments: List<JsProgramFragment>) {
     val coroutineTransformer = CoroutineTransformer()
     for (fragment in fragments) {
+        fragment.inlinedFunctionWrappers.values.forEach { coroutineTransformer.accept(it) }
         coroutineTransformer.accept(fragment.declarationBlock)
         coroutineTransformer.accept(fragment.initializerBlock)
     }

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

@@ -176,11 +176,21 @@ class FunctionReader(
         return functionCache.get(descriptor).let {
             if (it === NotFoundMarker) null else {
                 val (fn, info) = it as Pair<*, *>
-                renameModules(descriptor, fn as FunctionWithWrapper, info as ModuleInfo, fragment)
+                renameModules(descriptor, (fn as FunctionWithWrapper).deepCopy(), info as ModuleInfo, fragment)
             }
         }
     }
 
+    private fun FunctionWithWrapper.deepCopy(): FunctionWithWrapper {
+        return if (wrapperBody == null) {
+            FunctionWithWrapper(function.deepCopy(), null)
+        } else {
+            val newWrapper = wrapperBody.deepCopy()
+            val newFunction = (newWrapper.statements.last() as JsReturn).expression as JsFunction
+            FunctionWithWrapper(newFunction, newWrapper)
+        }
+    }
+
     private fun renameModules(
         descriptor: CallableDescriptor,
         fn: FunctionWithWrapper,

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

@@ -42,24 +42,32 @@ class InlinerCycleReporter(
         get() = if (namedFunctionsStack.empty()) null else namedFunctionsStack.peek()
 
 
-    fun startFunction(function: JsFunction) {
+    fun <T> withFunction(function: JsFunction, body: () -> T): T {
-        assert(!inProcessFunctions.contains(function)) { "Inliner has revisited function" }
-        inProcessFunctions.add(function)
-
         if (function in functionContext.functionsByFunctionNodes.keys) {
             namedFunctionsStack.push(function)
         }
-    }
 
-    fun endFunction(function: JsFunction) {
+        val result = body()
-        processedFunctions.add(function)
-
-        assert(inProcessFunctions.contains(function))
-        inProcessFunctions.remove(function)
 
         if (!namedFunctionsStack.empty() && namedFunctionsStack.peek() == function) {
             namedFunctionsStack.pop()
         }
+
+        return result
+    }
+
+    fun <T> withInlineFunctionDefinition(function: JsFunction, body: () -> T): T {
+        assert(!inProcessFunctions.contains(function)) { "Inliner has revisited function" }
+        inProcessFunctions.add(function)
+
+        val result = withFunction(function, body)
+
+        processedFunctions.add(function)
+
+        assert(function in inProcessFunctions)
+        inProcessFunctions.remove(function)
+
+        return result
     }
 
 

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

@@ -6,10 +6,12 @@
 package org.jetbrains.kotlin.js.inline
 
 import org.jetbrains.kotlin.js.backend.ast.*
+import org.jetbrains.kotlin.js.backend.ast.metadata.forcedReturnVariable
 import org.jetbrains.kotlin.js.backend.ast.metadata.inlineStrategy
-import org.jetbrains.kotlin.js.inline.context.FunctionContext
+import org.jetbrains.kotlin.js.inline.clean.FunctionPostProcessor
-import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper
+import org.jetbrains.kotlin.js.inline.clean.removeUnusedLocalFunctionDeclarations
 import org.jetbrains.kotlin.js.inline.util.extractFunction
+import org.jetbrains.kotlin.js.inline.util.refreshLabelNames
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 
@@ -56,6 +58,20 @@ class InlinerImpl(
         return super.visit(x, ctx)
     }
 
+    override fun visit(x: JsFunction, ctx: JsContext<*>): Boolean {
+        return jsInliner.cycleReporter.withFunction(x) {
+            super.visit(x, ctx)
+        }
+    }
+
+    override fun endVisit(function: JsFunction, ctx: JsContext<*>) {
+        // Cleanup
+        patchReturnsFromSecondaryConstructor(function)
+        refreshLabelNames(function.body, function.scope)
+        removeUnusedLocalFunctionDeclarations(function)
+        FunctionPostProcessor(function).apply()
+    }
+
 
     override fun endVisit(call: JsInvocation, ctx: JsContext<JsNode>) {
         if (hasToBeInlined(call)) {
@@ -94,4 +110,17 @@ class InlinerImpl(
         val strategy = call.inlineStrategy
         return if (strategy == null || !strategy.isInline) false else jsInliner.functionContext.hasFunctionDefinition(call, scope)
     }
+
+    private fun patchReturnsFromSecondaryConstructor(function: JsFunction) {
+        // Support non-local return from secondary constructor
+        // Returns from secondary constructors should return `$this` object.
+        // TODO This seems brittle
+        function.forcedReturnVariable?.let { returnVariable ->
+            function.body.accept(object : RecursiveJsVisitor() {
+                override fun visitReturn(x: JsReturn) {
+                    x.expression = returnVariable.makeRef()
+                }
+            })
+        }
+    }
 }

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

@@ -25,7 +25,7 @@ sealed class InliningScope {
 
     abstract fun addInlinedDeclaration(tag: String?, declaration: JsStatement)
 
-    abstract fun hasImport(tag: String): JsName?
+    abstract fun hasImport(name: JsName, tag: String): JsName?
 
     abstract fun addImport(tag: String, vars: JsVars)
 
@@ -33,11 +33,11 @@ sealed class InliningScope {
 
     abstract fun update()
 
-    private val publicFunctionCache = mutableMapOf<String, Map<JsName, JsNameRef>>()
+    private val publicFunctionCache = mutableMapOf<String, JsFunction>()
 
-    private val localFunctionCache = mutableMapOf<JsFunction, Map<JsName, JsNameRef>>()
+    private val localFunctionCache = mutableMapOf<JsFunction, JsFunction>()
 
-    private fun computeIfAbsent(tag: String?, function: JsFunction, fn: () -> Map<JsName, JsNameRef>): Map<JsName, JsNameRef> {
+    private fun computeIfAbsent(tag: String?, function: JsFunction, fn: () -> JsFunction): JsFunction {
         if (tag == null) return localFunctionCache.computeIfAbsent(function) { fn() }
 
         return publicFunctionCache.computeIfAbsent(tag) { fn() }
@@ -46,7 +46,9 @@ sealed class InliningScope {
     fun importFunctionDefinition(definition: InlineFunctionDefinition): JsFunction {
         // Apparently we should avoid this trick when we implement fair support for crossinline
         // That's because crossinline lambdas inline into the declaration block and specialize those.
-        val replacements = computeIfAbsent(definition.tag, definition.fn.function) {
+
+        // TODO cache the function itself instead
+        val result = computeIfAbsent(definition.tag, definition.fn.function) {
             val newReplacements = HashMap<JsName, JsNameRef>()
 
             val copiedStatements = ArrayList<JsStatement>()
@@ -63,15 +65,18 @@ sealed class InliningScope {
                             val tag = getImportTag(statement)
                             if (tag != null) {
                                 val name = statement.vars[0].name
-                                val existingName = name.localAlias ?: hasImport(tag) ?: JsScope.declareTemporaryName(name.ident).also {
+                                val existingName = hasImport(name, tag) ?: JsScope.declareTemporaryName(name.ident).also {
                                     it.copyMetadataFrom(name)
                                     importStatements[statement] = tag
+                                    copiedStatements.add(statement)
                                 }
 
                                 if (name !== existingName) {
                                     val replacement = JsAstUtils.pureFqn(existingName, null)
                                     newReplacements[name] = replacement
                                 }
+
+                                return@forEach
                             }
                         }
 
@@ -109,22 +114,24 @@ sealed class InliningScope {
                 }
             }
 
-            newReplacements
+            val result = definition.fn.function.deepCopy()
-        }
 
-        val paramMap = definition.fn.function.parameters.associate {
+            replaceNames(result, newReplacements)
+
+            result.body = transformSpecialFunctionsToCoroutineMetadata(result.body)
+
+            result
+        }.deepCopy()
+
+        // Copy parameter JsName's
+        val paramMap = result.parameters.associate {
             val alias = JsScope.declareTemporaryName(it.name.ident)
             alias.copyMetadataFrom(it.name)
             it.name to JsAstUtils.pureFqn(alias, null)
         }
 
-        val result = definition.fn.function.deepCopy()
-
-        replaceNames(result, replacements)
         replaceNames(result, paramMap)
 
-        result.body = transformSpecialFunctionsToCoroutineMetadata(result.body)
-
         return result
     }
 }
@@ -146,7 +153,7 @@ class ProgramFragmentInliningScope(
         fragment.declarationBlock.statements.addAll(0, additionalDeclarations)
 
         // post-processing
-        val block = JsBlock(fragment.declarationBlock, fragment.initializerBlock, fragment.exportBlock)
+        val block = JsBlock(JsBlock(fragment.inlinedFunctionWrappers.values.toList()), fragment.declarationBlock, fragment.initializerBlock, fragment.exportBlock)
         fragment.tests?.let { block.statements.add(it) }
         fragment.mainFunction?.let { block.statements.add(it) }
 
@@ -169,7 +176,8 @@ class ProgramFragmentInliningScope(
         }
     }
 
-    override fun hasImport(tag: String): JsName? = existingBindings[tag]
+    // TODO !!!!! This localAlias thing will get copied, inlined, and lost during IC =(
+    override fun hasImport(name: JsName, tag: String): JsName? = name.localAlias ?: existingBindings[tag]
 
     override fun addImport(tag: String, vars: JsVars) {
         val name = vars.vars[0].name
@@ -229,7 +237,7 @@ class PublicInlineFunctionInliningScope(
         additionalStatements.add(declaration)
     }
 
-    override fun hasImport(tag: String): JsName? {
+    override fun hasImport(name: JsName, tag: String): JsName? {
         return null // TODO
     }
 

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

@@ -7,15 +7,10 @@ package org.jetbrains.kotlin.js.inline
 
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.js.backend.ast.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.forcedReturnVariable
 import org.jetbrains.kotlin.js.backend.ast.metadata.synthetic
 import org.jetbrains.kotlin.js.config.JsConfig
-import org.jetbrains.kotlin.js.inline.clean.FunctionPostProcessor
-import org.jetbrains.kotlin.js.inline.clean.removeUnusedLocalFunctionDeclarations
 import org.jetbrains.kotlin.js.inline.context.FunctionContext
 import org.jetbrains.kotlin.js.inline.context.InliningContext
-import org.jetbrains.kotlin.js.inline.util.refreshLabelNames
-import org.jetbrains.kotlin.js.translate.declaration.transformSpecialFunctionsToCoroutineMetadata
 
 import org.jetbrains.kotlin.js.translate.general.AstGenerationResult
 
@@ -71,22 +66,15 @@ class JsInliner(
         if (cycleReporter.shouldProcess(inlineFn.fn, call)) {
             val (function, wrapperBody) = inlineFn.fn
 
-            cycleReporter.startFunction(function)
+            cycleReporter.withInlineFunctionDefinition(function) {
-
+                if (wrapperBody != null) {
-            if (wrapperBody != null) {
+                    val scope = PublicInlineFunctionInliningScope(function, wrapperBody, containingScope.fragment)
-                val scope = PublicInlineFunctionInliningScope(function, wrapperBody, containingScope.fragment)
+                    scope.process(wrapperBody)
-                scope.process(wrapperBody)
+                    scope.update()
-                scope.update()
+                } else {
-            } else {
+                    containingScope.process(function)
-                containingScope.process(function)
+                }
             }
-
-            // Cleanup
-            refreshLabelNames(function.body, function.scope)
-            removeUnusedLocalFunctionDeclarations(function)
-            FunctionPostProcessor(function).apply()
-
-            cycleReporter.endFunction(function)
         }
     }
 
@@ -109,25 +97,10 @@ class JsInliner(
             // body of inline function can contain call to lambdas that need to be inlined
             scope.process(inlineableBody)
 
-            patchReturnsFromSecondaryConstructor(inlineableBody)
-
             // TODO shouldn't we process the resultExpression qualifier along with the lambda inlining?
             resultExpression?.synthetic = true
 
             InlineableResult(JsBlock(inliningContext.previousStatements + inlineableBody), resultExpression)
         }
     }
-
-    private fun patchReturnsFromSecondaryConstructor(inlineableBody: JsStatement) {
-        // Support non-local return from secondary constructor
-        // Returns from secondary constructors should return `$this` object.
-        // TODO This seems brittle
-        cycleReporter.currentNamedFunction?.forcedReturnVariable?.let { returnVariable ->
-            inlineableBody.accept(object : RecursiveJsVisitor() {
-                override fun visitReturn(x: JsReturn) {
-                    x.expression = returnVariable.makeRef()
-                }
-            })
-        }
-    }
 }

js/js.serializer/src/org/jetbrains/kotlin/serialization/js/ast/JsAstSerializer.kt

@@ -471,7 +471,7 @@ class JsAstSerializer(private val pathResolver: (File) -> String) {
     private fun serialize(module: JsImportedModule): JsAstProtoBuf.JsImportedModule {
         val moduleBuilder = JsAstProtoBuf.JsImportedModule.newBuilder()
         moduleBuilder.externalName = module.externalName
-        moduleBuilder.internalName = serialize(module.externalName)
+        moduleBuilder.internalName = serialize(module.internalName)
         module.plainReference?.let {
             moduleBuilder.plainReference = serialize(it)
         }