From c1789a0ecc23abbe6be29fa91021d4d8cdb0e127 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Wed, 4 Dec 2019 17:29:02 +0300
Subject: [PATCH] FIR IDE: synchronous creation/registration of sessions

---
 .../fir/java/FirJavaModuleBasedSession.kt     |  3 +-
 .../idea/fir/FirIdeJavaModuleBasedSession.kt  |  1 -
 .../kotlin/idea/fir/FirResolveState.kt        | 33 +++++++++++++------
 3 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt b/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt
index 05f83575815..5d507eb0690 100644
--- a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt
+++ b/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt
@@ -72,7 +72,6 @@ class FirLibrarySession private constructor(
 
 
     init {
-        sessionProvider.sessionCache[moduleInfo] = this
         val javaSymbolProvider = JavaSymbolProvider(this, sessionProvider.project, scope)
 
         registerComponent(
@@ -98,6 +97,8 @@ class FirLibrarySession private constructor(
             FirCorrespondingSupertypesCache::class,
             FirCorrespondingSupertypesCache(this)
         )
+
+        sessionProvider.sessionCache[moduleInfo] = this
     }
 
     companion object {
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeJavaModuleBasedSession.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeJavaModuleBasedSession.kt
index 43efad31a52..81d72368cbb 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeJavaModuleBasedSession.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeJavaModuleBasedSession.kt
@@ -30,7 +30,6 @@ class FirIdeJavaModuleBasedSession(
 
 
     init {
-        sessionProvider.sessionCache[moduleInfo] = this
         registerComponent(
             FirProvider::class,
             IdeFirProvider(project, scope, RawFirBuilder(this, stubMode = false), this)
diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolveState.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolveState.kt
index 731be72a9f3..62574371c6c 100644
--- a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolveState.kt
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirResolveState.kt
@@ -25,22 +25,35 @@ private fun createLibrarySession(moduleInfo: IdeaModuleInfo, project: Project, p
     return FirLibrarySession.create(moduleInfo, provider, contentScope, project, IDEPackagePartProvider(contentScope))
 }
 
+private fun getOrCreateIdeSession(
+    sessionProvider: FirProjectSessionProvider,
+    project: Project,
+    moduleInfo: ModuleSourceInfo
+): FirSession {
+    sessionProvider.getSession(moduleInfo)?.let { return it }
+    return synchronized(moduleInfo.module) {
+        sessionProvider.getSession(moduleInfo) ?: FirIdeJavaModuleBasedSession(
+            project, moduleInfo, sessionProvider, moduleInfo.contentScope()
+        ).also { moduleBasedSession ->
+            val ideaModuleInfo = moduleInfo.cast<IdeaModuleInfo>()
+            ideaModuleInfo.dependenciesWithoutSelf().forEach {
+                if (it is IdeaModuleInfo && it.isLibraryClasses()) {
+                    // TODO: consider caching / synchronization here
+                    createLibrarySession(it, project, sessionProvider)
+                }
+            }
+            sessionProvider.sessionCache[moduleInfo] = moduleBasedSession
+        }
+    }
+}
+
 interface FirResolveState {
     val sessionProvider: FirSessionProvider
 
     fun getSession(psi: KtElement): FirSession {
         val sessionProvider = sessionProvider as FirProjectSessionProvider
         val moduleInfo = psi.getModuleInfo() as ModuleSourceInfo
-        return sessionProvider.getSession(moduleInfo) ?: FirIdeJavaModuleBasedSession(
-            psi.project, moduleInfo, sessionProvider, moduleInfo.contentScope()
-        ).also {
-            val ideaModuleInfo = moduleInfo.cast<IdeaModuleInfo>()
-            ideaModuleInfo.dependenciesWithoutSelf().forEach {
-                if (it is IdeaModuleInfo && it.isLibraryClasses()) {
-                    createLibrarySession(it, psi.project, sessionProvider)
-                }
-            }
-        }
+        return getOrCreateIdeSession(sessionProvider, psi.project, moduleInfo)
     }
 
     operator fun get(psi: KtElement): FirElement?