JS: support inline function definitions without wrappers

Apparently function literals are translated without hiding the local declarations
in the wrapFunction call.

Which makes sense. Apart from the fact that it would also make
sense to do the same for private inline functions and function
literals inside other inline functions.

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

@@ -27,13 +27,13 @@ class InlineAstVisitor(
     override fun visit(x: JsInvocation, ctx: JsContext<*>): Boolean {
         // Is it `defineInlineFunction('tag', ...)`?
         InlineMetadata.decompose(x)?.let {
-            jsInliner.process(InlineFunctionDefinition(it.function, it.tag.value), x, scope.fragment)
+            jsInliner.process(InlineFunctionDefinition(it.function, it.tag.value), x, scope.fragment, scope)
             return false
         }
 
         // Is it `wrapFunction(...)`?
         InlineMetadata.tryExtractFunction(x)?.let {
-            jsInliner.process(InlineFunctionDefinition(it, null), x, scope.fragment)
+            jsInliner.process(InlineFunctionDefinition(it, null), x, scope.fragment, scope)
             return false
         }
 
@@ -108,7 +108,7 @@ class InlineAstVisitor(
 
     private fun hasToBeInlined(call: JsInvocation): Boolean {
         val strategy = call.inlineStrategy
-        return if (strategy == null || !strategy.isInline) false else jsInliner.functionDefinitionLoader.hasFunctionDefinition(call, scope.fragment)
+        return if (strategy == null || !strategy.isInline) false else jsInliner.functionDefinitionLoader.hasFunctionDefinition(call, scope)
     }
 
     private fun patchReturnsFromSecondaryConstructor(function: JsFunction) {

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

@@ -33,23 +33,31 @@ class JsInliner(
         }
     }
 
-    fun process(inlineFn: InlineFunctionDefinition, call: JsInvocation?, definitionFragment: JsProgramFragment) {
+    fun process(
+        inlineFn: InlineFunctionDefinition,
+        call: JsInvocation?,
+        definitionFragment: JsProgramFragment,
+        callsiteScope: InliningScope
+    ) {
         // Old fragments are already processed
         if (definitionFragment !in translationResult.newFragments) return
 
         cycleReporter.processInlineFunction(inlineFn.fn, call) {
-            val wrapperBody = inlineFn.fn.wrapperBody
+            val (fn, wrapperBody) = inlineFn.fn
 
-            checkNotNull(wrapperBody) { "All unprocessed inline functions should be generated with wrappers" }
+            if (wrapperBody != null) {
-
+                ImportIntoWrapperInliningScope.process(wrapperBody, definitionFragment) { scope ->
-            ImportIntoWrapperInliningScope.process(wrapperBody, definitionFragment) { scope ->
+                    InlineAstVisitor(this, scope).accept(wrapperBody)
-                InlineAstVisitor(this, scope).accept(wrapperBody)
+                }
+            } else {
+                // e.g. lambda in a non-inline context
+                InlineAstVisitor(this, callsiteScope).accept(fn)
             }
         }
     }
 
     fun inline(scope: InliningScope, call: JsInvocation, currentStatement: JsStatement?): InlineableResult {
-        val definition = functionDefinitionLoader.getFunctionDefinition(call, scope.fragment)
+        val definition = functionDefinitionLoader.getFunctionDefinition(call, scope)
 
         val function = scope.importFunctionDefinition(definition)
 

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

@@ -28,12 +28,12 @@ import java.util.HashMap
 class FunctionDefinitionLoader(
     private val inliner: JsInliner
 ) {
-    fun getFunctionDefinition(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition {
+    fun getFunctionDefinition(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition {
-        return getFunctionDefinitionImpl(call, fragment)!!
+        return getFunctionDefinitionImpl(call, scope)!!
     }
 
-    fun hasFunctionDefinition(call: JsInvocation, fragment: JsProgramFragment): Boolean {
+    fun hasFunctionDefinition(call: JsInvocation, scope: InliningScope): Boolean {
-        return getFunctionDefinitionImpl(call, fragment) != null
+        return getFunctionDefinitionImpl(call, scope) != null
     }
 
     val functionsByFunctionNodes = HashMap<JsFunction, FunctionWithWrapper>()
@@ -68,11 +68,11 @@ class FunctionDefinitionLoader(
      * 5. Qualifier can be JsNameRef with ref to case [3]
      *    in case of local function with closure.
      */
-    private fun getFunctionDefinitionImpl(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition? {
+    private fun getFunctionDefinitionImpl(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
         // Ensure we have the local function information
-        loadFragment(fragment)
+        loadFragment(scope.fragment)
 
-        return lookUpFunctionDirect(call) ?: lookUpFunctionIndirect(call, fragment) ?: lookUpFunctionExternal(call, fragment)
+        return lookUpFunctionDirect(call, scope) ?: lookUpFunctionIndirect(call, scope) ?: lookUpFunctionExternal(call, scope.fragment)
     }
 
     private val functionReader = FunctionReader(inliner.reporter, inliner.config)
@@ -91,14 +91,15 @@ class FunctionDefinitionLoader(
         fragmentInfo[fragment]?.run { accessors[functionTag] }
 
 
-    private fun lookUpFunctionIndirect(call: JsInvocation, fragment: JsProgramFragment): InlineFunctionDefinition? {
+    private fun lookUpFunctionIndirect(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
         /** Try loading via `descriptor` metadata.
          *  In case of private inline properties that's the only way... */
         call.descriptor?.let { descriptor ->
-            fragmentInfo[fragment]?.let { info ->
+            fragmentInfo[scope.fragment]?.let { info ->
                 info.localAccessors[descriptor]?.let { fn ->
                     return InlineFunctionDefinition(fn, null).also { def ->
-                        inliner.process(def, call, fragment)
+                        // Definition fragment is that same as call site fragment
+                        inliner.process(def, call, scope.fragment, scope)
                     }
                 }
             }
@@ -116,10 +117,10 @@ class FunctionDefinitionLoader(
         return when (qualifier) {
             is JsInvocation -> {
                 tryExtractCallableReference(qualifier) ?: getSimpleName(qualifier)?.let { simpleName ->
-                    lookUpStaticFunction(simpleName, fragment)?.let { if (isFunctionCreator(it.function)) it else null }
+                    lookUpStaticFunction(simpleName, scope.fragment)?.let { if (isFunctionCreator(it.function)) it else null }
                 }
             }
-            is JsNameRef -> lookUpStaticFunction(qualifier.name, fragment)
+            is JsNameRef -> lookUpStaticFunction(qualifier.name, scope.fragment)
 
             // Since we could get functionWithWrapper as a simple function directly from staticRef (which always points on implementation)
             // we should check if we have a known wrapper for it
@@ -127,9 +128,8 @@ class FunctionDefinitionLoader(
             else -> null
         }?.let {
             InlineFunctionDefinition(it, null).also { definition ->
-                if (fragment in inliner.translationResult.newFragments) {
+                // Definition fragment is that same as call site fragment
-                    inliner.process(definition, call, fragment)
+                inliner.process(definition, call, scope.fragment, scope)
-                }
             }
         }
     }
@@ -154,7 +154,7 @@ class FunctionDefinitionLoader(
         }
     }
 
-    private fun lookUpFunctionDirect(call: JsInvocation): InlineFunctionDefinition? {
+    private fun lookUpFunctionDirect(call: JsInvocation, callsiteScope: InliningScope): InlineFunctionDefinition? {
         val descriptor = call.descriptor ?: return null
 
         val tag = Namer.getFunctionTag(descriptor, inliner.config)
@@ -166,7 +166,7 @@ class FunctionDefinitionLoader(
         val definition = InlineFunctionDefinition(fn, tag)
 
         // Make sure definition has it's own inline calls inlined.
-        inliner.process(definition, call, definitionFragment)
+        inliner.process(definition, call, definitionFragment, callsiteScope)
 
         return definition
     }