JS: split JsInliner

2018-11-19 23:14:21 +03:00
parent c65383fa3f
commit c83b6d46cc
9 changed files with 673 additions and 474 deletions
@@ -0,0 +1,139 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.js.inline
+
+import org.jetbrains.kotlin.js.backend.ast.*
+import org.jetbrains.kotlin.js.backend.ast.metadata.localAlias
+import org.jetbrains.kotlin.js.backend.ast.metadata.staticRef
+import org.jetbrains.kotlin.js.inline.util.collectDefinedNamesInAllScopes
+import org.jetbrains.kotlin.js.inline.util.collectNamedFunctions
+import org.jetbrains.kotlin.js.inline.util.getImportTag
+import org.jetbrains.kotlin.js.inline.util.replaceNames
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import java.util.ArrayList
+import java.util.HashMap
+
+fun applyWrapper(
+    wrapper: JsBlock, function: JsFunction, originalFunction: JsFunction,
+    inlineFunctionDepth: Int,
+    replacementsInducedByWrappers: MutableMap<JsFunction, Map<JsName, JsNameRef>>,
+    existingImports: MutableMap<String, JsName>,
+    additionalImports: MutableList<Triple<String, JsExpression, JsName>>,
+    existingNameBindings: MutableMap<JsName, String>,
+    additionalNameBindings: MutableList<JsNameBinding>,
+    inlinedModuleAliases: MutableSet<JsName>,
+    inverseNameBindings: Map<JsName, String>,
+    addPrevious: (JsStatement) -> Unit
+) {
+
+    // 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>()
+        val importStatements = ArrayList<Pair<String, JsVars.JsVar>>()
+        wrapper.statements.asSequence()
+            .filterNot { it is JsReturn }
+            .map { it.deepCopy() }
+            .forEach { statement ->
+                if (inlineFunctionDepth == 0) {
+                    replaceExpressionsWithLocalAliases(statement, inlinedModuleAliases)
+                }
+
+                if (statement is JsVars) {
+                    val tag = getImportTag(statement)
+                    if (tag != null) {
+                        // TODO handle JsVars with multiple vars?
+                        val name = statement.vars[0].name
+                        var existingName: JsName? = if (inlineFunctionDepth == 0) name.localAlias else null
+                        if (existingName == null) {
+                            existingName = existingImports.computeIfAbsent(tag) {
+                                importStatements.add(tag to statement.vars[0])
+                                val alias = JsScope.declareTemporaryName(name.ident)
+                                alias.copyMetadataFrom(name)
+                                newReplacements[name] = JsAstUtils.pureFqn(alias, null)
+                                alias
+                            }
+                        }
+
+                        if (name !== existingName) {
+                            val replacement = JsAstUtils.pureFqn(existingName, null)
+                            newReplacements[name] = replacement
+                        }
+
+                        return@forEach
+                    }
+                }
+
+                copiedStatements.add(statement)
+            }
+
+        val definedNames = (importStatements.asSequence().map { it.second } + 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 = JsAstUtils.pureFqn(alias, null)
+            newReplacements[name] = replacement
+        }
+
+        for ((tag, statement) in importStatements) {
+            val renamed = replaceNames(statement, newReplacements)
+            additionalImports.add(Triple(tag,renamed.initExpression, renamed.name))
+            additionalNameBindings.add(JsNameBinding(tag, renamed.name))
+        }
+
+        for (statement in copiedStatements) {
+            addPrevious(replaceNames(statement, newReplacements))
+        }
+
+        for ((key, value) in collectNamedFunctions(JsBlock(copiedStatements))) {
+            if (key.staticRef is JsFunction) {
+                key.staticRef = value
+            }
+        }
+
+        newReplacements
+    }
+
+    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))
+            }
+        }
+    }
+}
+
+private fun replaceExpressionsWithLocalAliases(statement: JsStatement, inlinedModuleAliases: MutableSet<JsName>) {
+    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)
+            }
+        }
+
+    }.accept(statement)
+}
@@ -59,7 +59,7 @@ class FunctionReader(
        private val reporter: JsConfig.Reporter,
        private val config: JsConfig,
        private val currentModuleName: JsName,
-        fragments: List<JsProgramFragment>
+        private val moduleNameMap: Map<String, JsExpression>
 ) {
    /**
     * fileContent: .js file content, that contains this module definition.
@@ -144,21 +144,9 @@ class FunctionReader(
        result
    }

-    private val moduleNameMap: Map<String, JsExpression>
    private val shouldRemapPathToRelativeForm = config.shouldGenerateRelativePathsInSourceMap()
    private val relativePathCalculator = config.configuration[JSConfigurationKeys.OUTPUT_DIR]?.let { RelativePathCalculator(it) }

-    init {
-        moduleNameMap = buildModuleNameMap(fragments)
-    }
-
-    // Since we compile each source file in its own context (and we may loose these context when performing incremental compilation)
-    // we don't use contexts to generate proper names for modules. Instead, we generate all necessary information during
-    // translation and rely on it here.
-    private fun buildModuleNameMap(fragments: List<JsProgramFragment>): Map<String, JsExpression> {
-        return fragments.flatMap { it.inlineModuleMap.entries }.associate { (k, v) -> k to v }
-    }
-
    private fun rewindToIdentifierStart(text: String, index: Int): Int {
        var result = index
        while (result > 0 && Character.isJavaIdentifierPart(text[result - 1])) {
@@ -0,0 +1,130 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.js.inline
+
+import org.jetbrains.kotlin.diagnostics.DiagnosticSink
+import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.js.backend.ast.JsBlock
+import org.jetbrains.kotlin.js.backend.ast.JsFunction
+import org.jetbrains.kotlin.js.backend.ast.JsInvocation
+import org.jetbrains.kotlin.js.backend.ast.JsName
+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.clean.FunctionPostProcessor
+import org.jetbrains.kotlin.js.inline.clean.removeUnusedLocalFunctionDeclarations
+import org.jetbrains.kotlin.js.inline.context.FunctionContext
+import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper
+import org.jetbrains.kotlin.js.inline.util.IdentitySet
+import org.jetbrains.kotlin.js.inline.util.refreshLabelNames
+import org.jetbrains.kotlin.resolve.inline.InlineStrategy
+import java.util.*
+
+class InlineDfsController(
+    val trace: DiagnosticSink,
+    val functions: Map<JsName, FunctionWithWrapper>,
+    val accessors: Map<String, FunctionWithWrapper>,
+    val functionContext: FunctionContext
+) {
+
+    val functionsByWrapperNodes =
+        HashMap<JsBlock, FunctionWithWrapper>()
+    val functionsByFunctionNodes =
+        HashMap<JsFunction, FunctionWithWrapper>()
+
+    init {
+        (functions.values.asSequence() + accessors.values.asSequence()).forEach { f ->
+            functionsByFunctionNodes[f.function] = f
+            if (f.wrapperBody != null) {
+                functionsByWrapperNodes[f.wrapperBody] = f
+            }
+        }
+    }
+
+    private val processedFunctions = IdentitySet<JsFunction>()
+    private val inProcessFunctions = IdentitySet<JsFunction>()
+    // these are needed for error reporting, when inliner detects cycle
+    private val namedFunctionsStack = Stack<JsFunction>()
+
+    private val inlineCallInfos = LinkedList<JsCallInfo>()
+
+    val currentNamedFunction: JsFunction?
+        get() = if (namedFunctionsStack.empty()) null else namedFunctionsStack.peek()
+
+
+    fun startFunction(function: JsFunction) {
+        assert(!inProcessFunctions.contains(function)) { "Inliner has revisited function" }
+        inProcessFunctions.add(function)
+
+        if (function in functionsByFunctionNodes.keys) {
+            namedFunctionsStack.push(function)
+        }
+    }
+
+    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()
+        }
+    }
+
+
+    // Return true iff the definition should be visited by the inliner
+    fun visitCall(call: JsInvocation): Boolean {
+        val definition = functionContext.getFunctionDefinition(call)
+
+        currentNamedFunction?.let {
+            inlineCallInfos.add(JsCallInfo(call, it))
+        }
+
+
+        if (inProcessFunctions.contains(definition.function)) {
+            reportInlineCycle(call, definition.function)
+        } else if (!processedFunctions.contains(definition.function)) {
+            return true
+        }
+
+        return false
+    }
+
+    fun endVisit(x: JsInvocation) {
+        if (!inlineCallInfos.isEmpty()) {
+            if (inlineCallInfos.last.call == x) {
+                inlineCallInfos.removeLast()
+            }
+        }
+    }
+
+    private fun reportInlineCycle(call: JsInvocation, calledFunction: JsFunction) {
+        call.inlineStrategy = InlineStrategy.NOT_INLINE
+        val it = inlineCallInfos.descendingIterator()
+
+        while (it.hasNext()) {
+            val callInfo = it.next()
+            val psiElement = callInfo.call.psiElement
+
+            val descriptor = callInfo.call.descriptor
+            if (psiElement != null && descriptor != null) {
+                trace.report(Errors.INLINE_CALL_CYCLE.on(psiElement, descriptor))
+            }
+
+            if (callInfo.containingFunction == calledFunction) {
+                break
+            }
+        }
+    }
+}
+
+private class JsCallInfo(val call: JsInvocation, val containingFunction: JsFunction)
@@ -0,0 +1,98 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.js.inline
+
+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.translate.expression.InlineMetadata
+import java.util.ArrayList
+import java.util.HashMap
+import java.util.HashSet
+
+class InlineSuspendFunctionSplitter(
+    val existingNameBindings: MutableMap<JsName, String>,
+    val existingImports: MutableMap<String, JsName>,
+    val inlineFunctionDepth: Int,
+    val inverseNameBindings: Map<JsName, String>,
+    val replacementsInducedByWrappers: MutableMap<JsFunction, Map<JsName, JsNameRef>>,
+    val addPrevious: (JsStatement) -> Unit
+) : JsVisitorWithContextImpl() {
+
+    val inlinedModuleAliases = HashSet<JsName>()
+
+    val additionalNameBindings = ArrayList<JsNameBinding>()
+
+    val additionalImports = mutableListOf<Triple<String, JsExpression, JsName>>()
+
+
+    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
+        }
+    }
+
+    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
+
+                // 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,
+                        inlineFunctionDepth,
+                        replacementsInducedByWrappers,
+                        existingImports,
+                        additionalImports,
+                        existingNameBindings,
+                        additionalNameBindings,
+                        inlinedModuleAliases,
+                        inverseNameBindings,
+                        addPrevious
+                    )
+                }
+
+                // 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
+    }
+}
@@ -0,0 +1,233 @@
+/*
+ * Copyright 2010-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+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.backend.ast.metadata.synthetic
+import org.jetbrains.kotlin.js.inline.context.FunctionContext
+import org.jetbrains.kotlin.js.inline.context.InliningContext
+import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper
+import org.jetbrains.kotlin.js.inline.util.getImportTag
+import org.jetbrains.kotlin.js.translate.declaration.transformSpecialFunctionsToCoroutineMetadata
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import java.util.ArrayList
+import java.util.HashMap
+import java.util.HashSet
+
+class InlinerImpl(
+    val existingNameBindings: MutableMap<JsName, String>,
+    val existingImports: MutableMap<String, JsName>,
+    val inlineFunctionDepth: Int,
+    val dfsController: InlineDfsController,
+    val inverseNameBindings: Map<JsName, String>,
+    val functionContext: FunctionContext,
+    val addPrevious: (JsStatement) -> Unit
+) : JsVisitorWithContextImpl() {
+    val inlinedModuleAliases = HashSet<JsName>()
+
+    val replacementsInducedByWrappers: MutableMap<JsFunction, Map<JsName, JsNameRef>> = mutableMapOf()
+
+    val additionalNameBindings = ArrayList<JsNameBinding>()
+
+    val additionalImports = mutableListOf<Triple<String, JsExpression, JsName>>()
+
+    override fun visit(function: JsFunction, context: JsContext<*>): Boolean {
+        val functionWithWrapper = dfsController.functionsByFunctionNodes[function]
+        if (functionWithWrapper != null) {
+            visit(functionWithWrapper)
+            return false
+        } else {
+            dfsController.startFunction(function)
+            return super.visit(function, context)
+        }
+    }
+
+    override fun endVisit(function: JsFunction, context: JsContext<*>) {
+        super.endVisit(function, context)
+        if (!dfsController.functionsByFunctionNodes.containsKey(function)) {
+            dfsController.endFunction(function)
+        }
+    }
+
+    override fun visit(x: JsBlock, ctx: JsContext<*>): Boolean {
+        val functionWithWrapper = dfsController.functionsByWrapperNodes[x]
+        if (functionWithWrapper != null) {
+            visit(functionWithWrapper)
+            return false
+        }
+        return super.visit(x, ctx)
+    }
+
+    private fun visit(functionWithWrapper: FunctionWithWrapper) {
+        dfsController.startFunction(functionWithWrapper.function)
+
+        val wrapperBody = functionWithWrapper.wrapperBody
+        if (wrapperBody != null) {
+            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,
+                dfsController,
+                inverseNameBindings,
+                functionContext
+            ) {
+                additionalStatements.add(it)
+            }
+            for (statement in wrapperBody.statements) {
+                if (statement !is JsReturn) {
+                    innerInliner.acceptStatement(statement)
+                } else {
+                    innerInliner.accept((statement.expression as JsFunction).body)
+                }
+            }
+
+            val importStatements = innerInliner.additionalImports.map { (_, importExpr, name) ->
+                JsAstUtils.newVar(name, importExpr)
+            }
+
+            // TODO keep order
+            wrapperBody.statements.addAll(0, importStatements + additionalStatements)
+        } else {
+            accept(functionWithWrapper.function.body)
+        }
+
+        dfsController.endFunction(functionWithWrapper.function)
+    }
+
+    override fun visit(call: JsInvocation, context: JsContext<*>): Boolean {
+        if (!hasToBeInlined(call)) return true
+
+        if (dfsController.visitCall(call)) {
+            val definition = functionContext.getFunctionDefinition(call)
+
+            for (i in 0 until call.arguments.size) {
+                val argument = call.arguments[i]
+                call.arguments[i] = accept(argument)
+            }
+            visit(definition)
+
+            return false
+        }
+
+        return true
+    }
+
+    override fun endVisit(x: JsInvocation, ctx: JsContext<JsNode>) {
+        if (hasToBeInlined(x)) {
+            inline(x, ctx)
+        }
+
+        dfsController.endVisit(x)
+    }
+
+    override fun doAcceptStatementList(statements: MutableList<JsStatement>) {
+        var i = 0
+
+        while (i < statements.size) {
+            val additionalStatements =
+                ExpressionDecomposer.preserveEvaluationOrder(statements[i]) { node ->
+                    node is JsInvocation && hasToBeInlined(
+                        node
+                    )
+                }
+            statements.addAll(i, additionalStatements)
+            i += additionalStatements.size + 1
+        }
+
+        super.doAcceptStatementList(statements)
+    }
+
+    private fun inline(call: JsInvocation, context: JsContext<JsNode>) {
+        var functionWithWrapper = 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
+        dfsController.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,
+                inlineFunctionDepth,
+                replacementsInducedByWrappers,
+                existingImports,
+                additionalImports,
+                existingNameBindings,
+                additionalNameBindings,
+                inlinedModuleAliases,
+                inverseNameBindings
+            ) {
+                addPrevious(accept(it))
+            }
+        }
+
+        val inliningContext = InliningContext(lastStatementLevelContext)
+
+        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.
+        // TODO This seems brittle
+        val currentFunction = dfsController.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(JsAstUtils.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 hasToBeInlined(call: JsInvocation): Boolean {
+        val strategy = call.inlineStrategy
+        return if (strategy == null || !strategy.isInline) false else functionContext.hasFunctionDefinition(call)
+    }
+}
@@ -6,458 +6,32 @@
 package org.jetbrains.kotlin.js.inline

 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
-import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.js.backend.ast.*
-import org.jetbrains.kotlin.js.backend.ast.metadata.*
 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.util.*
-import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
-import org.jetbrains.kotlin.resolve.inline.InlineStrategy

 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

 class JsInliner private constructor(
-    private val config: JsConfig,
-    private val functions: Map<JsName, FunctionWithWrapper>,
-    private val accessors: Map<String, FunctionWithWrapper>,
+    config: JsConfig,
+    functionReader: FunctionReader,
+    val functions: Map<JsName, FunctionWithWrapper>,
+    val accessors: Map<String, FunctionWithWrapper>,
    private val inverseNameBindings: Map<JsName, String>,
-    private val functionReader: FunctionReader,
-    private val trace: DiagnosticSink,
-    private val moduleMap: Map<JsName, JsImportedModule>
+    private val moduleMap: Map<JsName, JsImportedModule>,
+    private val trace: DiagnosticSink
 ) {
-    private val namedFunctionsSet: MutableSet<JsFunction> = functions.values.mapTo(IdentitySet()) { it.function }

-    private val processedFunctions = IdentitySet<JsFunction>()
-    private val inProcessFunctions = IdentitySet<JsFunction>()
+    private val functionContext = FunctionContext(functionReader, config, functions, accessors)

-    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
-            if (f.wrapperBody != null) {
-                functionsByWrapperNodes[f.wrapperBody] = f
-            }
-        }
-    }
-
-    // these are needed for error reporting, when inliner detects cycle
-    private val namedFunctionsStack = Stack<JsFunction>()
-    private val inlineCallInfos = LinkedList<JsCallInfo>()
-
-    private val currentNamedFunction: JsFunction?
-        get() = if (namedFunctionsStack.empty()) null else namedFunctionsStack.peek()
-
-    private inner class InlinerImpl(
-        val existingNameBindings: MutableMap<JsName, String>,
-        val existingImports: MutableMap<String, JsName>,
-        val inlineFunctionDepth: Int,
-        val addPrevious: (JsStatement) -> Unit
-    ) : JsVisitorWithContextImpl() {
-
-        val replacementsInducedByWrappers = HashMap<JsFunction, Map<JsName, JsNameRef>>()
-
-        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) {
-            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)
-            }
-
-            endFunction(functionWithWrapper.function)
-        }
-
-        override fun visit(call: JsInvocation, context: JsContext<*>): Boolean {
-            if (!hasToBeInlined(call)) return true
-
-            currentNamedFunction?.let {
-                inlineCallInfos.add(JsCallInfo(call, it))
-            }
-
-            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
-        }
-
-        override fun endVisit(x: JsInvocation, ctx: JsContext<JsNode>) {
-            if (hasToBeInlined(x)) {
-                inline(x, ctx)
-            }
-            if (!inlineCallInfos.isEmpty()) {
-                if (inlineCallInfos.last.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 ->
-                    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
-            }
-        }
-
-        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
-
-                    // 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)
-                    }
-
-                    // 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
-        }
-
-        override fun doAcceptStatementList(statements: MutableList<JsStatement>) {
-            var i = 0
-
-            while (i < statements.size) {
-                val additionalStatements =
-                    ExpressionDecomposer.preserveEvaluationOrder(statements[i]) { node -> node is JsInvocation && hasToBeInlined(node) }
-                statements.addAll(i, additionalStatements)
-                i += additionalStatements.size + 1
-            }
-
-            super.doAcceptStatementList(statements)
-        }
-
-        private fun inline(call: JsInvocation, context: JsContext<JsNode>) {
-            var functionWithWrapper = 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)
-            }
-
-            val inliningContext = InliningContext(lastStatementLevelContext)
-
-            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()
-                        }
-                    })
-                }
-            }
-
-            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)
-                    }
-
-                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) {
-                    addPrevious(accept(replaceNames(statement, newReplacements)))
-                }
-
-                for ((key, value) in collectNamedFunctions(JsBlock(copiedStatements))) {
-                    if (key.staticRef is JsFunction) {
-                        key.staticRef = value
-                    }
-                }
-
-                newReplacements
-            }
-
-            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))
-                    }
-                }
-            }
-        }
-
-        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)
-                    }
-                }
-
-            }.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) {
-        call.inlineStrategy = InlineStrategy.NOT_INLINE
-        val it = inlineCallInfos.descendingIterator()
-
-        while (it.hasNext()) {
-            val callInfo = it.next()
-            val psiElement = callInfo.call.psiElement
-
-            val descriptor = callInfo.call.descriptor
-            if (psiElement != null && descriptor != null) {
-                trace.report(Errors.INLINE_CALL_CYCLE.on(psiElement, descriptor))
-            }
-
-            if (callInfo.containingFunction == calledFunction) {
-                break
-            }
-        }
-    }
-
-
-    private fun addInlinedModules(fragment: JsProgramFragment, inliner: InlinerImpl) {
+    private fun addInlinedModules(fragment: JsProgramFragment, inlinedModuleAliases: Set<JsName>) {
        val existingModules = fragment.importedModules.mapTo(IdentitySet()) { it.internalName }
-        inliner.inlinedModuleAliases.forEach {
+        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.
@@ -467,24 +41,37 @@ class JsInliner private constructor(
        }
    }

-
-    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)
-
    fun process(fragment: JsProgramFragment, existingImports: MutableMap<String, JsName>) {
        val existingNameBindings = fragment.nameBindings.associateTo(IdentityHashMap()) { it.name to it.key }

+        val dfsController = InlineDfsController(trace, functions, accessors, functionContext)
+
        val additionalDeclarations = mutableListOf<JsStatement>()
-        val inliner = InlinerImpl(existingNameBindings, existingImports, 0) {
+        val inliner = InlinerImpl(
+            existingNameBindings,
+            existingImports,
+            0,
+            dfsController,
+            inverseNameBindings,
+            functionContext
+        ) {
+            additionalDeclarations.add(it)
+        }
+
+        val inlineSuspendFnSplitter = InlineSuspendFunctionSplitter(
+            existingNameBindings,
+            existingImports,
+            0,
+            inverseNameBindings,
+            inliner.replacementsInducedByWrappers
+        ) {
            additionalDeclarations.add(it)
        }

        inliner.acceptStatement(fragment.declarationBlock)
+
+        inlineSuspendFnSplitter.accept(fragment.declarationBlock)
+
        // Mostly for the sake of post-processor
        // TODO are inline function marked with @Test possible?
        if (fragment.tests != null) {
@@ -495,12 +82,28 @@ class JsInliner private constructor(
        // TODO fix the order
        fragment.declarationBlock.statements.addAll(0, additionalDeclarations)
        fragment.nameBindings.addAll(inliner.additionalNameBindings)
+        // TODO actually bogus
+        fragment.nameBindings.addAll(inlineSuspendFnSplitter.additionalNameBindings)

-        addInlinedModules(fragment, inliner)
+        (inliner.additionalImports.asSequence() + inlineSuspendFnSplitter.additionalImports.asSequence()).forEach { (tag, e, _) ->
+            fragment.imports[tag] = e
+        }
+
+        addInlinedModules(fragment, inliner.inlinedModuleAliases)
+        addInlinedModules(fragment, inlineSuspendFnSplitter.inlinedModuleAliases)
    }

    companion object {

+        // TODO decrypt
+        // Since we compile each source file in its own context (and we may loose these context when performing incremental compilation)
+        // we don't use contexts to generate proper names for modules. Instead, we generate all necessary information during
+        // translation and rely on it here.
+        private fun buildModuleNameMap(fragments: List<JsProgramFragment>): Map<String, JsExpression> {
+            return fragments.flatMap { it.inlineModuleMap.entries }.associate { (k, v) -> k to v }
+        }
+
+
        fun process(
            reporter: JsConfig.Reporter,
            config: JsConfig,
@@ -508,6 +111,7 @@ class JsInliner private constructor(
            translationResult: AstGenerationResult
        ) {
            val functions = collectNamedFunctionsAndWrappers(translationResult.newFragments)
+
            val accessors = collectAccessors(translationResult.fragments)

            val inverseNameBindings = inverseNameBindings(*translationResult.fragments.toTypedArray())
@@ -517,17 +121,23 @@ class JsInliner private constructor(
                accessorInvocationTransformer.accept<JsGlobalBlock>(fragment.declarationBlock)
                accessorInvocationTransformer.accept<JsGlobalBlock>(fragment.initializerBlock)
            }
-            val functionReader = FunctionReader(reporter, config, translationResult.innerModuleName, translationResult.fragments)
+            val functionReader = FunctionReader(reporter, config, translationResult.innerModuleName, buildModuleNameMap(translationResult.fragments))
+
            val moduleMap = translationResult.importedModuleList.associate { it.internalName to it }

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

            for (fragment in translationResult.newFragments) {
                val block = JsBlock(fragment.declarationBlock, fragment.initializerBlock, fragment.exportBlock)
-                removeUnusedImports(block)
+                val usedImports = removeUnusedImports(block)
+                for ((key, name) in fragment.nameBindings) {
+                    if (name !in usedImports && key in fragment.imports) {
+                        fragment.imports.remove(key)
+                    }
+                }
                simplifyWrappedFunctions(block)
                removeUnusedFunctionDefinitions(block, collectNamedFunctions(block))
            }
@@ -20,7 +20,8 @@ import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.coroutineMetadata
 import org.jetbrains.kotlin.js.backend.ast.metadata.imported

-fun removeUnusedImports(root: JsNode) {
+// Returns used imports
+fun removeUnusedImports(root: JsNode): Set<JsName> {
    val collector = UsedImportsCollector()
    root.accept(collector)
    NodeRemover(JsVars::class.java) { statement ->
@@ -32,6 +33,8 @@ fun removeUnusedImports(root: JsNode) {
            false
        }
    }.accept(root)
+
+    return collector.usedImports
 }

 private class UsedImportsCollector : RecursiveJsVisitor() {
@@ -26,10 +26,15 @@ import org.jetbrains.kotlin.js.inline.util.*
 import org.jetbrains.kotlin.js.translate.context.Namer
 import org.jetbrains.kotlin.js.translate.utils.JsDescriptorUtils.getModuleName

-abstract class FunctionContext(private val functionReader: FunctionReader, private val config: JsConfig) {
-    protected abstract fun lookUpStaticFunction(functionName: JsName?): FunctionWithWrapper?
+class FunctionContext(
+    private val functionReader: FunctionReader,
+    private val config: JsConfig,
+    private val functions: Map<JsName, FunctionWithWrapper>,
+    private val accessors: Map<String, FunctionWithWrapper>
+) {
+    fun lookUpStaticFunction(functionName: JsName?): FunctionWithWrapper? = functions[functionName]

-    protected abstract fun lookUpStaticFunctionByTag(functionTag: String): FunctionWithWrapper?
+    fun lookUpStaticFunctionByTag(functionTag: String): FunctionWithWrapper? = accessors[functionTag]

    fun getFunctionDefinition(call: JsInvocation): FunctionWithWrapper {
        return getFunctionDefinitionImpl(call)!!
@@ -83,8 +88,7 @@ abstract class FunctionContext(private val functionReader: FunctionReader, priva
        /** remove ending `()` */
        val callQualifier: JsExpression = if (isCallInvocation(call)) {
            (call.qualifier as JsNameRef).qualifier!!
-        }
-        else {
+        } else {
            call.qualifier
        }

@@ -136,12 +136,6 @@ class K2JSTranslator(private val config: JsConfig) {

        val program = translationResult.buildProgram()

-        removeUnusedImports(program)
-        checkCanceled()
-
-        removeDuplicateImports(program)
-        checkCanceled()
-
        program.resolveTemporaryNames()
        checkCanceled()