[IC] Segregate LazyStorage into LazyStorage and AppendableLazyStorage

#KT-56052 In Progress

build-common/src/org/jetbrains/kotlin/incremental/storage/BasicMap.kt

@@ -24,22 +24,11 @@ import org.jetbrains.kotlin.incremental.IncrementalCompilationContext
 import org.jetbrains.kotlin.utils.Printer
 import java.io.File
 
-abstract class BasicMap<K : Comparable<K>, V>(
+abstract class BasicMap<K : Comparable<K>, V, StorageType : LazyStorage<K, V>>(
     internal val storageFile: File,
-    keyDescriptor: KeyDescriptor<K>,
+    protected val storage: StorageType,
-    valueExternalizer: DataExternalizer<V>,
     protected val icContext: IncrementalCompilationContext,
 ) {
-    protected val storage: LazyStorage<K, V> = CachingLazyStorage(storageFile, keyDescriptor, valueExternalizer).let {
-        if (icContext.keepIncrementalCompilationCachesInMemory) {
-            DefaultInMemoryStorageWrapper(it).also { wrapper ->
-                icContext.transaction.registerInMemoryStorageWrapper(wrapper)
-            }
-        } else {
-            it
-        }
-    }
-
     protected val pathConverter
         get() = icContext.pathConverter
 
@@ -86,12 +75,34 @@ abstract class BasicMap<K : Comparable<K>, V>(
     protected abstract fun dumpValue(value: V): String
 }
 
+abstract class NonAppendableBasicMap<K : Comparable<K>, V>(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<K>,
+    valueExternalizer: DataExternalizer<V>,
+    icContext: IncrementalCompilationContext,
+) : BasicMap<K, V, LazyStorage<K, V>>(
+    storageFile,
+    createLazyStorage(storageFile, keyDescriptor, valueExternalizer, icContext),
+    icContext
+)
+
+abstract class AppendableBasicMap<K : Comparable<K>, V>(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<K>,
+    valueExternalizer: AppendableDataExternalizer<V>,
+    icContext: IncrementalCompilationContext,
+) : BasicMap<K, V, AppendableLazyStorage<K, V>>(
+    storageFile,
+    createLazyStorage(storageFile, keyDescriptor, valueExternalizer, icContext),
+    icContext
+)
+
 abstract class BasicStringMap<V>(
     storageFile: File,
     keyDescriptor: KeyDescriptor<String>,
     valueExternalizer: DataExternalizer<V>,
     icContext: IncrementalCompilationContext,
-) : BasicMap<String, V>(storageFile, keyDescriptor, valueExternalizer, icContext) {
+) : NonAppendableBasicMap<String, V>(storageFile, keyDescriptor, valueExternalizer, icContext) {
     constructor(
         storageFile: File,
         valueExternalizer: DataExternalizer<V>,
@@ -100,3 +111,18 @@ abstract class BasicStringMap<V>(
 
     override fun dumpKey(key: String): String = key
 }
+
+abstract class AppendableBasicStringMap<V>(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<String>,
+    valueExternalizer: AppendableDataExternalizer<V>,
+    icContext: IncrementalCompilationContext,
+) : AppendableBasicMap<String, V>(storageFile, keyDescriptor, valueExternalizer, icContext) {
+    constructor(
+        storageFile: File,
+        valueExternalizer: AppendableDataExternalizer<V>,
+        icContext: IncrementalCompilationContext,
+    ) : this(storageFile, EnumeratorStringDescriptor.INSTANCE, valueExternalizer, icContext)
+
+    override fun dumpKey(key: String): String = key
+}

build-common/src/org/jetbrains/kotlin/incremental/storage/BasicMapsOwner.kt

@@ -21,7 +21,7 @@ import java.io.File
 import java.io.IOException
 
 open class BasicMapsOwner(val cachesDir: File) {
-    private val maps = arrayListOf<BasicMap<*, *>>()
+    private val maps = arrayListOf<BasicMap<*, *, *>>()
 
     companion object {
         val CACHE_EXTENSION = "tab"
@@ -31,17 +31,17 @@ open class BasicMapsOwner(val cachesDir: File) {
         get() = File(cachesDir, this + "." + CACHE_EXTENSION)
 
     @Synchronized
-    protected fun <K, V, M : BasicMap<K, V>> registerMap(map: M): M {
+    protected fun <K, V, S, M : BasicMap<K, V, S>> registerMap(map: M): M {
         maps.add(map)
         return map
     }
 
     open fun clean() {
-        forEachMapSafe("clean", BasicMap<*, *>::clean)
+        forEachMapSafe("clean", BasicMap<*, *, *>::clean)
     }
 
     open fun close() {
-        forEachMapSafe("close", BasicMap<*, *>::close)
+        forEachMapSafe("close", BasicMap<*, *, *>::close)
     }
 
     open fun flush(memoryCachesOnly: Boolean) {
@@ -49,7 +49,7 @@ open class BasicMapsOwner(val cachesDir: File) {
     }
 
     @Synchronized
-    private fun forEachMapSafe(actionName: String, action: (BasicMap<*, *>) -> Unit) {
+    private fun forEachMapSafe(actionName: String, action: (BasicMap<*, *, *>) -> Unit) {
         val actionExceptions = LinkedHashMap<String, Exception>()
         maps.forEach {
             try {

build-common/src/org/jetbrains/kotlin/incremental/storage/CachingLazyStorage.kt

@@ -20,6 +20,7 @@ import com.intellij.util.io.DataExternalizer
 import com.intellij.util.io.IOUtil
 import com.intellij.util.io.KeyDescriptor
 import com.intellij.util.io.PersistentHashMap
+import org.jetbrains.kotlin.incremental.IncrementalCompilationContext
 import java.io.File
 import java.io.IOException
 
@@ -31,7 +32,7 @@ class CachingLazyStorage<K, V>(
     private val storageFile: File,
     private val keyDescriptor: KeyDescriptor<K>,
     private val valueExternalizer: DataExternalizer<V>
-) : LazyStorage<K, V> {
+) : AppendableLazyStorage<K, V> {
     private var storage: PersistentHashMap<K, V>? = null
     private var isStorageFileExist = true
 
@@ -80,6 +81,9 @@ class CachingLazyStorage<K, V>(
 
     @Synchronized
     override fun append(key: K, value: V) {
+        check(valueExternalizer is AppendableDataExternalizer<*>) {
+            "`valueExternalizer` should implement the `AppendableDataExternalizer` interface to be able to call `append`"
+        }
         getStorageOrCreateNew().appendData(key, { valueExternalizer.save(it, value) })
     }
 
@@ -119,3 +123,32 @@ class CachingLazyStorage<K, V>(
 
     private fun createMap(): PersistentHashMap<K, V> = PersistentHashMap(storageFile, keyDescriptor, valueExternalizer)
 }
+
+private fun <K, V> createLazyStorageImpl(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<K>,
+    valueExternalizer: DataExternalizer<V>,
+    icContext: IncrementalCompilationContext,
+) = CachingLazyStorage(storageFile, keyDescriptor, valueExternalizer).let {
+    if (icContext.keepIncrementalCompilationCachesInMemory) {
+        DefaultInMemoryStorageWrapper(it, valueExternalizer).also { wrapper ->
+            icContext.transaction.registerInMemoryStorageWrapper(wrapper)
+        }
+    } else {
+        it
+    }
+}
+
+fun <K, V> createLazyStorage(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<K>,
+    valueExternalizer: DataExternalizer<V>,
+    icContext: IncrementalCompilationContext,
+): LazyStorage<K, V> = createLazyStorageImpl(storageFile, keyDescriptor, valueExternalizer, icContext)
+
+fun <K, V> createLazyStorage(
+    storageFile: File,
+    keyDescriptor: KeyDescriptor<K>,
+    valueExternalizer: AppendableDataExternalizer<V>,
+    icContext: IncrementalCompilationContext,
+): AppendableLazyStorage<K, V> = createLazyStorageImpl(storageFile, keyDescriptor, valueExternalizer, icContext)

build-common/src/org/jetbrains/kotlin/incremental/storage/ClassOneToManyMap.kt

@@ -24,7 +24,7 @@ import java.io.File
 internal open class ClassOneToManyMap(
     storageFile: File,
     icContext: IncrementalCompilationContext,
-) : BasicStringMap<Collection<String>>(storageFile, StringCollectionExternalizer, icContext) {
+) : AppendableBasicStringMap<Collection<String>>(storageFile, StringCollectionExternalizer, icContext) {
     override fun dumpValue(value: Collection<String>): String = value.dumpCollection()
 
     @Synchronized

build-common/src/org/jetbrains/kotlin/incremental/storage/IdToFileMap.kt

@@ -24,7 +24,7 @@ import java.io.File
 internal class IdToFileMap(
     file: File,
     icContext: IncrementalCompilationContext,
-) : BasicMap<Int, String>(file, ExternalIntegerKeyDescriptor(), EnumeratorStringDescriptor.INSTANCE, icContext) {
+) : NonAppendableBasicMap<Int, String>(file, ExternalIntegerKeyDescriptor(), EnumeratorStringDescriptor.INSTANCE, icContext) {
     override fun dumpKey(key: Int): String = key.toString()
 
     override fun dumpValue(value: String): String = value

build-common/src/org/jetbrains/kotlin/incremental/storage/InMemoryStorageWrapper.kt

@@ -5,10 +5,11 @@
 
 package org.jetbrains.kotlin.incremental.storage
 
+import com.intellij.util.io.DataExternalizer
 import java.util.*
 import kotlin.collections.LinkedHashMap
 
-interface InMemoryStorageWrapper<K, V> : LazyStorage<K, V> {
+interface InMemoryStorageWrapper<K, V> : AppendableLazyStorage<K, V> {
     fun resetInMemoryChanges()
 }
 
@@ -17,7 +18,12 @@ interface InMemoryStorageWrapper<K, V> : LazyStorage<K, V> {
  * Flushes all the changes to the [origin] on [flush] invocation.
  * [resetInMemoryChanges] should be called to reset in-memory changes of this wrapper.
  */
-class DefaultInMemoryStorageWrapper<K, V>(private val origin: LazyStorage<K, V>) : InMemoryStorageWrapper<K, V> {
+class DefaultInMemoryStorageWrapper<K, V>(
+    private val origin: CachingLazyStorage<K, V>,
+    private val valueExternalizer: DataExternalizer<V>
+) :
+    InMemoryStorageWrapper<K, V> {
+    // These state properties keep the current diff that will be applied to the [origin] on flush if [resetInMemoryChanges] is not called
     private val inMemoryStorage = LinkedHashMap<K, ValueWrapper<V>>()
     private val removedKeys = hashSetOf<K>()
     private var isCleanRequested = false
@@ -65,25 +71,21 @@ class DefaultInMemoryStorageWrapper<K, V>(private val origin: LazyStorage<K, V>)
     }
 
     override fun append(key: K, value: V) {
-        @Suppress("UNCHECKED_CAST")
+        check(valueExternalizer is AppendableDataExternalizer<V>) {
-        when (value) {
+            "`valueExternalizer` should implement the `AppendableDataExternalizer` interface to be able to call `append`"
-            is Collection<*> -> {
+        }
-                when {
+        when {
-                    key in inMemoryStorage -> {
+            key in inMemoryStorage -> {
-                        // the key's value was set in-memory, so we keep changing it without additional marking as append
+                val currentEntry = inMemoryStorage.getValue(key)
-                        (inMemoryStorage.getValue(key).value as MutableCollection<Any?>).addAll(value)
+                // should we avoid new allocations by performing appends in-place?
-                    }
+                inMemoryStorage[key] = ValueWrapper(valueExternalizer.append(currentEntry.value, value), isAppend = currentEntry.isAppend)
-                    key !in removedKeys && key in origin -> {
+            }
-                        val collection = copyCollection(origin[key] as Collection<*>)
+            key !in removedKeys && key in origin -> {
-                        collection.addAll(value)
+                inMemoryStorage[key] = ValueWrapper(value, isAppend = true)
-                        inMemoryStorage[key] = ValueWrapper(collection as V, isAppend = true)
+            }
-                    }
+            else -> {
-                    else -> {
+                set(key, value)
-                        set(key, value)
-                    }
-                }
             }
-            else -> error("The value is expected to implement the Collection interface")
         }
     }
 
@@ -93,23 +95,26 @@ class DefaultInMemoryStorageWrapper<K, V>(private val origin: LazyStorage<K, V>)
     }
 
     override fun set(key: K, value: V) {
-        @Suppress("UNCHECKED_CAST")
+        inMemoryStorage[key] = ValueWrapper(value)
-        inMemoryStorage[key] = ValueWrapper(if (value is Collection<*>) copyCollection(value) as V else value)
     }
 
-    private fun copyCollection(collection: Collection<*>): MutableCollection<Any?> {
+    override fun get(key: K): V? {
-        val newCollection = when (collection) {
+        return when (key) {
-            is Set<*> -> TreeSet<Any?>()
+            in inMemoryStorage -> {
-            else -> ArrayList<Any?>(collection.size)
+                val entry = inMemoryStorage.getValue(key)
+                val originValue = origin[key]
+                if (entry.isAppend && originValue != null) {
+                    check(valueExternalizer is AppendableDataExternalizer<V>) {
+                        "`valueExternalizer` should implement the `AppendableDataExternalizer` interface to be able to handle `append`"
+                    }
+                    valueExternalizer.append(originValue, entry.value)
+                } else {
+                    entry.value
+                }
+            }
+            !in removedKeys -> origin[key]
+            else -> null
         }
-        newCollection.addAll(collection)
-        return newCollection
-    }
-
-    override fun get(key: K): V? = when (key) {
-        in inMemoryStorage -> inMemoryStorage.getValue(key).value
-        !in removedKeys -> origin[key]
-        else -> null
     }
 
     override fun contains(key: K) = key in inMemoryStorage || (key !in removedKeys && key in origin)

build-common/src/org/jetbrains/kotlin/incremental/storage/LazyStorage.kt

@@ -22,8 +22,11 @@ interface LazyStorage<K, V> {
     operator fun get(key: K): V?
     operator fun set(key: K, value: V)
     fun remove(key: K)
-    fun append(key: K, value: V)
     fun clean()
     fun flush(memoryCachesOnly: Boolean)
     fun close()
+}
+
+interface AppendableLazyStorage<K, V> : LazyStorage<K, V> {
+    fun append(key: K, value: V)
 }

build-common/src/org/jetbrains/kotlin/incremental/storage/LookupMap.kt

@@ -23,7 +23,7 @@ class LookupMap(
     storage: File,
     icContext: IncrementalCompilationContext,
 ) :
-    BasicMap<LookupSymbolKey, Collection<Int>>(
+    AppendableBasicMap<LookupSymbolKey, Collection<Int>>(
         storage,
         LookupSymbolKeyDescriptor(icContext.storeFullFqNamesInLookupCache),
         IntCollectionExternalizer,

build-common/src/org/jetbrains/kotlin/incremental/storage/SourceToJsOutputMap.kt

@@ -12,7 +12,7 @@ import java.io.File
 class SourceToJsOutputMap(
     storageFile: File,
     icContext: IncrementalCompilationContext,
-) : BasicStringMap<Collection<String>>(storageFile, StringCollectionExternalizer, icContext) {
+) : AppendableBasicStringMap<Collection<String>>(storageFile, StringCollectionExternalizer, icContext) {
     override fun dumpValue(value: Collection<String>): String = value.dumpCollection()
 
     @Synchronized

build-common/src/org/jetbrains/kotlin/incremental/storage/SourceToOutputMaps.kt

@@ -36,7 +36,7 @@ internal abstract class AbstractSourceToOutputMap<Name>(
     private val nameTransformer: NameTransformer<Name>,
     storageFile: File,
     icContext: IncrementalCompilationContext,
-) : BasicStringMap<Collection<String>>(storageFile, PathStringDescriptor, StringCollectionExternalizer, icContext) {
+) : AppendableBasicStringMap<Collection<String>>(storageFile, PathStringDescriptor, StringCollectionExternalizer, icContext) {
     fun clearOutputsForSource(sourceFile: File) {
         remove(pathConverter.toPath(sourceFile))
     }

build-common/src/org/jetbrains/kotlin/incremental/storage/externalizers.kt

@@ -31,6 +31,18 @@ import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.resolve.jvm.JvmClassName
 import java.io.*
 
+/**
+ * Externalizer that works correctly when [PersistentHashMap.appendData] is called
+ *
+ * Besides the [append] method, it should support incremental [save] and [read]. E.g. if [save] was called multiple times, [read] should be able to collect them together
+ */
+interface AppendableDataExternalizer<T> : DataExternalizer<T> {
+    /**
+     * Combines two non-serialized values
+     */
+    fun append(currentValue: T, appendData: T): T
+}
+
 class LookupSymbolKeyDescriptor(
     /** If `true`, original values are saved; if `false`, only hashes are saved. */
     private val storeFullFqNames: Boolean = false
@@ -276,7 +288,7 @@ class DelegateDataExternalizer<T>(
 open class CollectionExternalizer<T>(
     private val elementExternalizer: DataExternalizer<T>,
     private val newCollection: () -> MutableCollection<T>
-) : DataExternalizer<Collection<T>> {
+) : AppendableDataExternalizer<Collection<T>> {
     override fun read(input: DataInput): Collection<T> {
         val result = newCollection()
         val stream = input as DataInputStream
@@ -291,6 +303,8 @@ open class CollectionExternalizer<T>(
     override fun save(output: DataOutput, value: Collection<T>) {
         value.forEach { elementExternalizer.save(output, it) }
     }
+
+    override fun append(currentValue: Collection<T>, appendData: Collection<T>) = currentValue + appendData
 }
 
 object StringCollectionExternalizer : CollectionExternalizer<String>(EnumeratorStringDescriptor(), { HashSet() })