Extract LookupStorage from LookupTrackerImpl

jps-plugin/src/org/jetbrains/kotlin/jps/build/KotlinBuilder.kt

@@ -54,7 +54,6 @@ import org.jetbrains.kotlin.config.CompilerRunnerConstants.INTERNAL_ERROR_PREFIX
 import org.jetbrains.kotlin.config.IncrementalCompilation
 import org.jetbrains.kotlin.config.Services
 import org.jetbrains.kotlin.incremental.components.LookupTracker
-import org.jetbrains.kotlin.incremental.components.ScopeKind
 import org.jetbrains.kotlin.jps.JpsKotlinCompilerSettings
 import org.jetbrains.kotlin.jps.incremental.*
 import org.jetbrains.kotlin.load.kotlin.ModuleMapping
@@ -153,8 +152,8 @@ public class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR
         messageCollector.report(INFO, "Kotlin JPS plugin version " + KotlinVersion.VERSION, CompilerMessageLocation.NO_LOCATION)
 
         val project = projectDescriptor.project
-        val (lookupTracker, lookupStorage) = getLookupTrackerAndStorage(dataManager, project)
+        val lookupTracker = getLookupTracker(project)
-        val incrementalCaches = getIncrementalCaches(chunk, context, lookupStorage)
+        val incrementalCaches = getIncrementalCaches(chunk, context)
         val environment = createCompileEnvironment(incrementalCaches, lookupTracker, context)
         if (!environment.success()) {
             environment.reportErrorsTo(messageCollector)
@@ -166,7 +165,6 @@ public class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR
 
         val allCompiledFiles = getAllCompiledFilesContainer(context)
         val filesToCompile = KotlinSourceFileCollector.getDirtySourceFiles(dirtyFilesHolder)
-
         val start = System.nanoTime()
         val outputItemCollector = doCompileModuleChunk(allCompiledFiles, chunk, commonArguments, context, dirtyFilesHolder,
                                                        environment, filesToCompile, incrementalCaches, messageCollector, project)
@@ -177,7 +175,6 @@ public class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR
             return NOTHING_DONE
         }
 
-
         val compilationErrors = Utils.ERRORS_DETECTED_KEY[context, false]
         if (compilationErrors) {
             LOG.info("Compiled with errors")
@@ -199,6 +196,7 @@ public class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR
                 val generatedClasses = generatedFiles.filterIsInstance<GeneratedJvmClass>()
                 val info = updateKotlinIncrementalCache(compilationErrors, incrementalCaches, generatedFiles, chunk)
                 updateJavaMappings(chunk, compilationErrors, context, dirtyFilesHolder, filesToCompile, generatedClasses)
+                updateLookupStorage(chunk, lookupTracker, dataManager, dirtyFilesHolder, filesToCompile)
                 info
             }
         }
@@ -468,6 +466,26 @@ public class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR
         return changesInfo
     }
 
+    private fun updateLookupStorage(
+            chunk: ModuleChunk,
+            lookupTracker: LookupTracker,
+            dataManager: BuildDataManager,
+            dirtyFilesHolder: DirtyFilesHolder<JavaSourceRootDescriptor, ModuleBuildTarget>,
+            filesToCompile: MultiMap<ModuleBuildTarget, File>
+    ) {
+        if (!IncrementalCompilation.isExperimental()) return
+
+        if (lookupTracker !is LookupTrackerImpl) throw AssertionError("Lookup tracker is expected to be LookupTrackerImpl, got ${lookupTracker.javaClass}")
+
+        val lookupStorage = dataManager.getStorage(LOOKUP_TRACKER_TARGET, LOOKUP_TRACKER_STORAGE_PROVIDER)
+
+        filesToCompile.values().forEach { lookupStorage.removeLookupsFrom(it) }
+        val removedFiles = chunk.targets.flatMap { KotlinSourceFileCollector.getRemovedKotlinFiles(dirtyFilesHolder, it) }
+        removedFiles.forEach { lookupStorage.removeLookupsFrom(it) }
+
+        lookupTracker.lookups.entrySet().forEach { lookupStorage.add(it.key, it.value) }
+    }
+
     private fun File.isModuleMappingFile() = extension == ModuleMapping.MAPPING_FILE_EXT && parentFile.name == "META-INF"
 
     // if null is returned, nothing was done
