From d35a26a75bcff14af877ebe2aa1d980ce298d5b8 Mon Sep 17 00:00:00 2001
From: Mikhail Glukhikh <mikhail.glukhikh@jetbrains.com>
Date: Thu, 1 Aug 2019 12:25:25 +0300
Subject: [PATCH] Optimization: introduce FirIdeResolveStateService to cache
 FirResolveState

---
 .../idea/fir/FirIdeResolveStateService.kt     | 64 +++++++++++++++++++
 .../kotlin/idea/fir/FirResolveState.kt        | 11 +---
 idea/resources/META-INF/plugin-common.xml     |  3 +
 .../kotlin/idea/fir/FirResolveStateTest.kt    | 18 ++++++
 4 files changed, 87 insertions(+), 9 deletions(-)
 create mode 100644 idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeResolveStateService.kt
 create mode 100644 idea/tests/org/jetbrains/kotlin/idea/fir/FirResolveStateTest.kt

diff --git a/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeResolveStateService.kt b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeResolveStateService.kt
new file mode 100644
index 00000000000..32df81f5f16
--- /dev/null
+++ b/idea/idea-analysis/src/org/jetbrains/kotlin/idea/fir/FirIdeResolveStateService.kt
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir
+
+import com.intellij.openapi.components.ServiceManager
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.ModificationTracker
+import org.jetbrains.kotlin.analyzer.KotlinModificationTrackerService
+import org.jetbrains.kotlin.analyzer.TrackableModuleInfo
+import org.jetbrains.kotlin.fir.java.FirProjectSessionProvider
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+
+interface FirIdeResolveStateService {
+    companion object {
+        fun getInstance(project: Project): FirIdeResolveStateService =
+            ServiceManager.getService(project, FirIdeResolveStateService::class.java)!!
+    }
+
+    val fallbackModificationTracker: ModificationTracker?
+
+    fun getResolveState(moduleInfo: IdeaModuleInfo): FirResolveState
+}
+
+private class FirModuleData(val state: FirResolveState, val modificationTracker: ModificationTracker?) {
+    val modificationCount: Long = modificationTracker?.modificationCount ?: Long.MIN_VALUE
+
+    fun isOutOfDate(): Boolean {
+        val currentModCount = modificationTracker?.modificationCount
+        return currentModCount != null && currentModCount > modificationCount
+    }
+}
+
+class FirIdeResolveStateServiceImpl(val project: Project) : FirIdeResolveStateService {
+    private val stateCache = mutableMapOf<IdeaModuleInfo, FirModuleData>()
+
+    private fun createResolveState(): FirResolveState {
+        val provider = FirProjectSessionProvider(project)
+        return FirResolveStateImpl(provider)
+    }
+
+    private fun createModuleData(moduleInfo: IdeaModuleInfo): FirModuleData {
+        val state = createResolveState()
+        val modificationTracker = (moduleInfo as? TrackableModuleInfo)?.createModificationTracker() ?: fallbackModificationTracker
+        return FirModuleData(state, modificationTracker)
+    }
+
+    // TODO: multi thread protection
+    override fun getResolveState(moduleInfo: IdeaModuleInfo): FirResolveState {
+        var moduleData = stateCache.getOrPut(moduleInfo) {
+            createModuleData(moduleInfo)
+        }
+        if (moduleData.isOutOfDate()) {
+            moduleData = createModuleData(moduleInfo)
+            stateCache[moduleInfo] = moduleData
+        }
+        return moduleData.state
+    }
+
+    override val fallbackModificationTracker: ModificationTracker? =
+        KotlinModificationTrackerService.getInstance(project).outOfBlockModificationTracker
+}
\ No newline at end of file
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 8d836725284..731be72a9f3 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
@@ -58,12 +58,5 @@ class FirResolveStateImpl(override val sessionProvider: FirSessionProvider) : Fi
     }
 }
 
-// TODO: caching
-object FirIdeResolveFactory {
-    fun initiate(psi: KtElement): FirResolveState {
-        val provider = FirProjectSessionProvider(psi.project)
-        return FirResolveStateImpl(provider)
-    }
-}
-
-fun KtElement.firResolveState(): FirResolveState = FirIdeResolveFactory.initiate(this)
\ No newline at end of file
+fun KtElement.firResolveState(): FirResolveState =
+    FirIdeResolveStateService.getInstance(project).getResolveState(getModuleInfo())
\ No newline at end of file
diff --git a/idea/resources/META-INF/plugin-common.xml b/idea/resources/META-INF/plugin-common.xml
index 47863477335..956a5cdcb34 100644
--- a/idea/resources/META-INF/plugin-common.xml
+++ b/idea/resources/META-INF/plugin-common.xml
@@ -274,6 +274,9 @@
     <projectService serviceInterface="org.jetbrains.kotlin.caches.resolve.KotlinCacheService"
                     serviceImplementation="org.jetbrains.kotlin.idea.caches.resolve.KotlinCacheServiceImpl"/>
 
+    <projectService serviceInterface="org.jetbrains.kotlin.idea.fir.FirIdeResolveStateService"
+                    serviceImplementation="org.jetbrains.kotlin.idea.fir.FirIdeResolveStateServiceImpl"/>
+
     <projectService serviceInterface="org.jetbrains.kotlin.load.kotlin.MetadataFinderFactory"
                     serviceImplementation="org.jetbrains.kotlin.idea.vfilefinder.IDEVirtualFileFinderFactory"/>
 
diff --git a/idea/tests/org/jetbrains/kotlin/idea/fir/FirResolveStateTest.kt b/idea/tests/org/jetbrains/kotlin/idea/fir/FirResolveStateTest.kt
new file mode 100644
index 00000000000..669c1504aef
--- /dev/null
+++ b/idea/tests/org/jetbrains/kotlin/idea/fir/FirResolveStateTest.kt
@@ -0,0 +1,18 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.fir
+
+import org.jetbrains.kotlin.idea.AbstractResolveElementCacheTest
+
+class FirResolveStateTest : AbstractResolveElementCacheTest() {
+    fun testSimpleStateCaching() {
+        doTest {
+            val firstStatement = statements[0]
+            val secondStatement = statements[1]
+            assertSame(firstStatement.firResolveState(), secondStatement.firResolveState())
+        }
+    }
+}
\ No newline at end of file