JS: postpone JS AST merging

2018-11-15 15:58:38 +03:00
parent ec276dbea4
commit c65383fa3f
12 changed files with 613 additions and 643 deletions
@@ -16,4 +16,4 @@

 package org.jetbrains.kotlin.js.backend.ast

-class JsNameBinding(val key: String, var name: JsName)
+data class JsNameBinding(val key: String, var name: JsName)
@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.js.coroutine
 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.clean.LabeledBlockToDoWhileTransformation
 import org.jetbrains.kotlin.js.translate.declaration.transformCoroutineMetadataToSpecialFunctions
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
@@ -66,3 +67,11 @@ class CoroutineTransformer : JsVisitorWithContextImpl() {
        return super.visit(x, ctx)
    }
 }
+
+fun transformCoroutines(fragments: List<JsProgramFragment>) {
+    val coroutineTransformer = CoroutineTransformer()
+    for (fragment in fragments) {
+        coroutineTransformer.accept(fragment.declarationBlock)
+        coroutineTransformer.accept(fragment.initializerBlock)
+    }
+}
@@ -5,11 +5,6 @@

 package org.jetbrains.kotlin.js.inline

-import org.jetbrains.kotlin.backend.common.*
-import org.jetbrains.kotlin.config.*
-import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
-import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.js.backend.ast.*
@@ -18,7 +13,6 @@ import org.jetbrains.kotlin.js.config.JsConfig
 import org.jetbrains.kotlin.js.inline.clean.*
 import org.jetbrains.kotlin.js.inline.context.FunctionContext
 import org.jetbrains.kotlin.js.inline.context.InliningContext
-import org.jetbrains.kotlin.js.inline.context.NamingContext
 import org.jetbrains.kotlin.js.inline.util.*
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.resolve.inline.InlineStrategy
@@ -28,6 +22,7 @@ import java.util.*
 import org.jetbrains.kotlin.js.inline.util.getImportTag
 import org.jetbrains.kotlin.js.inline.util.IdentitySet
 import org.jetbrains.kotlin.js.translate.declaration.transformSpecialFunctionsToCoroutineMetadata
+import org.jetbrains.kotlin.js.translate.general.AstGenerationResult
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils.flattenStatement
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils.pureFqn