@@ -624,37 +642,23 @@ public class KotlinBuilder : ModuleLevelBuilder(BuilderCategory.SOURCE_PROCESSOR
 private val Iterable<BuildTarget<*>>.moduleTargets: Iterable<ModuleBuildTarget>
     get() = filterIsInstance(javaClass<ModuleBuildTarget>())
 
-private fun getLookupTrackerAndStorage(dataManager: BuildDataManager, project: JpsProject): Pair<LookupTracker, LookupStorage> {
+private fun getLookupTracker(project: JpsProject): LookupTracker {
     var lookupTracker = LookupTracker.DO_NOTHING
-    var lookupStorage = LookupStorage.DO_NOTHING
 
-    if (IncrementalCompilation.isExperimental()) {
+    if ("true".equals(System.getProperty("kotlin.jps.tests"), ignoreCase = true)) {
-        val lookupTrackerImpl = dataManager.getStorage(LOOKUP_TRACKER_TARGET, LOOKUP_TRACKER_STORAGE_PROVIDER)
-        lookupTracker = lookupTrackerImpl
-        lookupStorage = lookupTrackerImpl
-    }
-
-    val inTest = "true".equals(System.getProperty("kotlin.jps.tests"), ignoreCase = true)
-    if (inTest) {
         val testTracker = project.container.getChild(KotlinBuilder.LOOKUP_TRACKER)?.data
 
         if (testTracker != null) {
-            lookupTracker = CopyingLookupTracker(lookupTracker, testTracker)
+            lookupTracker = testTracker
         }
     }
 
-    return lookupTracker to lookupStorage
+    if (IncrementalCompilation.isExperimental()) return LookupTrackerImpl(lookupTracker)
+
+    return lookupTracker
 }
 
-private class CopyingLookupTracker(private val lookupTrackers: Collection<LookupTracker>) : LookupTracker {
+private fun getIncrementalCaches(chunk: ModuleChunk, context: CompileContext): Map<ModuleBuildTarget, IncrementalCacheImpl> {
-    constructor(vararg lookupTrackers: LookupTracker) : this(lookupTrackers.toList())
-
-    override fun record(lookupContainingFile: String, lookupLine: Int?, lookupColumn: Int?, scopeFqName: String, scopeKind: ScopeKind, name: String) {
-        lookupTrackers.forEach { it.record(lookupContainingFile, lookupLine, lookupColumn, scopeFqName, scopeKind, name) }
-    }
-}
-
-private fun getIncrementalCaches(chunk: ModuleChunk, context: CompileContext, lookupStorage: LookupStorage): Map<ModuleBuildTarget, IncrementalCacheImpl> {
     val dataManager = context.projectDescriptor.dataManager
     val targets = chunk.targets
 
@@ -684,8 +688,6 @@ private fun getIncrementalCaches(chunk: ModuleChunk, context: CompileContext, lo
         dependents[target]?.forEach {
             cache.addDependentCache(caches[it]!!)
         }
-
-        cache.setLookupStorage(lookupStorage)
     }
 
     return caches

jps-plugin/src/org/jetbrains/kotlin/jps/incremental/IncrementalCacheImpl.kt

@@ -87,11 +87,6 @@ public class IncrementalCacheImpl(
     private val cacheFormatVersion = CacheFormatVersion(targetDataRoot)
     private val dependents = arrayListOf<IncrementalCacheImpl>()
     private val outputDir = requireNotNull(target.outputDir) { "Target is expected to have output directory: $target" }
-    private var _lookupStorage: LookupStorage = LookupStorage.DO_NOTHING
-
-    public fun setLookupStorage(lookupStorage: LookupStorage) {
-        this._lookupStorage = lookupStorage
-    }
 
     override fun registerInline(fromPath: String, jvmSignature: String, toPath: String) {
         inlinedTo.add(fromPath, jvmSignature, toPath)
@@ -501,7 +496,6 @@ public class IncrementalCacheImpl(
 
         private fun remove(path: String) {
             storage.remove(path)
-            _lookupStorage.removeLookupsFrom(File(path))
         }
     }
 

jps-plugin/src/org/jetbrains/kotlin/jps/incremental/LookupStorage.kt

@@ -16,31 +16,20 @@
 
 package org.jetbrains.kotlin.jps.incremental
 
+import com.intellij.util.containers.MultiMap
 import org.jetbrains.jps.builders.storage.StorageProvider
+import org.jetbrains.kotlin.incremental.components.LocationInfo
 import org.jetbrains.kotlin.incremental.components.LookupTracker
 import org.jetbrains.kotlin.incremental.components.ScopeKind
-import org.jetbrains.kotlin.jps.incremental.storage.BasicMapsOwner
+import org.jetbrains.kotlin.jps.incremental.storage.*
-import org.jetbrains.kotlin.jps.incremental.storage.FileToIdMap
-import org.jetbrains.kotlin.jps.incremental.storage.IdToFileMap
-import org.jetbrains.kotlin.jps.incremental.storage.LookupMap
 import java.io.File
+import java.util.*
 
-object LOOKUP_TRACKER_STORAGE_PROVIDER : StorageProvider<LookupTrackerImpl>() {
+object LOOKUP_TRACKER_STORAGE_PROVIDER : StorageProvider<LookupStorage>() {
-    override fun createStorage(targetDataDir: File): LookupTrackerImpl = LookupTrackerImpl(targetDataDir)
+    override fun createStorage(targetDataDir: File): LookupStorage = LookupStorage(targetDataDir)
 }
 
-interface LookupStorage {
+class LookupStorage(private val targetDataDir: File) : BasicMapsOwner() {
-    fun removeLookupsFrom(file: File)
-
-    companion object {
-        val DO_NOTHING: LookupStorage = object : LookupStorage {
-            override fun removeLookupsFrom(file: File) {}
-        }
-    }
-}
-
-class LookupTrackerImpl(private val targetDataDir: File) : BasicMapsOwner(), LookupTracker, LookupStorage {
-
     companion object {
         private val DELETED_TO_SIZE_TRESHOLD = 0.5
         private val MINIMUM_GARBAGE_COLLECTIBLE_SIZE = 10000
@@ -64,13 +53,14 @@ class LookupTrackerImpl(private val targetDataDir: File) : BasicMapsOwner(), Loo
         }
     }
 
-    override fun record(lookupContainingFile: String, lookupLine: Int?, lookupColumn: Int?, scopeFqName: String, scopeKind: ScopeKind, name: String) {
+    public fun add(lookupSymbol: LookupSymbol, containingFiles: Collection<File>) {
-        val file = File(lookupContainingFile)
+        val key = lookupSymbol.toHashPair()
-        val fileId = fileToId[file] ?: addFile(file)
+        val fileIds = containingFiles.map { addFileIfNeeded(it) }.toHashSet()
-        lookupMap.add(name, scopeFqName, fileId)
+        fileIds.addAll(lookupMap[key] ?: emptySet())
+        lookupMap[key] = fileIds
     }
 
-    override fun removeLookupsFrom(file: File) {
+    public fun removeLookupsFrom(file: File) {
         val id = fileToId[file] ?: return
         idToFile.remove(id)
         fileToId.remove(file)
@@ -91,31 +81,72 @@ class LookupTrackerImpl(private val targetDataDir: File) : BasicMapsOwner(), Loo
     override fun flush(memoryCachesOnly: Boolean) {
         try {
             removeGarbageIfNeeded()
-            countersFile.writeText("$size\n$deletedCount")
+
+            if (size > 0) {
+                if (!countersFile.exists()) {
+                    countersFile.parentFile.mkdirs()
+                    countersFile.createNewFile()
+                }
+
+                countersFile.writeText("$size\n$deletedCount")
+            }
         }
         finally {
             super.flush(memoryCachesOnly)
         }
     }
 
-    private fun addFile(file: File): Int {
+    private fun addFileIfNeeded(file: File): Int {
+        val existing = fileToId[file]
+        if (existing != null) return existing
+
         val id = size++
         fileToId[file] = id
         idToFile[id] = file
         return id
     }
 
-    private fun removeGarbageIfNeeded() {
+    private fun removeGarbageIfNeeded(force: Boolean = false) {
-        if (size <= MINIMUM_GARBAGE_COLLECTIBLE_SIZE && deletedCount.toDouble() / size <= DELETED_TO_SIZE_TRESHOLD) return
+        if (!force && size <= MINIMUM_GARBAGE_COLLECTIBLE_SIZE && deletedCount.toDouble() / size <= DELETED_TO_SIZE_TRESHOLD) return
 
         for (hash in lookupMap.keys) {
             lookupMap[hash] = lookupMap[hash]!!.filter { it in idToFile }.toSet()
         }
 
+        val oldFileToId = fileToId.copyAsMap()
+        val oldIdToNewId = HashMap<Int, Int>(oldFileToId.size)
+        idToFile.clean()
+        fileToId.clean()
         size = 0
         deletedCount = 0
-        idToFile.clean()
 
-        fileToId.files.forEach { addFile(it) }
+        for ((file, oldId) in oldFileToId.entries) {
+            val newId = addFileIfNeeded(file)
+            oldIdToNewId[oldId] = newId
+        }
+
+        for (lookup in lookupMap.keys) {
+            val fileIds = lookupMap[lookup]!!.map { oldIdToNewId[it] }.filterNotNull().toSet()
+
+            if (fileIds.isEmpty()) {
+                lookupMap.remove(lookup)
+            }
+            else {
+                lookupMap[lookup] = fileIds
+            }
+        }
     }
 }
+
+class LookupTrackerImpl(private val delegate: LookupTracker) : LookupTracker {
+    val lookups = MultiMap<LookupSymbol, File>()
+
+    override fun record(locationInfo: LocationInfo, scopeFqName: String, scopeKind: ScopeKind, name: String) {
+        lookups.putValue(LookupSymbol(name, scopeFqName), File(locationInfo.filePath))
+        delegate.record(locationInfo, scopeFqName, scopeKind, name)
+    }
+}
+
+data class LookupSymbol(val name: String, val scope: String)
+
+fun LookupSymbol.toHashPair() = LookupHashPair(name, scope)

jps-plugin/src/org/jetbrains/kotlin/jps/incremental/storage/BasicMap.kt

@@ -21,6 +21,7 @@ import com.intellij.util.io.EnumeratorStringDescriptor
 import com.intellij.util.io.KeyDescriptor
 import org.jetbrains.annotations.TestOnly
 import org.jetbrains.kotlin.utils.Printer
+import org.jetbrains.kotlin.utils.keysToMap
 import java.io.File
 
 internal abstract class BasicMap<K : Comparable<K>, V>(
@@ -42,6 +43,8 @@ internal abstract class BasicMap<K : Comparable<K>, V>(
         storage.close()
     }
 
+    public fun copyAsMap(): Map<K, V> = storage.keys.keysToMap { storage[it]!! }
+
     @TestOnly
     fun dump(): String {
         return with(StringBuilder()) {

jps-plugin/src/org/jetbrains/kotlin/jps/incremental/storage/FileToIdMap.kt

@@ -32,7 +32,4 @@ class FileToIdMap(file: File) : BasicMap<File, Int>(file, FileKeyDescriptor, Int
     public fun remove(file: File) {
         storage.remove(file)
     }
-
-    public val files: Collection<File>
-        get() = storage.keys
 }

jps-plugin/src/org/jetbrains/kotlin/jps/incremental/storage/LookupMap.kt

@@ -33,6 +33,10 @@ class LookupMap(storage: File) : BasicMap<LookupHashPair, Set<Int>>(storage, Loo
         storage.set(key, fileIds)
     }
 
+    public fun remove(key: LookupHashPair) {
+        storage.remove(key)
+    }
+
     public val keys: Collection<LookupHashPair>
         get() = storage.keys
 }