JS: simplification in process

2018-12-12 19:02:48 +03:00
parent 2cbdc7ecb0
commit 75668826d3
12 changed files with 403 additions and 504 deletions
@@ -19,7 +19,6 @@ 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
@@ -38,7 +38,7 @@ private constructor(
    var resultExpr: JsNameRef? = null
    private var resultName: JsName? = null
    var breakLabel: JsLabel? = null
-    private val currentStatement = inliningContext.statementContext.currentNode
+    private val currentStatement = inliningContext.currentStatement

    init {
        invokedFunction = uncoverClosure(function.deepCopy())
@@ -112,7 +112,7 @@ private constructor(
        val breakName = JsScope.declareTemporaryName(getBreakLabel())
        this.breakLabel = JsLabel(breakName).apply { synthetic = true }

-        val visitor = ReturnReplacingVisitor(resultExpr as? JsNameRef, breakName.makeRef(), invokedFunction, call.isSuspend)
+        val visitor = ReturnReplacingVisitor(resultExpr, breakName.makeRef(), invokedFunction, call.isSuspend)
        visitor.accept(body)
    }

@@ -46,8 +46,9 @@ import java.io.StringReader
 */
 private val JS_IDENTIFIER_START = "\\p{Lu}\\p{Ll}\\p{Lt}\\p{Lm}\\p{Lo}\\p{Nl}\\\$_"
 private val JS_IDENTIFIER_PART = "$JS_IDENTIFIER_START\\p{Pc}\\p{Mc}\\p{Mn}\\d"
-private val JS_IDENTIFIER="[$JS_IDENTIFIER_START][$JS_IDENTIFIER_PART]*"
-private val DEFINE_MODULE_PATTERN = ("($JS_IDENTIFIER)\\.defineModule\\(\\s*(['\"])([^'\"]+)\\2\\s*,\\s*(\\w+)\\s*\\)").toRegex().toPattern()
+private val JS_IDENTIFIER = "[$JS_IDENTIFIER_START][$JS_IDENTIFIER_PART]*"
+private val DEFINE_MODULE_PATTERN =
+    ("($JS_IDENTIFIER)\\.defineModule\\(\\s*(['\"])([^'\"]+)\\2\\s*,\\s*(\\w+)\\s*\\)").toRegex().toPattern()
 private val DEFINE_MODULE_FIND_PATTERN = ".defineModule("

 private val specialFunctions = enumValues<SpecialFunction>().joinToString("|") { it.suggestedName }
@@ -55,9 +56,9 @@ private val specialFunctionsByName = enumValues<SpecialFunction>().associateBy {
 private val SPECIAL_FUNCTION_PATTERN = Regex("var\\s+($JS_IDENTIFIER)\\s*=\\s*($JS_IDENTIFIER)\\.($specialFunctions)\\s*;").toPattern()

 class FunctionReader(
-        private val reporter: JsConfig.Reporter,
-        private val config: JsConfig,
-        private val currentModuleName: JsName
+    private val reporter: JsConfig.Reporter,
+    private val config: JsConfig,
+    private val currentModuleName: JsName
 ) {
    /**
     * fileContent: .js file content, that contains this module definition.
@@ -70,20 +71,20 @@ class FunctionReader(
     *     The default variable is Kotlin, but it can be renamed by minifier.
     */
    class ModuleInfo(
-            val filePath: String,
-            val fileContent: String,
-            val moduleVariable: String,
-            val kotlinVariable: String,
-            val specialFunctions: Map<String, SpecialFunction>,
-            offsetToSourceMappingProvider: () -> OffsetToSourceMapping,
-            val sourceMap: SourceMap?,
-            val outputDir: File?
+        val filePath: String,
+        val fileContent: String,
+        val moduleVariable: String,
+        val kotlinVariable: String,
+        val specialFunctions: Map<String, SpecialFunction>,
+        offsetToSourceMappingProvider: () -> OffsetToSourceMapping,
+        val sourceMap: SourceMap?,
+        val outputDir: File?
    ) {
        val offsetToSourceMapping by lazy(offsetToSourceMappingProvider)

        val wrapFunctionRegex = specialFunctions.entries
-                .filter { (_, v) -> v == SpecialFunction.WRAP_FUNCTION } // TODO This is a hack! Investigate duplicates!
-                .map { Regex("\\s*${it.key}\\s*\\(\\s*").toPattern() }
+            .filter { (_, v) -> v == SpecialFunction.WRAP_FUNCTION } // TODO This is a hack! Investigate duplicates!
+            .map { Regex("\\s*${it.key}\\s*\\(\\s*").toPattern() }
    }

    private val moduleNameToInfo by lazy {
@@ -125,14 +126,14 @@ class FunctionReader(
                }

                val moduleInfo = ModuleInfo(
-                        filePath = path,
-                        fileContent = content,
-                        moduleVariable = moduleVariable,
-                        kotlinVariable = kotlinVariable,
-                        specialFunctions = specialFunctions,
-                        offsetToSourceMappingProvider = { OffsetToSourceMapping(content) },
-                        sourceMap = sourceMap,
-                        outputDir = file?.parentFile
+                    filePath = path,
+                    fileContent = content,
+                    moduleVariable = moduleVariable,
+                    kotlinVariable = kotlinVariable,
+                    specialFunctions = specialFunctions,
+                    offsetToSourceMappingProvider = { OffsetToSourceMapping(content) },
+                    sourceMap = sourceMap,
+                    outputDir = file?.parentFile
                )

                result.put(moduleName, moduleInfo)
@@ -164,37 +165,60 @@ class FunctionReader(
        override fun toString() = text.substring(offset)
    }

-    private object NotFoundMarker : Any()
+    object NotFoundMarker

    private val functionCache = object : SLRUCache<CallableDescriptor, Any>(50, 50) {
-        // LibraryInlineFunctionDefinition | NotFoundMarker
        override fun createValue(key: CallableDescriptor): Any =
            readFunction(key) ?: NotFoundMarker
    }

-    operator fun get(descriptor: CallableDescriptor): LibraryInlineFunctionDefinition? {
-        val existed = functionCache.get(descriptor)
-        return if (existed === NotFoundMarker) null else existed as LibraryInlineFunctionDefinition
+    operator fun get(descriptor: CallableDescriptor, fragment: JsProgramFragment): FunctionWithWrapper? {
+        return functionCache.get(descriptor).let {
+            if (it === NotFoundMarker) null else {
+                val (fn, info) = it as Pair<*, *>
+                renameModules(descriptor, fn as FunctionWithWrapper, info as ModuleInfo, fragment)
+            }
+        }
    }

-    private fun readFunction(descriptor: CallableDescriptor): LibraryInlineFunctionDefinition? {
+    private fun renameModules(
+        descriptor: CallableDescriptor,
+        fn: FunctionWithWrapper,
+        info: ModuleInfo,
+        fragment: JsProgramFragment
+    ): FunctionWithWrapper {
+        val tag = Namer.getFunctionTag(descriptor, config)
+        val moduleReference = fragment.inlineModuleMap[tag]?.deepCopy() ?: currentModuleName.makeRef()
+        val allDefinedNames = collectDefinedNamesInAllScopes(fn.function)
+        val replacements = hashMapOf(
+            info.moduleVariable to moduleReference,
+            info.kotlinVariable to Namer.kotlinObject()
+        )
+        replaceExternalNames(fn.function, replacements, allDefinedNames)
+        val wrapperStatements = fn.wrapperBody?.statements?.filter { it !is JsReturn }
+        wrapperStatements?.forEach { replaceExternalNames(it, replacements, allDefinedNames) }
+
+        return fn
+    }
+
+    private fun readFunction(descriptor: CallableDescriptor): Pair<FunctionWithWrapper, ModuleInfo>? {
        val moduleName = getModuleName(descriptor)

        if (moduleName !in moduleNameToInfo.keys()) return null

        for (info in moduleNameToInfo[moduleName]) {
            val function = readFunctionFromSource(descriptor, info)
-            if (function != null) return function
+            if (function != null) return function to info
        }

        return null
    }

    // TODO move renamings to a proper place
-    private fun readFunctionFromSource(descriptor: CallableDescriptor, info: ModuleInfo): LibraryInlineFunctionDefinition? {
+    private fun readFunctionFromSource(descriptor: CallableDescriptor, info: ModuleInfo): FunctionWithWrapper? {
        val source = info.fileContent
        var tag = Namer.getFunctionTag(descriptor, config)
-        val tagForModule = tag
+//        val tagForModule = tag
        var index = source.indexOf(tag)

        // Hack for compatibility with old versions of stdlib
@@ -226,16 +250,14 @@ class FunctionReader(
        val position = info.offsetToSourceMapping[offset]
        val jsScope = JsRootScope(JsProgram())
        val functionExpr = try {
-            parseFunction(source, info.filePath, position, offset, ThrowExceptionOnErrorReporter, jsScope) ?:
-            return null
+            parseFunction(source, info.filePath, position, offset, ThrowExceptionOnErrorReporter, jsScope) ?: return null
        } catch (t: Throwable) {
            throw Error("Exception while reading function '$tag' from ${info.filePath}", t)
        }
        functionExpr.fixForwardNameReferences()
        val (function, wrapper) = if (isWrapped) {
            InlineMetadata.decomposeWrapper(functionExpr) ?: return null
-        }
-        else {
+        } else {
            FunctionWithWrapper(functionExpr, null)
        }
 //        val moduleReference = moduleNameMap[tagForModule]?.deepCopy() ?: currentModuleName.makeRef()
@@ -262,8 +284,8 @@ class FunctionReader(
        markSpecialFunctions(function, allDefinedNames, info, jsScope)

        val namesWithoutSideEffects = wrapperStatements.orEmpty().asSequence()
-                .flatMap { collectDefinedNames(it).asSequence() }
-                .toSet()
+            .flatMap { collectDefinedNames(it).asSequence() }
+            .toSet()
        function.accept(object : RecursiveJsVisitor() {
            override fun visitNameRef(nameRef: JsNameRef) {
                if (nameRef.name in namesWithoutSideEffects && nameRef.qualifier == null) {
@@ -279,7 +301,7 @@ class FunctionReader(
            }
        }

-        return LibraryInlineFunctionDefinition(tagForModule, FunctionWithWrapper(function, wrapper), info)
+        return FunctionWithWrapper(function, wrapper)
    }

    private fun markSpecialFunctions(function: JsFunction, allDefinedNames: Set<JsName>, info: ModuleInfo, scope: JsScope) {
@@ -365,7 +387,7 @@ private fun JsFunction.markInlineArguments(descriptor: CallableDescriptor) {
        inlineFuns.add(paramsJs[i + offset].name)
    }

-    val visitor = object: JsVisitorWithContextImpl() {
+    val visitor = object : JsVisitorWithContextImpl() {
        override fun endVisit(x: JsInvocation, ctx: JsContext<*>) {
            val qualifier: JsExpression? = if (isCallInvocation(x)) {
                (x.qualifier as? JsNameRef)?.qualifier
@@ -385,7 +407,7 @@ private fun JsFunction.markInlineArguments(descriptor: CallableDescriptor) {
 }

 private fun replaceExternalNames(node: JsNode, replacements: Map<String, JsExpression>, definedNames: Set<JsName>) {
-    val visitor = object: JsVisitorWithContextImpl() {
+    val visitor = object : JsVisitorWithContextImpl() {
        override fun endVisit(x: JsNameRef, ctx: JsContext<JsNode>) {
            if (x.qualifier != null || x.name in definedNames) return

@@ -407,5 +429,5 @@ private class ShallowSubSequence(private val underlying: CharSequence, private v
    }

    override fun subSequence(startIndex: Int, endIndex: Int): CharSequence =
-            ShallowSubSequence(underlying, start + startIndex, start + endIndex)
+        ShallowSubSequence(underlying, start + startIndex, start + endIndex)
 }
@@ -5,66 +5,8 @@

 package org.jetbrains.kotlin.js.inline

-import org.jetbrains.kotlin.js.backend.ast.JsProgramFragment
 import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper

-sealed class InlineFunctionDefinition {
-    abstract val functionWithWrapper: FunctionWithWrapper
-
-    abstract val tag: String?
-
-    open fun process() {}
-
-// imports / nameBindings
-    // modules
-    // renameFor scope?
-
-
-}
-
-
-// Current module, new fragment. Expressed through `defineInlineFunction`. Should be self-contained
-open class PublicInlineFunctionDefinition(
-    override val tag: String,
-    override val functionWithWrapper: FunctionWithWrapper,
-    val fragment: JsProgramFragment,
-    val scope: ProgramFragmentInliningScope
-) : InlineFunctionDefinition() {
-    override fun process() {
-        scope.process()
-    }
-}
-
-// Current module, from Binary AST
-class BinaryInlineFunctionDefinition(
-    override val tag: String,
-    override val functionWithWrapper: FunctionWithWrapper,
-    val fragment: JsProgramFragment
-): InlineFunctionDefinition() {
-
-
-}
-
-// Current module, new fragment, used within scope only. Private functions, lambdas.
-class LocalInlineFunctionDefinition(
-    override val functionWithWrapper: FunctionWithWrapper,
-    val scope: InliningScope
-) : InlineFunctionDefinition() {
-
-    override val tag = null
-
-    override fun process() {
-        // TODO this is incorrect!
-        scope.process()
-    }
-}
-
-// Deserialized from a binary dependency (<module>.js file)
-open class LibraryInlineFunctionDefinition(
-    override val tag: String,
-    override val functionWithWrapper: FunctionWithWrapper,
-    val moduleInfo: FunctionReader.ModuleInfo
-): InlineFunctionDefinition() {
-
-
-}
+class InlineFunctionDefinition(
+    val fn: FunctionWithWrapper,
+    val tag: String?)
@@ -47,10 +47,10 @@ class InlineSuspendFunctionSplitter(
                    val statementContext = lastStatementLevelContext

                    // This function will be exported to JS
-                    val function = scope.importFunctionDefinition(PublicInlineFunctionDefinition(inlineMetadata.tag.value, inlineMetadata.function, scope.fragment, scope))
+                    val function = scope.importFunctionDefinition(InlineFunctionDefinition(inlineMetadata.function, inlineMetadata.tag.value))

                    // Original function should be not be transformed into a state machine
-                    f.function.setName(null)
+                    f.function.name = null
                    f.function.coroutineMetadata = null
                    f.function.isInlineableCoroutineBody = true

@@ -12,12 +12,9 @@ import org.jetbrains.kotlin.js.backend.ast.JsInvocation
 import org.jetbrains.kotlin.js.backend.ast.metadata.descriptor
 import org.jetbrains.kotlin.js.backend.ast.metadata.inlineStrategy
 import org.jetbrains.kotlin.js.backend.ast.metadata.psiElement
-import org.jetbrains.kotlin.js.inline.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.*

@@ -66,13 +63,24 @@ class InlinerCycleReporter(
    }


-    // Return true iff the definition should be visited by the inliner
-    fun shouldProcess(definition: FunctionWithWrapper, call: JsInvocation): Boolean {
-
+    fun <T> withInlining(call: JsInvocation, body: () -> T): T {
        currentNamedFunction?.let {
            inlineCallInfos.add(JsCallInfo(call, it))
        }

+        val result = body()
+
+        if (!inlineCallInfos.isEmpty()) {
+            if (inlineCallInfos.last.call == call) {
+                inlineCallInfos.removeLast()
+            }
+        }
+
+        return result
+    }
+
+    // Return true iff the definition should be visited by the inliner
+    fun shouldProcess(definition: FunctionWithWrapper, call: JsInvocation?): Boolean {
        if (definition.function in inProcessFunctions) {
            reportInlineCycle(call, definition.function)
        } else if (definition.function !in processedFunctions) {
@@ -82,16 +90,8 @@ class InlinerCycleReporter(
        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
+    private fun reportInlineCycle(call: JsInvocation?, calledFunction: JsFunction) {
+        call?.inlineStrategy = InlineStrategy.NOT_INLINE
        val it = inlineCallInfos.descendingIterator()

        while (it.hasNext()) {
@@ -6,76 +6,71 @@
 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.clean.FunctionPostProcessor
-import org.jetbrains.kotlin.js.inline.clean.removeUnusedLocalFunctionDeclarations
 import org.jetbrains.kotlin.js.inline.context.FunctionContext
-import org.jetbrains.kotlin.js.inline.context.InliningContext
-import org.jetbrains.kotlin.js.inline.util.refreshLabelNames
-import org.jetbrains.kotlin.js.translate.declaration.transformSpecialFunctionsToCoroutineMetadata
+import org.jetbrains.kotlin.js.inline.util.FunctionWithWrapper
+import org.jetbrains.kotlin.js.inline.util.extractFunction
+import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils

 // TODO stateless?
 class InlinerImpl(
-    val cycleReporter: InlinerCycleReporter,
-    val functionContext: FunctionContext,
-    // TODO other way around? Need to find a correct inliner by function declaration.
+    val jsInliner: JsInliner,
    val scope: InliningScope
 ) : JsVisitorWithContextImpl() {

-    override fun visit(function: JsFunction, context: JsContext<*>): Boolean {
-        functionContext.functionsByFunctionNodes[function]?.let { (function, wrapper) ->
-            visit(function, wrapper)
-            return false
+    override fun visit(x: JsBinaryOperation, ctx: JsContext<*>): Boolean {
+        val assignment = JsAstUtils.decomposeAssignment(x)
+        if (assignment != null) {
+            val (left, right) = assignment
+            if (left is JsNameRef) {
+                val name = left.name
+                if (name != null) {
+                    extractFunction(right)?.let { function ->
+                        jsInliner.process(InlineFunctionDefinition(function, null), null, scope)
+                    }
+                }
+            }
        }
-        visit(function, null)
-        return true
-    }

-    override fun visit(x: JsBlock, ctx: JsContext<*>): Boolean {
-        // TODO Seems like a very roundabout way. Probably should reuse same approach as in CoroutineTransformer and ...Splitter
-        // TODO That approach might be missing inline properties. Check!
-        functionContext.functionsByWrapperNodes[x]?.let { (function, wrapper) ->
-            visit(function, wrapper)
-            return false
-        }
        return super.visit(x, ctx)
    }

-    private fun visit(function: JsFunction, wrapperBody: JsBlock?) {
-        cycleReporter.startFunction(function)
-
-        // TODO Different visitors?
-        if (wrapperBody != null && scope is ProgramFragmentInliningScope) {
-            PublicInlineFunctionInliningScope(scope.fragment, cycleReporter, functionContext, function, wrapperBody).process()
-        } else {
-            // TODO this is still not super-clear
-            accept(function.body)
+    override fun visit(x: JsVars.JsVar, ctx: JsContext<*>): Boolean {
+        val initializer = x.initExpression
+        val name = x.name
+        if (initializer != null && name != null) {
+            extractFunction(initializer)?.let { function ->
+                jsInliner.process(InlineFunctionDefinition(function, null), null, scope)
+            }
        }

-        // Cleanup
-        refreshLabelNames(function.body, function.scope)
-        removeUnusedLocalFunctionDeclarations(function)
-        FunctionPostProcessor(function).apply()
-
-        cycleReporter.endFunction(function)
+        return super.visit(x, ctx)
    }

+    override fun visit(x: JsInvocation, ctx: JsContext<*>): Boolean {
+        InlineMetadata.decompose(x)?.let {
+            jsInliner.process(InlineFunctionDefinition(it.function, it.tag.value), x, scope)
+        }
+
+        return super.visit(x, ctx)
+    }
+
+
    override fun endVisit(call: JsInvocation, ctx: JsContext<JsNode>) {
        if (hasToBeInlined(call)) {
+            val (inlineableBody, resultExpression) = jsInliner.inline(scope, call, lastStatementLevelContext.currentNode)

-            val definition = functionContext.getFunctionDefinition(call, scope)
+            lastStatementLevelContext.addPrevious(JsAstUtils.flattenStatement(inlineableBody))

-            if (cycleReporter.shouldProcess(definition.functionWithWrapper, call)) {
-                definition.process()
+            // Assumes, that resultExpression == null, when result is not needed.
+            // @see FunctionInlineMutator.isResultNeeded()
+            if (resultExpression == null) {
+                lastStatementLevelContext.removeMe()
+            } else {
+                ctx.replaceMe(resultExpression)
            }
-
-            inline(call, definition, ctx)
        }
-
-        cycleReporter.endVisit(call)
    }

    // TODO This could be extracted into a separate pass.
@@ -94,62 +89,9 @@ class InlinerImpl(
        super.doAcceptStatementList(statements)
    }

-    // TODO a lot of code... Probably a bad sign
-    private fun inline(call: JsInvocation, definition: InlineFunctionDefinition, context: JsContext<JsNode>) {
-
-        // ---------------
-        // This should be isolated
-
-        val function = scope.importFunctionDefinition(definition)
-        // TODO This should be done inside the importer
-        function.body = transformSpecialFunctionsToCoroutineMetadata(function.body)
-
-        // -------------------
-
-        val statementContext = lastStatementLevelContext
-
-        val (inlineableBody, resultExpression) =
-                FunctionInlineMutator.getInlineableCallReplacement(call, function, InliningContext(statementContext))
-
-        // body of inline function can contain call to lambdas that need to be inlined
-        val inlineableBodyWithLambdasInlined = accept(inlineableBody)
-        assert(inlineableBody === inlineableBodyWithLambdasInlined)
-
-        patchReturnsFromSecondaryConstructor(inlineableBody)
-
-        statementContext.addPrevious(JsAstUtils.flattenStatement(inlineableBody))
-
-        /*
-         * Assumes, that resultExpression == null, when result is not needed.
-         * @see FunctionInlineMutator.isResultNeeded()
-         */
-        if (resultExpression == null) {
-            statementContext.removeMe()
-            return
-        }
-
-        // TODO Why accept? Seems unnecessary... Some inline call in the qualifier? Shouldn't this be done along with the lambdas then?
-        accept(resultExpression)?.let {
-            it.synthetic = true
-            context.replaceMe(it)
-        }
-    }
-
-    private fun patchReturnsFromSecondaryConstructor(inlineableBody: JsStatement) {
-        // Support non-local return from secondary constructor
-        // Returns from secondary constructors should return `$this` object.
-        // TODO This seems brittle
-        cycleReporter.currentNamedFunction?.forcedReturnVariable?.let { returnVariable ->
-            inlineableBody.accept(object : RecursiveJsVisitor() {
-                override fun visitReturn(x: JsReturn) {
-                    x.expression = returnVariable.makeRef()
-                }
-            })
-        }
-    }

    private fun hasToBeInlined(call: JsInvocation): Boolean {
        val strategy = call.inlineStrategy
-        return if (strategy == null || !strategy.isInline) false else functionContext.hasFunctionDefinition(call, scope)
+        return if (strategy == null || !strategy.isInline) false else jsInliner.functionContext.hasFunctionDefinition(call, scope)
    }
 }
@@ -12,75 +12,134 @@ import org.jetbrains.kotlin.js.backend.ast.metadata.staticRef
 import org.jetbrains.kotlin.js.inline.clean.removeUnusedFunctionDefinitions
 import org.jetbrains.kotlin.js.inline.clean.removeUnusedImports
 import org.jetbrains.kotlin.js.inline.clean.simplifyWrappedFunctions
-import org.jetbrains.kotlin.js.inline.context.FunctionContext
 import org.jetbrains.kotlin.js.inline.util.*
+import org.jetbrains.kotlin.js.translate.declaration.transformSpecialFunctionsToCoroutineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import java.util.*

-
 // Handles interpreting an inline function in terms of the current context.
 // Either an program fragment, or a public inline function
 sealed class InliningScope {

-    private val cache = mutableMapOf<String, Map<JsName, JsNameRef>>()
+    abstract val fragment: JsProgramFragment

-    protected fun computeIfAbsent(tag: String?, fn: () -> Map<JsName, JsNameRef>): Map<JsName, JsNameRef> {
-        if (tag == null) return fn()
+    abstract fun addInlinedDeclaration(tag: String?, declaration: JsStatement)

-        return cache.computeIfAbsent(tag) { fn() }
+    abstract fun hasImport(tag: String): JsName?
+
+    abstract fun addImport(tag: String, vars: JsVars)
+
+    open fun preprocess(statement: JsStatement) {}
+
+    abstract fun update()
+
+    private val publicFunctionCache = mutableMapOf<String, Map<JsName, JsNameRef>>()
+
+    private val localFunctionCache = mutableMapOf<JsFunction, Map<JsName, JsNameRef>>()
+
+    private fun computeIfAbsent(tag: String?, function: JsFunction, fn: () -> Map<JsName, JsNameRef>): Map<JsName, JsNameRef> {
+        if (tag == null) return localFunctionCache.computeIfAbsent(function) { fn() }
+
+        return publicFunctionCache.computeIfAbsent(tag) { fn() }
    }

-    abstract fun importFunctionDefinition(f: InlineFunctionDefinition): JsFunction
+    fun importFunctionDefinition(definition: InlineFunctionDefinition): JsFunction {
+        // Apparently we should avoid this trick when we implement fair support for crossinline
+        // That's because crossinline lambdas inline into the declaration block and specialize those.
+        val replacements = computeIfAbsent(definition.tag, definition.fn.function) {
+            val newReplacements = HashMap<JsName, JsNameRef>()

-    abstract fun process()
+            val copiedStatements = ArrayList<JsStatement>()
+            val importStatements = mutableMapOf<JsVars, String>()

-    abstract val fragment: JsProgramFragment
+            definition.fn.wrapperBody?.let {
+                it.statements.asSequence()
+                    .filterNot { it is JsReturn }
+                    .map { it.deepCopy() }
+                    .forEach { statement ->
+                        preprocess(statement)
+
+                        if (statement is JsVars) {
+                            val tag = getImportTag(statement)
+                            if (tag != null) {
+                                val name = statement.vars[0].name
+                                val existingName = name.localAlias ?: hasImport(tag) ?: JsScope.declareTemporaryName(name.ident).also {
+                                    it.copyMetadataFrom(name)
+                                    importStatements[statement] = tag
+                                }
+
+                                if (name !== existingName) {
+                                    val replacement = JsAstUtils.pureFqn(existingName, null)
+                                    newReplacements[name] = replacement
+                                }
+                            }
+                        }
+
+                        copiedStatements.add(statement)
+                    }
+            }
+
+            copiedStatements.asSequence()
+                .flatMap { node -> collectDefinedNamesInAllScopes(node).asSequence() }
+                .filter { name -> !newReplacements.containsKey(name) }
+                .forEach { name ->
+                    val alias = JsScope.declareTemporaryName(name.ident)
+                    alias.copyMetadataFrom(name)
+                    val replacement = JsAstUtils.pureFqn(alias, null)
+                    newReplacements[name] = replacement
+                }
+
+            // Apply renaming and restore the static ref links
+            JsBlock(copiedStatements).let {
+                replaceNames(it, newReplacements)
+
+                // Restore the staticRef links
+                for ((key, value) in collectNamedFunctions(it)) {
+                    if (key.staticRef is JsFunction) {
+                        key.staticRef = value
+                    }
+                }
+            }
+
+            copiedStatements.forEach {
+                if (it is JsVars && it in importStatements) {
+                    addImport(importStatements[it]!!, it)
+                } else {
+                    addInlinedDeclaration(definition.tag, it)
+                }
+            }
+
+            newReplacements
+        }
+
+        val paramMap = definition.fn.function.parameters.associate {
+            val alias = JsScope.declareTemporaryName(it.name.ident)
+            alias.copyMetadataFrom(it.name)
+            it.name to JsAstUtils.pureFqn(alias, null)
+        }
+
+        val result = definition.fn.function.deepCopy()
+
+        replaceNames(result, replacements)
+        replaceNames(result, paramMap)
+
+        result.body = transformSpecialFunctionsToCoroutineMetadata(result.body)
+
+        return result
+    }
 }

 class ProgramFragmentInliningScope(
-    override val fragment: JsProgramFragment,
-    val functionContext: FunctionContext,
-    val rootInliner: JsInliner
+    override val fragment: JsProgramFragment
 ) : InliningScope() {

-    private val existingModules = fragment.importedModules.mapTo(IdentitySet()) { it.internalName }
+    private val existingModules = fragment.importedModules.associateTo(mutableMapOf()) { it.key to it }

-    private val existingImports = fragment.nameBindings.associateTo(mutableMapOf()) { it.key to it.name }
-
-    private val existingNameBindings = fragment.nameBindings.associateTo(IdentityHashMap()) { it.name to it.key }
+    private val existingBindings = fragment.nameBindings.associateTo(mutableMapOf()) { it.key to it.name }

    private val additionalDeclarations = mutableListOf<JsStatement>()

-    private var processed = false
-
-    override fun process() {
-        if (!processed) {
-            // TODO is this even needed?
-            processed = true
-
-            val inliner = InlinerImpl(rootInliner.cycleReporter, functionContext, this)
-
-            // TODO any way and/or need to visit everything inside the fragment?
-            inliner.acceptStatement(fragment.declarationBlock)
-
-            // TODO Atm it's placed after inliner in order not to perform the body inlining twice. Is that OK?
-            // Ideally it could be moved to the coroutine transformers. The info regarding which inline function wrappers have been imported
-            // on top level should be persisted for that sake. Also it going to be needed in order to avoid duplicate code.
-            InlineSuspendFunctionSplitter(this).accept(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)
-            }
-            // TODO wrap in a function in order to do the post-processing
-            inliner.acceptStatement(fragment.initializerBlock)
-
-            updateProgramFragment()
-        }
-    }
-
-    private fun updateProgramFragment() {
+    override fun update() {
        // TODO fix the order
        // TODO this probably will be replaced with a special tag -> block map for the imported stuff, so that we can merge same imports.
        // TODO in that case this method will become obsolete
@@ -110,125 +169,25 @@ class ProgramFragmentInliningScope(
        }
    }

-    private fun addInlinedModule(module: JsImportedModule) {
-//        if (moduleName !in existingModules) {
-//            fragment.importedModules.add(moduleMap[moduleName]!!.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 hasImport(tag: String): JsName? = existingBindings[tag]
+
+    override fun addImport(tag: String, vars: JsVars) {
+        val name = vars.vars[0].name
+        val expr = vars.vars[0].initExpression
+        fragment.imports[tag] = expr
+        fragment.nameBindings.add(JsNameBinding(tag, name))
+        existingBindings[tag] = name
    }

-    private fun addImport(tag: String, e: JsExpression) {
-        fragment.imports[tag] = e
-    }
-
-    private fun addNameBinding(binding: JsNameBinding) {
-        fragment.nameBindings.add(binding)
-        existingNameBindings[binding.name] = binding.key
-    }
-
-    override fun importFunctionDefinition(f: InlineFunctionDefinition): JsFunction {
-
-        // Apparently we should avoid this trick when we implement fair support for crossinline
-        // That's because crossinline lambdas inline into the declaration block and specialize those.
-        val replacements = computeIfAbsent(f.tag) {
-            val newReplacements = HashMap<JsName, JsNameRef>()
-
-            val copiedStatements = ArrayList<JsStatement>()
-            val importStatements = ArrayList<Pair<String, JsVars.JsVar>>()
-
-            f.functionWithWrapper.wrapperBody?.let {
-                it.statements.asSequence()
-                    .filterNot { it is JsReturn }
-                    .map { it.deepCopy() }
-                    .forEach { statement ->
-                        replaceExpressionsWithLocalAliases(statement)
-
-                        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? = name.localAlias
-                                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)
-                    }
-            }
-
-            (importStatements.asSequence().map { JsVars(it.second) } + copiedStatements.asSequence())
-                .flatMap { node -> collectDefinedNamesInAllScopes(node).asSequence() }
-                .filter { name -> !newReplacements.containsKey(name) }
-                .forEach { name ->
-                    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)
-                // TODO shouldn't this be done at `existingImports.computeIfAbsent` moment?
-                addImport(tag, renamed.initExpression)
-
-                addNameBinding(JsNameBinding(tag, renamed.name))
-            }
-
-            if (f.tag != null) {
-                fragment.inlinedFunctionWrappers[f.tag!!] = JsGlobalBlock().also {
-                    copiedStatements.mapTo(it.statements) { replaceNames(it, newReplacements) }
-                }
-            } else {
-                // TODO Handle it better?
-                for (statement in copiedStatements) {
-                    additionalDeclarations.add(replaceNames(statement, newReplacements))
-                }
-            }
-
-            // TODO shouldn't this be moved to renamer?
-            for ((key, value) in collectNamedFunctions(JsBlock(copiedStatements))) {
-                if (key.staticRef is JsFunction) {
-                    key.staticRef = value
-                }
-            }
-
-            newReplacements
+    override fun addInlinedDeclaration(tag: String?, declaration: JsStatement) {
+        if (tag != null) {
+            fragment.inlinedFunctionWrappers.computeIfAbsent(tag) { JsGlobalBlock() }.statements.add(declaration)
+        } else {
+            additionalDeclarations.add(declaration)
        }
-
-        val paramMap = f.functionWithWrapper.function.parameters.associate {
-            val alias = JsScope.declareTemporaryName(it.name.ident)
-            alias.copyMetadataFrom(it.name)
-            it.name to JsAstUtils.pureFqn(alias, null)
-        }
-
-        val result = f.functionWithWrapper.function.deepCopy()
-
-        replaceNames(result, replacements)
-        replaceNames(result, paramMap)
-
-        return result
    }

-    private fun replaceExpressionsWithLocalAliases(statement: JsStatement) {
+    override fun preprocess(statement: JsStatement) {
        object : JsVisitorWithContextImpl() {
            override fun endVisit(x: JsNameRef, ctx: JsContext<JsNode>) {
                replaceIfNecessary(x, ctx)
@@ -241,97 +200,44 @@ class ProgramFragmentInliningScope(
            private fun replaceIfNecessary(expression: JsExpression, ctx: JsContext<JsNode>) {
                val alias = expression.localAlias
                if (alias != null) {
-                    // TODO wrong!
-                    ctx.replaceMe(alias.internalName.makeRef())
-                    addInlinedModule(alias)
+                    ctx.replaceMe(addInlinedModule(alias).makeRef())
                }
            }

        }.accept(statement)
    }
+
+    private fun addInlinedModule(module: JsImportedModule): JsName {
+        return existingModules.computeIfAbsent(module.key) {
+            // Copy so that the Merger.kt doesn't operate on the same instance in different fragments.
+            // TODO What about nameBindings?
+            JsImportedModule(module.externalName, module.internalName, module.plainReference).also {
+                fragment.importedModules.add(it)
+            }
+        }.internalName
+    }
 }

 class PublicInlineFunctionInliningScope(
-    override val fragment: JsProgramFragment,
-    cycleReporter: InlinerCycleReporter,
-    functionContext: FunctionContext,
    val function: JsFunction,
-    val wrapperBody: JsBlock
+    val wrapperBody: JsBlock,
+    override val fragment: JsProgramFragment
 ) : InliningScope() {
-
    val additionalStatements = mutableListOf<JsStatement>()

-    val innerInliner = InlinerImpl(
-        cycleReporter,
-        functionContext,
-        this
-    )
+    override fun addInlinedDeclaration(tag: String?, declaration: JsStatement) {
+        additionalStatements.add(declaration)
+    }

-    override fun process() {
-        for (statement in wrapperBody.statements) {
-            if (statement !is JsReturn) {
-                innerInliner.acceptStatement(statement)
-            } else {
-                innerInliner.accept((statement.expression as JsFunction).body)
-            }
-        }
+    override fun hasImport(tag: String): JsName? {
+        return null // TODO
+    }

-        // TODO keep order
+    override fun addImport(tag: String, vars: JsVars) {
+        additionalStatements.add(vars) // TODO
+    }
+
+    override fun update() {
        wrapperBody.statements.addAll(0, additionalStatements)
    }
-
-    private fun addPrevious(statement: JsStatement) {
-        // TODO Is this correct?
-        additionalStatements.add(innerInliner.accept(statement))
-    }
-
-    override fun importFunctionDefinition(f: InlineFunctionDefinition): JsFunction {
-        // TODO Decrypt the comment below
-        // Apparently we should avoid this trick when we implement fair support for crossinline
-        val replacements = computeIfAbsent(f.tag) {
-            val newReplacements = HashMap<JsName, JsNameRef>()
-
-            // TODO Why don't we collect existing imports?
-
-            val copiedStatements = f.functionWithWrapper.wrapperBody!!.statements.asSequence()
-                .filterNot { it is JsReturn }
-                .map { it.deepCopy() }.toList()
-
-            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 = JsAstUtils.pureFqn(alias, null)
-                newReplacements[name] = replacement
-            }
-
-            for (statement in copiedStatements) {
-                addPrevious(replaceNames(statement, newReplacements))
-            }
-
-            for ((key, value) in collectNamedFunctions(JsBlock(copiedStatements))) {
-                if (key.staticRef is JsFunction) {
-                    key.staticRef = value
-                }
-            }
-
-            newReplacements
-        }
-
-        val paramMap = f.functionWithWrapper.function.parameters.associate {
-            val alias = JsScope.declareTemporaryName(it.name.ident)
-            alias.copyMetadataFrom(it.name)
-            it.name to JsAstUtils.pureFqn(alias, null)
-        }
-
-        val result = f.functionWithWrapper.function.deepCopy()
-
-        replaceNames(result, replacements)
-        replaceNames(result, paramMap)
-
-        return result
-    }
 }
@@ -7,8 +7,15 @@ package org.jetbrains.kotlin.js.inline

 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.js.backend.ast.*
+import org.jetbrains.kotlin.js.backend.ast.metadata.forcedReturnVariable
+import org.jetbrains.kotlin.js.backend.ast.metadata.synthetic
 import org.jetbrains.kotlin.js.config.JsConfig
+import org.jetbrains.kotlin.js.inline.clean.FunctionPostProcessor
+import org.jetbrains.kotlin.js.inline.clean.removeUnusedLocalFunctionDeclarations
 import org.jetbrains.kotlin.js.inline.context.FunctionContext
+import org.jetbrains.kotlin.js.inline.context.InliningContext
+import org.jetbrains.kotlin.js.inline.util.refreshLabelNames
+import org.jetbrains.kotlin.js.translate.declaration.transformSpecialFunctionsToCoroutineMetadata

 import org.jetbrains.kotlin.js.translate.general.AstGenerationResult

@@ -35,7 +42,92 @@ class JsInliner(

    fun process() {
        for (fragment in translationResult.newFragments) {
-            functionContext.scopeForFragment(fragment).process()
+            process(fragment)
+        }
+    }
+
+    fun process(fragment: JsProgramFragment) {
+        val fragmentScope = functionContext.scopeForFragment(fragment) ?: return
+
+        // TODO any way and/or need to visit everything inside the fragment?
+        fragmentScope.process(fragment.declarationBlock)
+
+        // TODO Atm it's placed after inliner in order not to perform the body inlining twice. Is that OK?
+        // Ideally it could be moved to the coroutine transformers. The info regarding which inline function wrappers have been imported
+        // on top level should be persisted for that sake. Also it going to be needed in order to avoid duplicate code.
+        InlineSuspendFunctionSplitter(fragmentScope).accept(fragment.declarationBlock)
+
+        // Mostly for the sake of post-processor
+        // TODO are inline function marked with @Test possible?
+        fragment.tests?.let { fragmentScope.process(it) }
+
+        // TODO wrap in a function in order to do the post-processing
+        fragmentScope.process(fragment.initializerBlock)
+
+        fragmentScope.update()
+    }
+
+    fun process(inlineFn: InlineFunctionDefinition, call: JsInvocation?, containingScope: InliningScope) {
+        if (cycleReporter.shouldProcess(inlineFn.fn, call)) {
+            val (function, wrapperBody) = inlineFn.fn
+
+            cycleReporter.startFunction(function)
+
+            if (wrapperBody != null) {
+                val scope = PublicInlineFunctionInliningScope(function, wrapperBody, containingScope.fragment)
+                scope.process(wrapperBody)
+                scope.update()
+            } else {
+                containingScope.process(function)
+            }
+
+            // Cleanup
+            refreshLabelNames(function.body, function.scope)
+            removeUnusedLocalFunctionDeclarations(function)
+            FunctionPostProcessor(function).apply()
+
+            cycleReporter.endFunction(function)
+        }
+    }
+
+    private fun InliningScope.process(node: JsNode) {
+        InlinerImpl(this@JsInliner, this).accept(node)
+    }
+
+    // TODO a lot of code... Probably a bad sign
+    fun inline(scope: InliningScope, call: JsInvocation, currentStatement: JsStatement?): InlineableResult {
+        val definition = functionContext.getFunctionDefinition(call, scope)
+
+        return cycleReporter.withInlining(call) {
+
+            val function = scope.importFunctionDefinition(definition)
+
+            val inliningContext = InliningContext(currentStatement)
+
+            val (inlineableBody, resultExpression) = FunctionInlineMutator.getInlineableCallReplacement(call, function, inliningContext)
+
+            // body of inline function can contain call to lambdas that need to be inlined
+            scope.process(inlineableBody)
+
+            patchReturnsFromSecondaryConstructor(inlineableBody)
+
+            // TODO shouldn't we process the resultExpression qualifier along with the lambda inlining?
+            resultExpression?.synthetic = true
+
+            InlineableResult(JsBlock(inliningContext.previousStatements + inlineableBody), resultExpression)
+        }
+    }
+
+    private fun patchReturnsFromSecondaryConstructor(inlineableBody: JsStatement) {
+        // Support non-local return from secondary constructor
+        // Returns from secondary constructors should return `$this` object.
+        // TODO This seems brittle
+        cycleReporter.currentNamedFunction?.forcedReturnVariable?.let { returnVariable ->
+            inlineableBody.accept(object : RecursiveJsVisitor() {
+                override fun visitReturn(x: JsReturn) {
+                    x.expression = returnVariable.makeRef()
+                }
+            })
        }
    }
 }
@@ -16,15 +16,12 @@

 package org.jetbrains.kotlin.js.inline.context

-import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.js.backend.ast.*
 import org.jetbrains.kotlin.js.backend.ast.metadata.isCallableReference
 import org.jetbrains.kotlin.js.backend.ast.metadata.descriptor
-import org.jetbrains.kotlin.js.config.JsConfig
 import org.jetbrains.kotlin.js.inline.*
 import org.jetbrains.kotlin.js.inline.util.*
 import org.jetbrains.kotlin.js.translate.context.Namer
-import org.jetbrains.kotlin.js.translate.general.AstGenerationResult
 import java.util.HashMap

 class FunctionContext(
@@ -52,8 +49,6 @@ class FunctionContext(
        return getFunctionDefinitionImpl(call, scope) != null
    }

-    val functionsByWrapperNodes = HashMap<JsBlock, FunctionWithWrapper>()
-
    val functionsByFunctionNodes = HashMap<JsFunction, FunctionWithWrapper>()

    /**
@@ -88,17 +83,13 @@ class FunctionContext(
     */
    private fun getFunctionDefinitionImpl(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
        // Ensure we have the local function information
+        // TODO is this necessary?
        loadFragment(scope.fragment)

-        val descriptor = call.descriptor
-        if (descriptor != null) {
-            return lookUpFunctionDirect(descriptor) ?: lookUpFunctionIndirect(call, scope) ?: lookUpFunctionExternal(descriptor)
-        }
-
-        return lookUpFunctionIndirect(call, scope)
+        return lookUpFunctionDirect(call) ?: lookUpFunctionIndirect(call, scope) ?: lookUpFunctionExternal(call, scope.fragment)
    }

-    private fun lookUpFunctionIndirect(call: JsInvocation, scope: InliningScope): LocalInlineFunctionDefinition? {
+    private fun lookUpFunctionIndirect(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
        /** remove ending `()` */
        val callQualifier: JsExpression = if (isCallInvocation(call)) {
            (call.qualifier as JsNameRef).qualifier!!
@@ -121,21 +112,23 @@ class FunctionContext(
            is JsFunction -> functionsByFunctionNodes[qualifier] ?: FunctionWithWrapper(qualifier, null)
            else -> null
        }?.let {
-            LocalInlineFunctionDefinition(it, scope)
+            InlineFunctionDefinition(it, null).also { definition ->
+                if (scope.fragment in newFragments) {
+                    inliner.process(definition, call, scope)
+                }
+            }
        }
    }

-    fun functionTag(call: JsInvocation): String? {
-        return call.descriptor?.let { Namer.getFunctionTag(it, inliner.config) }
-    }
-
-    private val newFragmentSet = inliner.translationResult.newFragments.toIdentitySet()
-
    private val inliningScopeCache = mutableMapOf<JsProgramFragment, ProgramFragmentInliningScope>()

-    fun scopeForFragment(fragment: JsProgramFragment) = inliningScopeCache.computeIfAbsent(fragment) {
-        ProgramFragmentInliningScope(fragment, this, inliner)
-    }
+    private val newFragments = inliner.translationResult.newFragments.toIdentitySet()
+
+    fun scopeForFragment(fragment: JsProgramFragment) = if (fragment in newFragments) {
+        inliningScopeCache.computeIfAbsent(fragment) {
+            ProgramFragmentInliningScope(fragment)
+        }
+    } else null

    private fun loadFragment(fragment: JsProgramFragment) {
        fragmentInfo.computeIfAbsent(fragment) {
@@ -145,36 +138,38 @@ class FunctionContext(
            ).also { (functions, accessors) ->
                (functions.values.asSequence() + accessors.values.asSequence()).forEach { f ->
                    functionsByFunctionNodes[f.function] = f
-                    if (f.wrapperBody != null) {
-                        functionsByWrapperNodes[f.wrapperBody] = f
-                    }
                }
            }
        }
    }

-    private fun fragmentByTag(tag: String): JsProgramFragment? {
+    fun fragmentByTag(tag: String): JsProgramFragment? {
        return inliner.translationResult.inlineFunctionTagMap[tag]?.let { unit ->
            inliner.translationResult.translate(unit).fragment.also { loadFragment(it) }
        }
    }

-    private fun lookUpFunctionDirect(descriptor: CallableDescriptor): InlineFunctionDefinition? =
-        Namer.getFunctionTag(descriptor, inliner.config).let { tag ->
-            fragmentByTag(tag)?.let { fragment ->
-                lookUpStaticFunctionByTag(tag, fragment)?.let {
-                    if (fragment !in newFragmentSet) {
-                        BinaryInlineFunctionDefinition(tag, it, fragment)
-                    } else {
-                        // TODO This is a wrong scope =(
-                        PublicInlineFunctionDefinition(tag, it, fragment, scopeForFragment(fragment))
+    private fun lookUpFunctionDirect(call: JsInvocation): InlineFunctionDefinition? =
+        call.descriptor?.let { descriptor ->
+            Namer.getFunctionTag(descriptor, inliner.config).let { tag ->
+                fragmentByTag(tag)?.let { definitionFragment ->
+                    return lookUpStaticFunctionByTag(tag, definitionFragment)?.let { fn ->
+                        InlineFunctionDefinition(fn, tag).also { definition ->
+                            scopeForFragment(definitionFragment)?.let { definitionScope ->
+                                inliner.process(definition, call, definitionScope)
+                            }
+                        }
                    }
                }
            }
        }

-
-    private fun lookUpFunctionExternal(descriptor: CallableDescriptor): LibraryInlineFunctionDefinition? = functionReader[descriptor]
+    private fun lookUpFunctionExternal(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition? =
+        call.descriptor?.let { descriptor ->
+            functionReader[descriptor, fragment]?.let {
+                InlineFunctionDefinition(it, Namer.getFunctionTag(descriptor, inliner.config))
+            }
+        }

    private fun tryExtractCallableReference(invocation: JsInvocation): FunctionWithWrapper? {
        if (invocation.isCallableReference) {
@@ -16,9 +16,10 @@

 package org.jetbrains.kotlin.js.inline.context

-import org.jetbrains.kotlin.js.backend.ast.JsContext
 import org.jetbrains.kotlin.js.backend.ast.JsStatement

-class InliningContext(val statementContext: JsContext<JsStatement>) {
-    fun newNamingContext() = NamingContext(statementContext)
+class InliningContext(val currentStatement: JsStatement?) {
+    val previousStatements = mutableListOf<JsStatement>()
+
+    fun newNamingContext() = NamingContext(previousStatements)
 }
@@ -21,14 +21,14 @@ import org.jetbrains.kotlin.js.backend.ast.metadata.synthetic
 import org.jetbrains.kotlin.js.inline.util.replaceNames
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils

-class NamingContext(private val statementContext: JsContext<JsStatement>) {
+class NamingContext(private val previousStatements: MutableList<JsStatement>) {
    private val renamings = mutableMapOf<JsName, JsNameRef>()
    private val declarations = mutableListOf<JsVars>()
    private var addedDeclarations = false

    fun applyRenameTo(target: JsNode): JsNode {
        if (!addedDeclarations) {
-            statementContext.addPrevious(declarations)
+            previousStatements.addAll(declarations)
            addedDeclarations = true
        }