From b7bea3242e3cea1a7cf0736da552ec39b3b058c4 Mon Sep 17 00:00:00 2001
From: Anton Bannykh <anton.bannykh@jetbrains.com>
Date: Sun, 10 Feb 2019 20:52:17 +0300
Subject: [PATCH] JS: fix fragment info loading

The very first fragment didn't got processed, thus inline cycle
reporter was not failing tests
---
 .../context/FunctionDefinitionLoader.kt       | 39 ++++++++++---------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/js/js.inliner/src/org/jetbrains/kotlin/js/inline/context/FunctionDefinitionLoader.kt b/js/js.inliner/src/org/jetbrains/kotlin/js/inline/context/FunctionDefinitionLoader.kt
index a79ec55ed59..196f7962c73 100644
--- a/js/js.inliner/src/org/jetbrains/kotlin/js/inline/context/FunctionDefinitionLoader.kt
+++ b/js/js.inliner/src/org/jetbrains/kotlin/js/inline/context/FunctionDefinitionLoader.kt
@@ -70,7 +70,7 @@ class FunctionDefinitionLoader(
      */
     private fun getFunctionDefinitionImpl(call: JsInvocation, scope: InliningScope): InlineFunctionDefinition? {
         // Ensure we have the local function information
-        loadFragment(scope.fragment)
+        assert(scope.fragment in fragmentInfo)
 
         return lookUpFunctionDirect(call, scope) ?: lookUpFunctionIndirect(call, scope) ?: lookUpFunctionExternal(call, scope.fragment)
     }
@@ -80,9 +80,22 @@ class FunctionDefinitionLoader(
     private data class FragmentInfo(
         val functions: Map<JsName, FunctionWithWrapper>,
         val accessors: Map<String, FunctionWithWrapper>,
-        val localAccessors: Map<CallableDescriptor, FunctionWithWrapper>)
+        val localAccessors: Map<CallableDescriptor, FunctionWithWrapper>
+    )
 
-    private val fragmentInfo = mutableMapOf<JsProgramFragment, FragmentInfo>()
+    private fun JsProgramFragment.loadInfo(): FragmentInfo {
+        return FragmentInfo(
+            collectNamedFunctionsAndWrappers(listOf(this)),
+            collectAccessors(listOf(this)),
+            collectLocalFunctions(listOf(this))
+        ).also { (functions, accessors) ->
+            (functions.values.asSequence() + accessors.values.asSequence()).forEach { f ->
+                functionsByFunctionNodes[f.function] = f
+            }
+        }
+    }
+
+    private val fragmentInfo = inliner.translationResult.newFragments.associateTo(mutableMapOf()) { it to it.loadInfo() }
 
     private fun lookUpStaticFunction(functionName: JsName?, fragment: JsProgramFragment): FunctionWithWrapper? =
         fragmentInfo[fragment]?.run { functions[functionName] }
@@ -134,23 +147,13 @@ class FunctionDefinitionLoader(
         }
     }
 
-    private fun loadFragment(fragment: JsProgramFragment) {
-        fragmentInfo.computeIfAbsent(fragment) {
-            FragmentInfo(
-                collectNamedFunctionsAndWrappers(listOf(fragment)),
-                collectAccessors(listOf(fragment)),
-                collectLocalFunctions(listOf(fragment))
-            ).also { (functions, accessors) ->
-                (functions.values.asSequence() + accessors.values.asSequence()).forEach { f ->
-                    functionsByFunctionNodes[f.function] = f
-                }
-            }
-        }
-    }
-
     private fun fragmentByTag(tag: String): JsProgramFragment? {
         return inliner.translationResult.inlineFunctionTagMap[tag]?.let { unit ->
-            inliner.translationResult.getTranslationResult(unit).fragment.also { loadFragment(it) }
+            inliner.translationResult.getTranslationResult(unit).fragment.also { fragment ->
+                fragmentInfo.computeIfAbsent(fragment) {
+                    fragment.loadInfo()
+                }
+            }
         }
     }