FIR IDE: synchronous creation/registration of sessions

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 {

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)

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(
+        return getOrCreateIdeSession(sessionProvider, psi.project, moduleInfo)
-            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)
-                }
-            }
-        }
     }
 
     operator fun get(psi: KtElement): FirElement?