[Workspace model] IDEA-338037 Speed up deserialization in...

[Workspace model] IDEA-338037 Speed up deserialization in `EntityStorageSerializerImpl` using custom reference resolver `CachedReferenceResolver`

`CachedReferenceResolver` uses metadata of types to determine if this type is an enum or not. It also caches calculated results from default `ReferenceResolver`

GitOrigin-RevId: 501cd6bfe6a330f51d9a98696406ca09c412c6d6

platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/serialization/EntityStorageSerializerImpl.kt

@@ -16,13 +16,16 @@ import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.*
 import com.intellij.platform.workspace.storage.impl.containers.BidirectionalLongMultiMap
 import com.intellij.platform.workspace.storage.impl.indices.*
+import com.intellij.platform.workspace.storage.impl.serialization.kryo.CachedReferenceResolver
 import com.intellij.platform.workspace.storage.impl.serialization.registration.StorageClassesRegistrar
 import com.intellij.platform.workspace.storage.impl.serialization.registration.StorageRegistrar
 import com.intellij.platform.workspace.storage.impl.serialization.registration.registerEntitiesClasses
 import com.intellij.platform.workspace.storage.impl.serialization.serializer.StorageSerializerUtil
 import com.intellij.platform.workspace.storage.metadata.diff.CacheMetadataComparator
 import com.intellij.platform.workspace.storage.metadata.diff.ComparisonResult
+import com.intellij.platform.workspace.storage.metadata.model.StorageClassMetadata
 import com.intellij.platform.workspace.storage.metadata.model.StorageTypeMetadata
+import com.intellij.platform.workspace.storage.metadata.utils.collectClassesByFqn
 import com.intellij.platform.workspace.storage.url.UrlRelativizer
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
 import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
@@ -105,10 +108,13 @@ public class EntityStorageSerializerImpl(
       // Save version
       output.writeString(serializerDataFormatVersion)
 
-      val entitiesMetadata = getCacheMetadata(storage, typesResolver)
-      kryo.writeObject(output, entitiesMetadata)// Serialize all Entities, Entity Source and Symbolic id metadata from the storage
+      val cacheMetadata = getCacheMetadata(storage, typesResolver)
 
-      writeAndRegisterClasses(kryo, output, storage, entitiesMetadata, classCache) // Register entities classes
+      setCachedReferenceResolver(kryo, cacheMetadata)
+
+      kryo.writeObject(output, cacheMetadata)// Serialize all Entities, Entity Source and Symbolic id metadata from the storage
+
+      writeAndRegisterClasses(kryo, output, storage, cacheMetadata, classCache) // Register entities classes
 
       // Write entity data and references
       kryo.writeClassAndObject(output, storage.entitiesByType)
@@ -164,6 +170,7 @@ public class EntityStorageSerializerImpl(
 
         time = logAndResetTime(time) { measuredTime -> "Read cache metadata and compare it with the existing metadata: $measuredTime ns" }
 
+        setCachedReferenceResolver(kryo, cacheMetadata)
 
         readAndRegisterClasses(kryo, input, cacheMetadata, classCache)
         time = logAndResetTime(time) { measuredTime -> "Read and register classes: $measuredTime ns" }
@@ -263,6 +270,21 @@ public class EntityStorageSerializerImpl(
     }
   }
 
+  private fun setCachedReferenceResolver(kryo: Kryo, cacheMetadata: CacheMetadata) {
+    val classesByFqn: MutableMap<String, StorageClassMetadata> = hashMapOf()
+
+    val cacheMetadataAsList = cacheMetadata.toList()
+    cacheMetadataAsList.forEach {
+      it.collectClassesByFqn(classesByFqn)
+    }
+
+    val entitiesWithCollectedClasses = arrayListOf<StorageTypeMetadata>()
+    entitiesWithCollectedClasses.addAll(cacheMetadataAsList)
+    entitiesWithCollectedClasses.addAll(classesByFqn.values)
+
+    kryo.referenceResolver = CachedReferenceResolver(kryo.referenceResolver, entitiesWithCollectedClasses)
+  }
+
   internal val Class<*>.typeInfo: TypeInfo
     get() = getTypeInfo(this, interner, typesResolver)
 

platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/serialization/kryo/ReferenceResolver.kt

+// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.platform.workspace.storage.impl.serialization.kryo
+
+import com.esotericsoftware.kryo.kryo5.ReferenceResolver
+import com.intellij.platform.workspace.storage.EntityStorage
+import com.intellij.platform.workspace.storage.metadata.model.EntityMetadata
+import com.intellij.platform.workspace.storage.metadata.model.FinalClassMetadata
+import com.intellij.platform.workspace.storage.metadata.model.StorageTypeMetadata
+
+
+/**
+ * Used to speed up [ReferenceResolver.useReferences] method during deserialization:
+ * * It checks if metadata of that type exists, in this case it uses metadata to determine if this type is an enum or not
+ * * If there is no metadata of this type, then it calculates result in the delegate and caches it
+ */
+internal class CachedReferenceResolver(private val referenceResolver: ReferenceResolver,
+                                       metadata: List<StorageTypeMetadata>) : ReferenceResolver by referenceResolver {
+  private val cachedTypes: MutableMap<String, Boolean> = hashMapOf()
+
+  private val metadataByFqn = metadata.associate {
+    if (it is EntityMetadata) {
+      it.entityDataFqName to it
+    }
+    else {
+      it.fqName to it
+    }
+  }
+
+  override fun useReferences(type: Class<*>): Boolean {
+    val fqn = type.name
+    val typeMetadata = metadataByFqn[fqn]
+    if (typeMetadata != null) {
+      return typeMetadata !is FinalClassMetadata.EnumClassMetadata
+    }
+    return cachedTypes.getOrPut(fqn) { referenceResolver.useReferences(type) }
+  }
+}
\ No newline at end of file