@@ -37,17 +32,17 @@ class JsInliner private constructor(
    private val accessors: Map<String, FunctionWithWrapper>,
    private val inverseNameBindings: Map<JsName, String>,
    private val functionReader: FunctionReader,
-    private val trace: DiagnosticSink
-) : JsVisitorWithContextImpl() {
+    private val trace: DiagnosticSink,
+    private val moduleMap: Map<JsName, JsImportedModule>
+) {
    private val namedFunctionsSet: MutableSet<JsFunction> = functions.values.mapTo(IdentitySet()) { it.function }
-    private val inliningContexts = Stack<JsInliningContext>()
+
    private val processedFunctions = IdentitySet<JsFunction>()
    private val inProcessFunctions = IdentitySet<JsFunction>()
-    private var existingImports: MutableMap<String, JsName> = HashMap()
-    private var statementContextForInline: JsContext<JsStatement>? = null

    private val functionsByWrapperNodes = HashMap<JsBlock, FunctionWithWrapper>()
    private val functionsByFunctionNodes = HashMap<JsFunction, FunctionWithWrapper>()
+
    init {
        (functions.values.asSequence() + accessors.values.asSequence()).forEach { f ->
            functionsByFunctionNodes[f.function] = f
@@ -60,442 +55,384 @@ class JsInliner private constructor(
    // these are needed for error reporting, when inliner detects cycle
    private val namedFunctionsStack = Stack<JsFunction>()
    private val inlineCallInfos = LinkedList<JsCallInfo>()
-    private val canBeExtractedByInliner: (JsNode) -> Boolean = { node -> node is JsInvocation && hasToBeInlined(node) }
-    private var inlineFunctionDepth: Int = 0
-
-    private val replacementsInducedByWrappers = HashMap<JsWrapperKey, Map<JsName, JsNameRef>>()
-
-    private var existingNameBindings: MutableMap<JsName, String> = HashMap()
-
-    private val additionalNameBindings = ArrayList<JsNameBinding>()
-
-    private val inlinedModuleAliases = HashSet<JsName>()
-
-    private val inliningContext: JsInliningContext
-        get() = inliningContexts.peek()
-
-    private val functionContext: FunctionContext
-        get() = inliningContext.functionContext

    private val currentNamedFunction: JsFunction?
        get() = if (namedFunctionsStack.empty()) null else namedFunctionsStack.peek()

-    private fun addInlinedModules(fragment: JsProgramFragment, moduleMap: Map<JsName, JsImportedModule>) {
-        val localMap = buildModuleMap(listOf(fragment)).keys
-        for (inlinedModuleName in inlinedModuleAliases) {
-            if (!localMap.contains(inlinedModuleName)) {
-                fragment.importedModules.add(moduleMap[inlinedModuleName]!!)
-            }
-        }
-    }
+    private inner class InlinerImpl(
+        val existingNameBindings: MutableMap<JsName, String>,
+        val existingImports: MutableMap<String, JsName>,
+        val inlineFunctionDepth: Int,
+        val addPrevious: (JsStatement) -> Unit
+    ) : JsVisitorWithContextImpl() {

-    private fun processImportStatement(statement: JsStatement) {
-        if (statement is JsVars) {
-            val tag = getImportTag(statement)
-            if (tag != null) {
-                existingImports[tag] = statement.vars[0].name
-            }
-        }
-    }
+        val replacementsInducedByWrappers = HashMap<JsFunction, Map<JsName, JsNameRef>>()

-    override fun visit(function: JsFunction, context: JsContext<*>): Boolean {
-        val functionWithWrapper = functionsByFunctionNodes[function]
-        if (functionWithWrapper != null) {
-            visit(functionWithWrapper)
-            return false
-        } else {
-            if (statementContextForInline == null) {
-                statementContextForInline = lastStatementLevelContext
-                startFunction(function)
-                val result = super.visit(function, context)
-                statementContextForInline = null
-                return result
+        val inlinedModuleAliases = HashSet<JsName>()
+
+        val additionalNameBindings = ArrayList<JsNameBinding>()
+
+        override fun visit(function: JsFunction, context: JsContext<*>): Boolean {
+            val functionWithWrapper = functionsByFunctionNodes[function]
+            if (functionWithWrapper != null) {
+                visit(functionWithWrapper)
+                return false
            } else {
                startFunction(function)
                return super.visit(function, context)
            }
        }
-    }

-    override fun endVisit(function: JsFunction, context: JsContext<*>) {
-        super.endVisit(function, context)
-        if (!functionsByFunctionNodes.containsKey(function)) {
-            endFunction(function)
-        }
-    }
-
-    private fun startFunction(function: JsFunction) {
-        inliningContexts.push(JsInliningContext(statementContextForInline!!))
-
-        assert(!inProcessFunctions.contains(function)) { "Inliner has revisited function" }
-        inProcessFunctions.add(function)
-
-        if (namedFunctionsSet.contains(function)) {
-            namedFunctionsStack.push(function)
-        }
-    }
-
-    private fun endFunction(function: JsFunction) {
-        refreshLabelNames(function.body, function.scope)
-
-        removeUnusedLocalFunctionDeclarations(function)
-        processedFunctions.add(function)
-
-        FunctionPostProcessor(function).apply()
-
-        assert(inProcessFunctions.contains(function))
-        inProcessFunctions.remove(function)
-
-        inliningContexts.pop()
-
-        if (!namedFunctionsStack.empty() && namedFunctionsStack.peek() == function) {
-            namedFunctionsStack.pop()
-        }
-    }
-
-    override fun visit(x: JsBlock, ctx: JsContext<*>): Boolean {
-        val functionWithWrapper = functionsByWrapperNodes[x]
-        if (functionWithWrapper != null) {
-            visit(functionWithWrapper)
-            return false
-        }
-        return super.visit(x, ctx)
-    }
-
-    private fun visit(functionWithWrapper: FunctionWithWrapper) {
-        val oldContextForInline = statementContextForInline
-        val oldExistingImports = existingImports
-        val oldInlineFunctionDepth = inlineFunctionDepth
-
-        val innerContext = ListContext<JsStatement>()
-
-        val wrapperBody = functionWithWrapper.wrapperBody
-        var statements: MutableList<JsStatement>? = null
-        if (wrapperBody != null) {
-            existingImports = HashMap()
-            statementContexts.push(innerContext)
-            statementContextForInline = innerContext
-            inlineFunctionDepth++
-
-            for (statement in wrapperBody.statements) {
-                processImportStatement(statement)
+        override fun endVisit(function: JsFunction, context: JsContext<*>) {
+            super.endVisit(function, context)
+            if (!functionsByFunctionNodes.containsKey(function)) {
+                endFunction(function)
            }
-            statements = wrapperBody.statements
-            if (!statements.isEmpty() && statements[statements.size - 1] is JsReturn) {
-                statements = statements.subList(0, statements.size - 1)
+        }
+
+
+        private fun startFunction(function: JsFunction) {
+            assert(!inProcessFunctions.contains(function)) { "Inliner has revisited function" }
+            inProcessFunctions.add(function)
+
+            if (namedFunctionsSet.contains(function)) {
+                namedFunctionsStack.push(function)
+            }
+        }
+
+        private fun endFunction(function: JsFunction) {
+            refreshLabelNames(function.body, function.scope)
+
+            removeUnusedLocalFunctionDeclarations(function)
+            processedFunctions.add(function)
+
+            FunctionPostProcessor(function).apply()
+
+            assert(inProcessFunctions.contains(function))
+            inProcessFunctions.remove(function)
+
+            if (!namedFunctionsStack.empty() && namedFunctionsStack.peek() == function) {
+                namedFunctionsStack.pop()
+            }
+        }
+
+        override fun visit(x: JsBlock, ctx: JsContext<*>): Boolean {
+            val functionWithWrapper = functionsByWrapperNodes[x]
+            if (functionWithWrapper != null) {
+                visit(functionWithWrapper)
+                return false
+            }
+            return super.visit(x, ctx)
+        }
+
+        private fun visit(functionWithWrapper: FunctionWithWrapper) {
+            startFunction(functionWithWrapper.function)
+
+            val wrapperBody = functionWithWrapper.wrapperBody
+            if (wrapperBody != null) {
+//                statementContexts.push(ListContext())
+
+                val existingImports = HashMap<String, JsName>()
+
+                for (statement in wrapperBody.statements) {
+                    if (statement is JsVars) {
+                        val tag = getImportTag(statement)
+                        if (tag != null) {
+                            existingImports[tag] = statement.vars[0].name
+                        }
+                    }
+                }
+
+                val additionalStatements = mutableListOf<JsStatement>()
+                val innerInliner = InlinerImpl(existingNameBindings, existingImports, inlineFunctionDepth + 1) {
+                    additionalStatements.add(it)
+                }
+                for (statement in wrapperBody.statements) {
+                    if (statement !is JsReturn) {
+                        innerInliner.acceptStatement(statement)
+                    } else {
+                        innerInliner.accept((statement.expression as JsFunction).body)
+                    }
+                }
+
+//                innerInliner.acceptStatement(wrapperBody)
+                // TODO keep order
+                wrapperBody.statements.addAll(0, additionalStatements)
+
+//                statementContexts.pop()
+            } else {
+                accept(functionWithWrapper.function.body)
            }

-            innerContext.traverse(statements)
-            statementContexts.pop()
-        } else {
-            if (statementContextForInline == null) statementContextForInline = lastStatementLevelContext
+            endFunction(functionWithWrapper.function)
        }

-        startFunction(functionWithWrapper.function)
+        override fun visit(call: JsInvocation, context: JsContext<*>): Boolean {
+            if (!hasToBeInlined(call)) return true

-        val block = JsBlock(functionWithWrapper.function.body)
-        innerContext.traverse(block.statements)
-        functionWithWrapper.function.body.traverse(this, innerContext)
-
-        endFunction(functionWithWrapper.function)
-
-        statements?.addAll(block.statements.subList(0, block.statements.size - 1))
-
-        statementContextForInline = oldContextForInline
-        existingImports = oldExistingImports
-        inlineFunctionDepth = oldInlineFunctionDepth
-    }
-
-    override fun visit(call: JsInvocation, context: JsContext<*>): Boolean {
-        if (!hasToBeInlined(call)) return true
-
-        val containingFunction = currentNamedFunction
-
-        if (containingFunction != null) {
-            inlineCallInfos.add(JsCallInfo(call, containingFunction))
-        }
-
-        val definition = functionContext.getFunctionDefinition(call)
-
-        if (inProcessFunctions.contains(definition.function)) {
-            reportInlineCycle(call, definition.function)
-        } else if (!processedFunctions.contains(definition.function)) {
-            for (i in 0 until call.arguments.size) {
-                val argument = call.arguments[i]
-                call.arguments[i] = accept(argument)
+            currentNamedFunction?.let {
+                inlineCallInfos.add(JsCallInfo(call, it))
            }
-            visit(definition)
-            return false
+
+            val definition = functionContext.getFunctionDefinition(call)
+
+            if (inProcessFunctions.contains(definition.function)) {
+                reportInlineCycle(call, definition.function)
+            } else if (!processedFunctions.contains(definition.function)) {
+                for (i in 0 until call.arguments.size) {
+                    val argument = call.arguments[i]
+                    call.arguments[i] = accept(argument)
+                }
+                visit(definition)
+                return false
+            }
+
+            return true
        }

-        return true
-    }
-
-    override fun endVisit(x: JsInvocation, ctx: JsContext<JsNode>) {
-        if (hasToBeInlined(x)) {
-            inline(x, ctx)
+        override fun endVisit(x: JsInvocation, ctx: JsContext<JsNode>) {
+            if (hasToBeInlined(x)) {
+                inline(x, ctx)
+            }
+            if (!inlineCallInfos.isEmpty()) {
+                if (inlineCallInfos.last.call == x) {
+                    inlineCallInfos.removeLast()
+                }
+            }
        }

-        var lastCallInfo: JsCallInfo? = null
-
-        if (!inlineCallInfos.isEmpty()) {
-            lastCallInfo = inlineCallInfos.last
-        }
-
-        if (lastCallInfo != null && lastCallInfo.call == x) {
-            inlineCallInfos.removeLast()
-        }
-    }
-
-    override fun endVisit(x: JsExpressionStatement, ctx: JsContext<*>) {
-        val e = x.expression
-        if (e is JsBinaryOperation) {
-            if (e.operator == JsBinaryOperator.ASG) {
-                e.arg2?.let { argument2 ->
+        override fun endVisit(x: JsExpressionStatement, ctx: JsContext<*>) {
+            val e = x.expression
+            if (e is JsBinaryOperation) {
+                if (e.operator == JsBinaryOperator.ASG) {
+                    e.arg2?.let { argument2 ->
                    val splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(argument2)
                    if (splitSuspendInlineFunction != null) {
                        e.arg2 = splitSuspendInlineFunction
                    }
                }
+            }}
+
+            super.endVisit(x, ctx)
+        }
+
+        override fun endVisit(x: JsVars.JsVar, ctx: JsContext<*>) {
+            val initExpression = x.initExpression ?: return
+
+            val splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(initExpression)
+            if (splitSuspendInlineFunction != null) {
+                x.initExpression = splitSuspendInlineFunction
            }
        }

-        super.endVisit(x, ctx)
-    }
+        private fun splitExportedSuspendInlineFunctionDeclarations(expression: JsExpression): JsFunction? {
+            val inlineMetadata = InlineMetadata.decompose(expression)
+            if (inlineMetadata != null) {
+                val (originalFunction, wrapperBody) = inlineMetadata.function
+                if (originalFunction.coroutineMetadata != null) {
+                    val statementContext = lastStatementLevelContext

-    override fun endVisit(x: JsVars.JsVar, ctx: JsContext<*>) {
-        val initExpression = x.initExpression ?: return
+                    // This function will be exported to JS
+                    val function = originalFunction.deepCopy()

-        val splitSuspendInlineFunction = splitExportedSuspendInlineFunctionDeclarations(initExpression)
-        if (splitSuspendInlineFunction != null) {
-            x.initExpression = splitSuspendInlineFunction
-        }
-    }
+                    // Original function should be not be transformed into a state machine
+                    originalFunction.setName(null)
+                    originalFunction.coroutineMetadata = null
+                    originalFunction.isInlineableCoroutineBody = true
+                    if (wrapperBody != null) {
+                        // Extract local declarations
+                        applyWrapper(wrapperBody, function, originalFunction)
+                    }

-    private fun splitExportedSuspendInlineFunctionDeclarations(expression: JsExpression): JsFunction? {
-        val inlineMetadata = InlineMetadata.decompose(expression)
-        if (inlineMetadata != null) {
-            val (originalFunction, wrapperBody) = inlineMetadata.function
-            if (originalFunction.coroutineMetadata != null) {
-                val statementContext = lastStatementLevelContext
+                    // Keep the `defineInlineFunction` for the inliner to find
+                    statementContext.addNext(expression.makeStmt())

-                // This function will be exported to JS
-                val function = originalFunction.deepCopy()
-
-                // Original function should be not be transformed into a state machine
-                originalFunction.setName(null)
-                originalFunction.coroutineMetadata = null
-                originalFunction.isInlineableCoroutineBody = true
-                if (wrapperBody != null) {
-                    // Extract local declarations
-                    applyWrapper(wrapperBody, function, originalFunction, JsInliningContext(statementContext))
+                    // Return the function body to be used without inlining.
+                    return function
                }
-
-                // Keep the `defineInlineFunction` for the inliner to find
-                statementContext.addNext(expression.makeStmt())
-
-                // Return the function body to be used without inlining.
-                return function
            }
+            return null
        }
-        return null
-    }

-    override fun doAcceptStatementList(statements: MutableList<JsStatement>) {
-        // at top level of js ast, contexts stack can be empty,
-        // but there is no inline calls anyway
-        if (!inliningContexts.isEmpty()) {
+        override fun doAcceptStatementList(statements: MutableList<JsStatement>) {
            var i = 0

            while (i < statements.size) {
-                val additionalStatements = ExpressionDecomposer.preserveEvaluationOrder(statements[i], canBeExtractedByInliner)
+                val additionalStatements =
+                    ExpressionDecomposer.preserveEvaluationOrder(statements[i]) { node -> node is JsInvocation && hasToBeInlined(node) }
                statements.addAll(i, additionalStatements)
                i += additionalStatements.size + 1
            }
+
+            super.doAcceptStatementList(statements)
        }

-        super.doAcceptStatementList(statements)
-    }
+        private fun inline(call: JsInvocation, context: JsContext<JsNode>) {
+            var functionWithWrapper = functionContext.getFunctionDefinition(call)

-    private fun inline(call: JsInvocation, context: JsContext<JsNode>) {
-        val callDescriptor = call.descriptor
-        if (isSuspendWithCurrentContinuation(
-                callDescriptor,
-                config.configuration.languageVersionSettings
-            )
-        ) {
-            inlineSuspendWithCurrentContinuation(call, context)
-            return
-        }
-
-        val inliningContext = inliningContext
-        var functionWithWrapper = inliningContext.functionContext.getFunctionDefinition(call)
-
-        // 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
-        functionsByFunctionNodes[functionWithWrapper.function]?.let {
-            functionWithWrapper = it
-        }
-
-        val function = functionWithWrapper.function.deepCopy()
-        function.body = transformSpecialFunctionsToCoroutineMetadata(function.body)
-        if (functionWithWrapper.wrapperBody != null) {
-            applyWrapper(functionWithWrapper.wrapperBody!!, function, functionWithWrapper.function, inliningContext)
-        }
-        val inlineableResult = FunctionInlineMutator.getInlineableCallReplacement(call, function, inliningContext)
-
-        val inlineableBody = inlineableResult.inlineableBody
-        var resultExpression = inlineableResult.resultExpression
-        val statementContext = inliningContext.statementContext
-        // body of inline function can contain call to lambdas that need to be inlined
-        val inlineableBodyWithLambdasInlined = accept(inlineableBody)
-        assert(inlineableBody === inlineableBodyWithLambdasInlined)
-
-        // Support non-local return from secondary constructor
-        // Returns from secondary constructors should return `$this` object.
-        val currentFunction = currentNamedFunction
-        if (currentFunction != null) {
-            val returnVariable = currentFunction.forcedReturnVariable
-            if (returnVariable != null) {
-                inlineableBody.accept(object : RecursiveJsVisitor() {
-                    override fun visitReturn(x: JsReturn) {
-                        x.expression = returnVariable.makeRef()
-                    }
-                })
+            // 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
+            functionsByFunctionNodes[functionWithWrapper.function]?.let {
+                functionWithWrapper = it
            }
-        }

-        statementContext.addPrevious(flattenStatement(inlineableBody))
+            val function = functionWithWrapper.function.deepCopy()
+            function.body = transformSpecialFunctionsToCoroutineMetadata(function.body)
+            if (functionWithWrapper.wrapperBody != null) {
+                applyWrapper(functionWithWrapper.wrapperBody!!, function, functionWithWrapper.function)
+            }

-        /*
-         * Assumes, that resultExpression == null, when result is not needed.
-         * @see FunctionInlineMutator.isResultNeeded()
-         */
-        if (resultExpression == null) {
-            statementContext.removeMe()
-            return
-        }
+            val inliningContext = InliningContext(lastStatementLevelContext)

-        resultExpression = accept(resultExpression)
-        resultExpression.synthetic = true
-        context.replaceMe(resultExpression)
-    }
+            val inlineableResult = FunctionInlineMutator.getInlineableCallReplacement(call, function, inliningContext)

-    private fun applyWrapper(
-        wrapper: JsBlock, function: JsFunction, originalFunction: JsFunction,
-        inliningContext: InliningContext
-    ) {
-        val key = JsWrapperKey(inliningContext.statementContextBeforeCurrentFunction, originalFunction)
+            val inlineableBody = inlineableResult.inlineableBody
+            var resultExpression = inlineableResult.resultExpression
+            val statementContext = inliningContext.statementContext
+            // body of inline function can contain call to lambdas that need to be inlined
+            val inlineableBodyWithLambdasInlined = accept(inlineableBody)
+            assert(inlineableBody === inlineableBodyWithLambdasInlined)

-        // Apparently we should avoid this trick when we implement fair support for crossinline
-        val replacements = replacementsInducedByWrappers.computeIfAbsent(key) { k ->
-            val ctx = k.context
-
-            val newReplacements = HashMap<JsName, JsNameRef>()
-
-            val copiedStatements = ArrayList<JsStatement>()
-            wrapper.statements.asSequence()
-                .filterNot { it is JsReturn }
-                .map { it.deepCopy() }
-                .forEach { statement ->
-                    if (inlineFunctionDepth == 0) {
-                        replaceExpressionsWithLocalAliases(statement)
-                    }
-
-                    if (statement is JsVars) {
-                        val tag = getImportTag(statement)
-                        if (tag != null) {
-                            val name = statement.vars[0].name
-                            var existingName: JsName? = if (inlineFunctionDepth == 0) name.localAlias else null
-                            if (existingName == null) {
-                                existingName = existingImports.computeIfAbsent(tag) {
-                                    copiedStatements.add(statement)
-                                    val alias = JsScope.declareTemporaryName(name.ident)
-                                    alias.copyMetadataFrom(name)
-                                    newReplacements[name] = pureFqn(alias, null)
-                                    alias
-                                }
-                            }
-
-                            if (name !== existingName) {
-                                val replacement = pureFqn(existingName, null)
-                                newReplacements[name] = replacement
-                            }
-
-                            return@forEach
+            // Support non-local return from secondary constructor
+            // Returns from secondary constructors should return `$this` object.
+            val currentFunction = currentNamedFunction
+            if (currentFunction != null) {
+                val returnVariable = currentFunction.forcedReturnVariable
+                if (returnVariable != null) {
+                    inlineableBody.accept(object : RecursiveJsVisitor() {
+                        override fun visitReturn(x: JsReturn) {
+                            x.expression = returnVariable.makeRef()
                        }
+                    })
+                }
+            }
+
+            statementContext.addPrevious(flattenStatement(inlineableBody))
+
+            /*
+             * Assumes, that resultExpression == null, when result is not needed.
+             * @see FunctionInlineMutator.isResultNeeded()
+             */
+            if (resultExpression == null) {
+                statementContext.removeMe()
+                return
+            }
+
+            resultExpression = accept(resultExpression)
+            resultExpression.synthetic = true
+            context.replaceMe(resultExpression)
+        }
+
+
+        private fun applyWrapper(
+            wrapper: JsBlock, function: JsFunction, originalFunction: JsFunction
+        ) {
+
+            // TODO Decrypt the comment below
+            // Apparently we should avoid this trick when we implement fair support for crossinline
+            val replacements = replacementsInducedByWrappers.computeIfAbsent(originalFunction) { k ->
+                val newReplacements = HashMap<JsName, JsNameRef>()
+
+                val copiedStatements = ArrayList<JsStatement>()
+                wrapper.statements.asSequence()
+                    .filterNot { it is JsReturn }
+                    .map { it.deepCopy() }
+                    .forEach { statement ->
+                        if (inlineFunctionDepth == 0) {
+                            replaceExpressionsWithLocalAliases(statement)
+                        }
+
+                        if (statement is JsVars) {
+                            val tag = getImportTag(statement)
+                            if (tag != null) {
+                                val name = statement.vars[0].name
+                                var existingName: JsName? = if (inlineFunctionDepth == 0) name.localAlias else null
+                                if (existingName == null) {
+                                    existingName = existingImports.computeIfAbsent(tag) {
+                                        copiedStatements.add(statement)
+                                        val alias = JsScope.declareTemporaryName(name.ident)
+                                        alias.copyMetadataFrom(name)
+                                        newReplacements[name] = pureFqn(alias, null)
+                                        alias
+                                    }
+                                }
+
+                                if (name !== existingName) {
+                                    val replacement = pureFqn(existingName, null)
+                                    newReplacements[name] = replacement
+                                }
+
+                                return@forEach
+                            }
+                        }
+
+                        copiedStatements.add(statement)
                    }

-                    copiedStatements.add(statement)
+                val definedNames = copiedStatements.asSequence()
+                    .flatMap { node -> collectDefinedNamesInAllScopes(node).asSequence() }
+                    .filter { name -> !newReplacements.containsKey(name) }
+                    .toSet()
+                for (name in definedNames) {
+                    val alias = JsScope.declareTemporaryName(name.ident)
+                    alias.copyMetadataFrom(name)
+                    val replacement = pureFqn(alias, null)
+                    newReplacements[name] = replacement
                }

-            val definedNames = copiedStatements.asSequence()
-                .flatMap { node -> collectDefinedNamesInAllScopes(node).asSequence() }
-                .filter { name -> !newReplacements.containsKey(name) }
-                .toSet()
-            for (name in definedNames) {
-                val alias = JsScope.declareTemporaryName(name.ident)
-                alias.copyMetadataFrom(name)
-                val replacement = pureFqn(alias, null)
-                newReplacements[name] = replacement
-            }
-
-            for (statement in copiedStatements) {
-                ctx.addPrevious(accept(replaceNames(statement, newReplacements)))
-            }
-
-            for ((key, value) in collectNamedFunctions(JsBlock(copiedStatements))) {
-                if (key.staticRef is JsFunction) {
-                    key.staticRef = value
+                for (statement in copiedStatements) {
+                    addPrevious(accept(replaceNames(statement, newReplacements)))
                }
+
+                for ((key, value) in collectNamedFunctions(JsBlock(copiedStatements))) {
+                    if (key.staticRef is JsFunction) {
+                        key.staticRef = value
+                    }
+                }
+
+                newReplacements
            }

-            newReplacements
-        }
+            replaceNames(function, replacements)

-        replaceNames(function, replacements)
-
-        // Copy nameBinding's for inlined localAlias'es
-        for (nameRef in replacements.values) {
-            val name = nameRef.name
-            if (name != null && !existingNameBindings.containsKey(name)) {
-                val tag = inverseNameBindings[name]
-                if (tag != null) {
-                    existingNameBindings[name] = tag
-                    additionalNameBindings.add(JsNameBinding(tag, name))
+            // Copy nameBinding's for inlined localAlias'es
+            for (nameRef in replacements.values) {
+                val name = nameRef.name
+                if (name != null && !existingNameBindings.containsKey(name)) {
+                    val tag = inverseNameBindings[name]
+                    if (tag != null) {
+                        existingNameBindings[name] = tag
+                        additionalNameBindings.add(JsNameBinding(tag, name))
+                    }
                }
            }
        }
-    }

-    private fun replaceExpressionsWithLocalAliases(statement: JsStatement) {
-        object : JsVisitorWithContextImpl() {
-            override fun endVisit(x: JsNameRef, ctx: JsContext<JsNode>) {
-                replaceIfNecessary(x, ctx)
-            }
-
-            override fun endVisit(x: JsArrayAccess, ctx: JsContext<JsNode>) {
-                replaceIfNecessary(x, ctx)
-            }
-
-            private fun replaceIfNecessary(expression: JsExpression, ctx: JsContext<JsNode>) {
-                val alias = expression.localAlias
-                if (alias != null) {
-                    ctx.replaceMe(alias.makeRef())
-                    inlinedModuleAliases.add(alias)
+        private fun replaceExpressionsWithLocalAliases(statement: JsStatement) {
+            object : JsVisitorWithContextImpl() {
+                override fun endVisit(x: JsNameRef, ctx: JsContext<JsNode>) {
+                    replaceIfNecessary(x, ctx)
                }
-            }

-        }.accept(statement)
-    }
+                override fun endVisit(x: JsArrayAccess, ctx: JsContext<JsNode>) {
+                    replaceIfNecessary(x, ctx)
+                }

-    private fun inlineSuspendWithCurrentContinuation(call: JsInvocation, context: JsContext<JsNode>) {
-        val lambda = call.arguments[0]
-        val continuationArg = call.arguments[call.arguments.size - 1]
+                private fun replaceIfNecessary(expression: JsExpression, ctx: JsContext<JsNode>) {
+                    val alias = expression.localAlias
+                    if (alias != null) {
+                        ctx.replaceMe(alias.makeRef())
+                        inlinedModuleAliases.add(alias)
+                    }
+                }

-        val invocation = JsInvocation(lambda, continuationArg)
-        invocation.isSuspend = true
-        context.replaceMe(accept(invocation))
+            }.accept(statement)
+        }
+
+        private fun hasToBeInlined(call: JsInvocation): Boolean {
+            val strategy = call.inlineStrategy
+            return if (strategy == null || !strategy.isInline) false else functionContext.hasFunctionDefinition(call)
+        }
    }

    private fun reportInlineCycle(call: JsInvocation, calledFunction: JsFunction) {
@@ -517,50 +454,49 @@ class JsInliner private constructor(
        }
    }

-    private fun hasToBeInlined(call: JsInvocation): Boolean {
-        val strategy = call.inlineStrategy
-        return if (strategy == null || !strategy.isInline) false else functionContext.hasFunctionDefinition(call)

-    }
-
-    private inner class JsInliningContext internal constructor(override val statementContextBeforeCurrentFunction: JsContext<JsStatement>) :
-        InliningContext {
-        override val functionContext: FunctionContext
-
-        override val statementContext: JsContext<JsStatement>
-            get() = lastStatementLevelContext
-
-        init {
-            functionContext = object : FunctionContext(functionReader, config) {
-                override fun lookUpStaticFunction(functionName: JsName?): FunctionWithWrapper? {
-                    return functions[functionName]
-                }
-
-                override fun lookUpStaticFunctionByTag(functionTag: String): FunctionWithWrapper? {
-                    return accessors[functionTag]
-                }
+    private fun addInlinedModules(fragment: JsProgramFragment, inliner: InlinerImpl) {
+        val existingModules = fragment.importedModules.mapTo(IdentitySet()) { it.internalName }
+        inliner.inlinedModuleAliases.forEach {
+            if (it !in existingModules) {
+                fragment.importedModules.add(moduleMap[it]!!.let {
+                    // Copy so that the Merger.kt doesn't operate on the same instance in different fragments.
+                    JsImportedModule(it.externalName, it.internalName, it.plainReference)
+                })
            }
        }
+    }

-        override fun newNamingContext(): NamingContext {
-            return NamingContext(statementContext)
-        }
+
+    val functionContext: FunctionContext = object : FunctionContext(functionReader, config) {
+        override fun lookUpStaticFunction(functionName: JsName?): FunctionWithWrapper? = functions[functionName]
+
+        override fun lookUpStaticFunctionByTag(functionTag: String): FunctionWithWrapper? = accessors[functionTag]
    }

    private class JsCallInfo(val call: JsInvocation, val containingFunction: JsFunction)

-    internal class JsWrapperKey(val context: JsContext<JsStatement>, private val function: JsFunction) {
+    fun process(fragment: JsProgramFragment, existingImports: MutableMap<String, JsName>) {
+        val existingNameBindings = fragment.nameBindings.associateTo(IdentityHashMap()) { it.name to it.key }

-        override fun equals(o: Any?): Boolean {
-            if (this === o) return true
-            if (o == null || javaClass != o.javaClass) return false
-            val key = o as JsWrapperKey?
-            return context == key!!.context && function == key.function
+        val additionalDeclarations = mutableListOf<JsStatement>()
+        val inliner = InlinerImpl(existingNameBindings, existingImports, 0) {
+            additionalDeclarations.add(it)
        }

-        override fun hashCode(): Int {
-            return Objects.hash(context, function)
+        inliner.acceptStatement(fragment.declarationBlock)
+        // Mostly for the sake of post-processor
+        // TODO are inline function marked with @Test possible?
+        if (fragment.tests != null) {
+            inliner.acceptStatement(fragment.tests)
        }
+        inliner.acceptStatement(fragment.initializerBlock)
+
+        // TODO fix the order
+        fragment.declarationBlock.statements.addAll(0, additionalDeclarations)
+        fragment.nameBindings.addAll(inliner.additionalNameBindings)
+
+        addInlinedModules(fragment, inliner)
    }

    companion object {
@@ -569,55 +505,27 @@ class JsInliner private constructor(
            reporter: JsConfig.Reporter,
            config: JsConfig,
            trace: DiagnosticSink,
-            currentModuleName: JsName,
-            fragments: List<JsProgramFragment>,
-            fragmentsToProcess: List<JsProgramFragment>,
-            importStatements: List<JsStatement>
+            translationResult: AstGenerationResult
        ) {
-            val functions = collectNamedFunctionsAndWrappers(fragments)
-            val accessors = collectAccessors(fragments)
-            val inverseNameBindings = collectNameBindings(fragments)
+            val functions = collectNamedFunctionsAndWrappers(translationResult.newFragments)
+            val accessors = collectAccessors(translationResult.fragments)
+
+            val inverseNameBindings = inverseNameBindings(*translationResult.fragments.toTypedArray())

            val accessorInvocationTransformer = DummyAccessorInvocationTransformer()
-            for (fragment in fragmentsToProcess) {
+            for (fragment in translationResult.newFragments) {
                accessorInvocationTransformer.accept<JsGlobalBlock>(fragment.declarationBlock)
                accessorInvocationTransformer.accept<JsGlobalBlock>(fragment.initializerBlock)
            }
-            val functionReader = FunctionReader(reporter, config, currentModuleName, fragments)
-            val inliner = JsInliner(config, functions, accessors, inverseNameBindings, functionReader, trace)
+            val functionReader = FunctionReader(reporter, config, translationResult.innerModuleName, translationResult.fragments)
+            val moduleMap = translationResult.importedModuleList.associate { it.internalName to it }

-            for (statement in importStatements) {
-                inliner.processImportStatement(statement)
+            val inliner = JsInliner(config, functions, accessors, inverseNameBindings, functionReader, trace, moduleMap)
+            for (fragment in translationResult.newFragments) {
+                inliner.process(fragment, inverseNameBindings(fragment).entries.associateTo(mutableMapOf()) { (name, tag) -> tag to name })
            }

-            val moduleMap = fillModuleMap(buildModuleMap(fragments), fragmentsToProcess)
-
-            for (fragment in fragmentsToProcess) {
-                inliner.existingImports.clear()
-                inliner.additionalNameBindings.clear()
-                inliner.inlinedModuleAliases.clear()
-                inliner.existingNameBindings = collectNameBindings(listOf(fragment))
-
-                inliner.acceptStatement<JsGlobalBlock>(fragment.declarationBlock)
-                // Mostly for the sake of post-processor
-                // TODO are inline function marked with @Test possible?
-                if (fragment.tests != null) {
-                    inliner.acceptStatement(fragment.tests)
-                }
-
-                // There can be inlined function in top-level initializers, we need to optimize them as well
-                val fakeInitFunction = JsFunction(JsDynamicScope, fragment.initializerBlock, "")
-                val initWrapper = JsGlobalBlock()
-                initWrapper.statements.add(JsExpressionStatement(fakeInitFunction))
-                inliner.accept(initWrapper)
-                initWrapper.statements.removeAt(initWrapper.statements.size - 1)
-
-                fragment.initializerBlock.getStatements().addAll(0, initWrapper.statements)
-                fragment.nameBindings.addAll(inliner.additionalNameBindings)
-                inliner.addInlinedModules(fragment, moduleMap)
-            }
-
-            for (fragment in fragmentsToProcess) {
+            for (fragment in translationResult.newFragments) {
                val block = JsBlock(fragment.declarationBlock, fragment.initializerBlock, fragment.exportBlock)
                removeUnusedImports(block)
                simplifyWrappedFunctions(block)
@@ -625,29 +533,16 @@ class JsInliner private constructor(
            }
        }

-        private fun buildModuleMap(fragments: List<JsProgramFragment>): MutableMap<JsName, JsImportedModule> {
-            return fillModuleMap(HashMap(), fragments)
-        }
-
-        private fun fillModuleMap(
-            map: MutableMap<JsName, JsImportedModule>,
-            fragments: List<JsProgramFragment>
-        ): MutableMap<JsName, JsImportedModule> {
-            for (fragment in fragments) {
-                for (module in fragment.importedModules) {
-                    map[module.internalName] = module
+        private fun inverseNameBindings(vararg fragments: JsProgramFragment): MutableMap<JsName, String> {
+            val name2Tag = mutableMapOf<JsName, String>()
+            fragments.forEach {
+                it.nameBindings.forEach { (tag, name) ->
+                    name2Tag[name] = tag
                }
            }
-            return map
+
+            return name2Tag
        }

-        private fun isSuspendWithCurrentContinuation(
-            descriptor: DeclarationDescriptor?,
-            languageVersionSettings: LanguageVersionSettings
-        ): Boolean {
-            return (descriptor as? FunctionDescriptor)?.original?.isBuiltInSuspendCoroutineUninterceptedOrReturn(
-                languageVersionSettings
-            ) ?: false
-        }
    }
 }
@@ -21,13 +21,6 @@ import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*

 object LabeledBlockToDoWhileTransformation {
-    fun apply(fragments: List<JsProgramFragment>) {
-        for (fragment in fragments) {
-            apply(fragment.declarationBlock)
-            apply(fragment.initializerBlock)
-        }
-    }
-
    fun apply(root: JsNode) {
        object : JsVisitorWithContextImpl() {
            val loopOrSwitchStack = Stack<JsStatement>()
@@ -120,3 +113,10 @@ object LabeledBlockToDoWhileTransformation {
        }.accept(loopOrSwitch)
    }
 }
+
+fun transformLabeledBlockToDoWhile(fragments: List<JsProgramFragment>) {
+    for (fragment in fragments) {
+        LabeledBlockToDoWhileTransformation.apply(fragment.declarationBlock)
+        LabeledBlockToDoWhileTransformation.apply(fragment.initializerBlock)
+    }
+}
@@ -19,12 +19,6 @@ package org.jetbrains.kotlin.js.inline.context
 import org.jetbrains.kotlin.js.backend.ast.JsContext
 import org.jetbrains.kotlin.js.backend.ast.JsStatement

-interface InliningContext {
-    val statementContext: JsContext<JsStatement>
-
-    val statementContextBeforeCurrentFunction: JsContext<JsStatement>
-
-    val functionContext: FunctionContext
-
-    fun newNamingContext(): NamingContext
+class InliningContext(val statementContext: JsContext<JsStatement>) {
+    fun newNamingContext() = NamingContext(statementContext)
 }
@@ -230,16 +230,6 @@ fun collectAccessors(fragments: List<JsProgramFragment>): Map<String, FunctionWi
    return result
 }

-fun collectNameBindings(fragments: List<JsProgramFragment>): MutableMap<JsName, String> {
-    val result = mutableMapOf<JsName, String>()
-    for (fragment in fragments) {
-        for (binding in fragment.nameBindings) {
-            result[binding.name] = binding.key
-        }
-    }
-    return result
-}
-
 fun extractFunction(expression: JsExpression) = when (expression) {
    is JsFunction -> FunctionWithWrapper(expression, null)
    else -> InlineMetadata.decompose(expression)?.function ?: InlineMetadata.tryExtractFunction(expression)
@@ -37,7 +37,7 @@ class JsAstSerializer(private val pathResolver: (File) -> String) {
    private val importedNames = mutableSetOf<JsName>()

    fun serialize(fragment: JsProgramFragment, output: OutputStream) {
-        val namesBySignature = fragment.nameBindings.associate { it.key to it.name }
+        val namesBySignature = fragment.nameBindings.associateTo(mutableMapOf()) { it.key to it.name }
        importedNames.clear()
        importedNames += fragment.imports.map { namesBySignature[it.key]!! }
        serialize(fragment).writeTo(output)
@@ -19,7 +19,7 @@ package org.jetbrains.kotlin.js.facade
 import com.intellij.openapi.vfs.VfsUtilCore
 import org.jetbrains.kotlin.config.CommonConfigurationKeys
 import org.jetbrains.kotlin.config.*
-import org.jetbrains.kotlin.incremental.js.IncrementalResultsConsumer
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.js.analyze.TopDownAnalyzerFacadeForJS
 import org.jetbrains.kotlin.js.analyzer.JsAnalysisResult
 import org.jetbrains.kotlin.js.config.JSConfigurationKeys
@@ -43,6 +43,11 @@ import java.io.IOException
 import java.util.ArrayList

 import org.jetbrains.kotlin.diagnostics.DiagnosticUtils.hasError
+import org.jetbrains.kotlin.js.backend.ast.JsGlobalBlock
+import org.jetbrains.kotlin.js.backend.ast.JsProgramFragment
+import org.jetbrains.kotlin.js.coroutine.transformCoroutines
+import org.jetbrains.kotlin.js.translate.general.AstGenerationResult
+import org.jetbrains.kotlin.resolve.BindingTrace

 /**
 * An entry point of translator.
@@ -74,7 +79,6 @@ class K2JSTranslator(private val config: JsConfig) {
        mainCallParameters: MainCallParameters,
        analysisResult: JsAnalysisResult? = null
    ): TranslationResult {
-        var analysisResult = analysisResult
        val files = ArrayList<KtFile>()
        for (unit in units) {
            if (unit is TranslationUnit.SourceFile) {
@@ -82,46 +86,93 @@ class K2JSTranslator(private val config: JsConfig) {
            }
        }

-        if (analysisResult == null) {
-            analysisResult = TopDownAnalyzerFacadeForJS.analyzeFiles(files, config)
-            ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
-        }
+        val actualAnalysisResult = analysisResult ?: TopDownAnalyzerFacadeForJS.analyzeFiles(files, config)

+        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+
+        return translate(reporter, files, units, mainCallParameters, actualAnalysisResult)
+    }
+
+    @Throws(TranslationException::class)
+    private fun translate(
+        reporter: JsConfig.Reporter,
+        files: List<KtFile>,
+        allUnits: List<TranslationUnit>,
+        mainCallParameters: MainCallParameters,
+        analysisResult: JsAnalysisResult
+    ): TranslationResult {
        val bindingTrace = analysisResult.bindingTrace
        TopDownAnalyzerFacadeForJS.checkForErrors(files, bindingTrace.bindingContext)
        val moduleDescriptor = analysisResult.moduleDescriptor
        val diagnostics = bindingTrace.bindingContext.diagnostics
-
        val pathResolver = SourceFilePathResolver.create(config)

-        val translationResult = Translation.generateAst(
-            bindingTrace, units, mainCallParameters, moduleDescriptor, config, pathResolver
-        )
-        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+        val translationResult = Translation.generateAst(bindingTrace, allUnits, mainCallParameters, moduleDescriptor, config, pathResolver)
        if (hasError(diagnostics)) return TranslationResult.Fail(diagnostics)
-
-        val newFragments = ArrayList(translationResult.newFragments)
-        val allFragments = ArrayList(translationResult.fragments)
+        checkCanceled()

        JsInliner.process(
-            reporter, config, analysisResult.bindingTrace, translationResult.innerModuleName,
-            allFragments, newFragments, translationResult.importStatements
+            reporter,
+            config,
+            analysisResult.bindingTrace,
+            translationResult
        )
-
-        LabeledBlockToDoWhileTransformation.apply(newFragments)
-
-        val coroutineTransformer = CoroutineTransformer()
-        for (fragment in newFragments) {
-            coroutineTransformer.accept(fragment.declarationBlock)
-            coroutineTransformer.accept(fragment.initializerBlock)
-        }
-        removeUnusedImports(translationResult.program)
-
-        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
        if (hasError(diagnostics)) return TranslationResult.Fail(diagnostics)
+        checkCanceled()

-        expandIsCalls(newFragments)
+        transformLabeledBlockToDoWhile(translationResult.newFragments)
+        checkCanceled()
+
+        transformCoroutines(translationResult.newFragments)
+        checkCanceled()
+
+        expandIsCalls(translationResult.newFragments)
+        checkCanceled()
+
+        trySaveIncrementalData(files, translationResult, pathResolver, bindingTrace, moduleDescriptor)
+        checkCanceled()
+
+        // Global phases
+
+        val program = translationResult.buildProgram()
+
+        removeUnusedImports(program)
+        checkCanceled()
+
+        removeDuplicateImports(program)
+        checkCanceled()
+
+        program.resolveTemporaryNames()
+        checkCanceled()
+
+        return if (hasError(diagnostics)) {
+            TranslationResult.Fail(diagnostics)
+        } else {
+            TranslationResult.Success(
+                config,
+                files,
+                program,
+                diagnostics,
+                translationResult.importedModuleList.map { it.externalName },
+                moduleDescriptor,
+                bindingTrace.bindingContext
+            )
+        }
+    }
+
+    private fun checkCanceled() {
        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+    }
+
+    private fun trySaveIncrementalData(
+        files: List<KtFile>,
+        translationResult: AstGenerationResult,
+        pathResolver: SourceFilePathResolver,
+        bindingTrace: BindingTrace,
+        moduleDescriptor: ModuleDescriptor
+    ) {
+        if (incrementalResults == null) return
+
        val serializer = JsAstSerializer { file ->
            try {
                pathResolver.getPathRelativeToSourceRoots(file)
@@ -130,45 +181,25 @@ class K2JSTranslator(private val config: JsConfig) {
            }
        }

-        if (incrementalResults != null) {
-            val serializationUtil = KotlinJavascriptSerializationUtil
+        for (file in files) {
+            val fragment = translationResult.fragmentMap[file] ?: error("Could not find AST for file: $file")
+            val output = ByteArrayOutputStream()
+            serializer.serialize(fragment, output)
+            val binaryAst = output.toByteArray()

-            for (file in files) {
-                val fragment = translationResult.fragmentMap[file] ?: error("Could not find AST for file: $file")
-                val output = ByteArrayOutputStream()
-                serializer.serialize(fragment, output)
-                val binaryAst = output.toByteArray()
+            val scope = translationResult.fileMemberScopes[file] ?: error("Could not find descriptors for file: $file")
+            val metadataVersion = config.configuration.get(CommonConfigurationKeys.METADATA_VERSION)
+            val packagePart = KotlinJavascriptSerializationUtil.serializeDescriptors(
+                bindingTrace.bindingContext, moduleDescriptor, scope, file.packageFqName,
+                config.configuration.languageVersionSettings,
+                metadataVersion ?: JsMetadataVersion.INSTANCE
+            )

-                val scope = translationResult.fileMemberScopes[file] ?: error("Could not find descriptors for file: $file")
-                val metadataVersion = config.configuration.get(CommonConfigurationKeys.METADATA_VERSION)
-                val packagePart = serializationUtil.serializeDescriptors(
-                    bindingTrace.bindingContext, moduleDescriptor, scope, file.packageFqName,
-                    config.configuration.languageVersionSettings,
-                    metadataVersion ?: JsMetadataVersion.INSTANCE
-                )
-
-                val ioFile = VfsUtilCore.virtualToIoFile(file.virtualFile)
-                incrementalResults.processPackagePart(ioFile, packagePart.toByteArray(), binaryAst)
-            }
-
-            val settings = config.configuration.languageVersionSettings
-            incrementalResults.processHeader(serializationUtil.serializeHeader(moduleDescriptor, null, settings).toByteArray())
+            val ioFile = VfsUtilCore.virtualToIoFile(file.virtualFile)
+            incrementalResults.processPackagePart(ioFile, packagePart.toByteArray(), binaryAst)
        }

-        removeDuplicateImports(translationResult.program)
-        translationResult.program.resolveTemporaryNames()
-
-        ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
-        if (hasError(diagnostics)) return TranslationResult.Fail(diagnostics)
-
-        val importedModules = ArrayList<String>()
-        for (module in translationResult.importedModuleList) {
-            importedModules.add(module.externalName)
-        }
-
-        return TranslationResult.Success(
-            config, files, translationResult.program, diagnostics, importedModules,
-            moduleDescriptor, bindingTrace.bindingContext
-        )
+        val settings = config.configuration.languageVersionSettings
+        incrementalResults.processHeader(KotlinJavascriptSerializationUtil.serializeHeader(moduleDescriptor, null, settings).toByteArray())
    }
 }
@@ -17,16 +17,30 @@
 package org.jetbrains.kotlin.js.translate.general

 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
+import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.serialization.js.ModuleKind

 class AstGenerationResult(
-        val program: JsProgram,
-        val innerModuleName: JsName,
-        val fragments: List<JsProgramFragment>,
-        val fragmentMap: Map<KtFile, JsProgramFragment>,
-        val newFragments: List<JsProgramFragment>,
-        val importStatements: List<JsStatement>,
-        val fileMemberScopes: Map<KtFile, List<DeclarationDescriptor>>,
-        val importedModuleList: List<JsImportedModule>
-)
+    val fragments: List<JsProgramFragment>,
+    val fragmentMap: Map<KtFile, JsProgramFragment>,
+    val newFragments: List<JsProgramFragment>,
+    val fileMemberScopes: Map<KtFile, List<DeclarationDescriptor>>,
+    private val program: JsProgram,
+    private val moduleDescriptor: ModuleDescriptor,
+    private val moduleId: String,
+    private val moduleKind: ModuleKind
+) {
+
+    val innerModuleName = program.getScope().declareName("_");
+
+    val importedModuleList: List<JsImportedModule>
+        get() = fragments.flatMap { it.importedModules }
+
+    fun buildProgram(): JsProgram {
+        val merger = Merger(program, innerModuleName, moduleDescriptor, moduleId, moduleKind)
+        fragments.forEach { merger.addFragment(it) }
+        return merger.buildProgram()
+    }
+}
@@ -26,12 +26,24 @@ import org.jetbrains.kotlin.js.translate.context.Namer
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import org.jetbrains.kotlin.js.translate.utils.createPrototypeStatements
 import org.jetbrains.kotlin.js.translate.utils.definePackageAlias
+import org.jetbrains.kotlin.serialization.js.ModuleKind
+
+class Merger(
+    val program: JsProgram,
+    val internalModuleName: JsName,
+    val module: ModuleDescriptor,
+    val moduleId: String,
+    val moduleKind: ModuleKind
+) {
+    private val rootFunction = JsFunction(program.rootScope, JsBlock(), "root function").also {
+        it.parameters.add(JsParameter(internalModuleName))
+    }

-class Merger(private val rootFunction: JsFunction, val internalModuleName: JsName, val module: ModuleDescriptor) {
    // Maps unique signature (see generateSignature) to names
    private val nameTable = mutableMapOf<String, JsName>()
+
    private val importedModuleTable = mutableMapOf<JsImportedModuleKey, JsName>()
-    val importBlock = JsGlobalBlock()
+    private val importBlock = JsGlobalBlock()
    private val declarationBlock = JsGlobalBlock()
    private val initializerBlock = JsGlobalBlock()
    private val testsMap = mutableMapOf<String, JsStatement>()
@@ -43,25 +55,11 @@ class Merger(private val rootFunction: JsFunction, val internalModuleName: JsNam
    private val exportedPackages = mutableMapOf<String, JsName>()
    private val exportedTags = mutableSetOf<String>()

+    private val fragments = mutableListOf<JsProgramFragment>()
+
    // Add declaration and initialization statements from program fragment to resulting single program
    fun addFragment(fragment: JsProgramFragment) {
-        val nameMap = buildNameMap(fragment)
-        nameMap.rename(fragment)
-
-        for ((key, importExpr) in fragment.imports) {
-            if (declaredImports.add(key)) {
-                val name = nameTable[key]!!
-                importBlock.statements += JsAstUtils.newVar(name, importExpr)
-            }
-        }
-
-        declarationBlock.statements += fragment.declarationBlock
-        initializerBlock.statements += fragment.initializerBlock
-        fragment.tryUpdateTests()
-        fragment.tryUpdateMain()
-        addExportStatements(fragment)
-
-        classes += fragment.classes
+        fragments.add(fragment)
    }

    private fun JsProgramFragment.tryUpdateTests() {
@@ -205,8 +203,32 @@ class Merger(private val rootFunction: JsFunction, val internalModuleName: JsNam
        return rootNode
    }

+    private fun mergeNames() {
+        for (fragment in fragments) {
+            val nameMap = buildNameMap(fragment)
+            nameMap.rename(fragment)
+
+            for ((key, importExpr) in fragment.imports) {
+                if (declaredImports.add(key)) {
+                    val name = nameTable[key]!!
+                    importBlock.statements += JsAstUtils.newVar(name, importExpr)
+                }
+            }
+
+            declarationBlock.statements += fragment.declarationBlock
+            initializerBlock.statements += fragment.initializerBlock
+            fragment.tryUpdateTests()
+            fragment.tryUpdateMain()
+            addExportStatements(fragment)
+
+            classes += fragment.classes
+        }
+    }
+
    // Adds different boilerplate code (like imports, class prototypes, etc) to resulting program.
    fun merge() {
+        mergeNames()
+
        rootFunction.body.statements.apply {
            addImportForInlineDeclarationIfNecessary()
            this += importBlock.statements
@@ -220,6 +242,37 @@ class Merger(private val rootFunction: JsFunction, val internalModuleName: JsNam
        }
    }

+    fun buildProgram(): JsProgram {
+        merge()
+
+        val rootBlock = rootFunction.getBody()
+
+        val statements = rootBlock.getStatements()
+
+        statements.add(0, JsStringLiteral("use strict").makeStmt())
+        if (!isBuiltinModule(fragments)) {
+            defineModule(program, statements, moduleId)
+        }
+
+        // Invoke function passing modules as arguments
+        // This should help minifier tool to recognize references to these modules as local variables and make them shorter.
+        for (importedModule in importedModules) {
+            rootFunction.parameters.add(JsParameter(importedModule.internalName))
+        }
+
+        statements.add(JsReturn(internalModuleName.makeRef()))
+
+        val block = program.globalBlock
+        block.statements.addAll(
+            ModuleWrapperTranslation.wrapIfNecessary(
+                moduleId, rootFunction, importedModules, program,
+                moduleKind
+            )
+        )
+
+        return program
+    }
+
    private fun MutableList<JsStatement>.addImportForInlineDeclarationIfNecessary() {
        val importsForInlineName = nameTable[Namer.IMPORTS_FOR_INLINE_PROPERTY] ?: return
        this += definePackageAlias(
@@ -267,4 +320,34 @@ class Merger(private val rootFunction: JsFunction, val internalModuleName: JsNam
        cls.interfaces.forEach { addClassPostDeclarations(it, visited, statements) }
        statements += cls.postDeclarationBlock.statements
    }
+
+    companion object {
+        private val ENUM_SIGNATURE = "kotlin\$Enum"
+
+        // TODO is there no better way?
+        private fun isBuiltinModule(fragments: List<JsProgramFragment>): Boolean {
+            for (fragment in fragments) {
+                for (nameBinding in fragment.nameBindings) {
+                    if (nameBinding.key == ENUM_SIGNATURE && !fragment.imports.containsKey(ENUM_SIGNATURE)) {
+                        return true
+                    }
+                }
+            }
+            return false
+        }
+
+        private fun defineModule(program: JsProgram, statements: MutableList<JsStatement>, moduleId: String) {
+            val rootPackageName = program.scope.findName(Namer.getRootPackageName())
+            if (rootPackageName != null) {
+                val namer = Namer.newInstance(program.scope)
+                statements.add(
+                    JsInvocation(
+                        namer.kotlin("defineModule"),
+                        JsStringLiteral(moduleId),
+                        rootPackageName.makeRef()
+                    ).makeStmt()
+                )
+            }
+        }
+    }
 }
@@ -324,10 +324,8 @@ public final class Translation {
            @NotNull JsConfig config,
            @NotNull SourceFilePathResolver sourceFilePathResolver
    ) {
+
        JsProgram program = new JsProgram();
-        JsFunction rootFunction = new JsFunction(program.getRootScope(), new JsBlock(), "root function");
-        JsName internalModuleName = program.getScope().declareName("_");
-        Merger merger = new Merger(rootFunction, internalModuleName, moduleDescriptor);

        Map<KtFile, JsProgramFragment> fragmentMap = new HashMap<>();
        List<JsProgramFragment> fragments = new ArrayList<>();
@@ -354,56 +352,22 @@ public final class Translation {
                newFragments.add(fragment);
                fragmentMap.put(file, fragment);
                fileMemberScopes.put(file, fileMemberScope);
-                merger.addFragment(fragment);
            }
            else if (unit instanceof TranslationUnit.BinaryAst) {
                byte[] astData = ((TranslationUnit.BinaryAst) unit).getData();
                JsProgramFragment fragment = deserializer.deserialize(new ByteArrayInputStream(astData));
-                merger.addFragment(fragment);
                fragments.add(fragment);
            }
        }

-        rootFunction.getParameters().add(new JsParameter(internalModuleName));
-
-        merger.merge();
-
-        JsBlock rootBlock = rootFunction.getBody();
-
-        List<JsStatement> statements = rootBlock.getStatements();
-
-        statements.add(0, new JsStringLiteral("use strict").makeStmt());
-        if (!isBuiltinModule(fragments)) {
-            defineModule(program, statements, config.getModuleId());
-        }
-
-        // Invoke function passing modules as arguments
-        // This should help minifier tool to recognize references to these modules as local variables and make them shorter.
-        List<JsImportedModule> importedModuleList = merger.getImportedModules();
-
-        for (JsImportedModule importedModule : importedModuleList) {
-            rootFunction.getParameters().add(new JsParameter(importedModule.getInternalName()));
-        }
-
-        statements.add(new JsReturn(internalModuleName.makeRef()));
-
-        JsBlock block = program.getGlobalBlock();
-        block.getStatements().addAll(wrapIfNecessary(config.getModuleId(), rootFunction, importedModuleList, program,
-                                                     config.getModuleKind()));
-
-        return new AstGenerationResult(program, internalModuleName, fragments, fragmentMap, newFragments,
-                                       merger.getImportBlock().getStatements(), fileMemberScopes, importedModuleList);
-    }
-
-    private static boolean isBuiltinModule(@NotNull List<JsProgramFragment> fragments) {
-        for (JsProgramFragment fragment : fragments) {
-            for (JsNameBinding nameBinding : fragment.getNameBindings()) {
-                if (nameBinding.getKey().equals(ENUM_SIGNATURE) && !fragment.getImports().containsKey(ENUM_SIGNATURE)) {
-                    return true;
-                }
-            }
-        }
-        return false;
+        return new AstGenerationResult(fragments,
+                                       fragmentMap,
+                                       newFragments,
+                                       fileMemberScopes,
+                                       program,
+                                       moduleDescriptor,
+                                       config.getModuleId(),
+                                       config.getModuleKind());
    }

    private static void translateFile(
@@ -430,15 +394,6 @@ public final class Translation {
        }
    }

-    private static void defineModule(@NotNull JsProgram program, @NotNull List<JsStatement> statements, @NotNull String moduleId) {
-        JsName rootPackageName = program.getScope().findName(Namer.getRootPackageName());
-        if (rootPackageName != null) {
-            Namer namer = Namer.newInstance(program.getScope());
-            statements.add(new JsInvocation(namer.kotlin("defineModule"), new JsStringLiteral(moduleId),
-                                            rootPackageName.makeRef()).makeStmt());
-        }
-    }
-
    @Nullable
    private static JsStatement mayBeGenerateTests(
            @NotNull TranslationContext context,
@@ -1,6 +1,5 @@
 // KJS_WITH_FULL_RUNTIME
 // EXPECTED_REACHABLE_NODES: 1625
-// MODULE: lib
 // FILE: lib.kt


@@ -8,8 +7,8 @@ inline fun tenUInt() = 10U

 inline fun tenULong() = 10UL

-// MODULE: main(lib)
 // FILE: main.kt
+// RECOMPILE

 fun box(): String {