[workspace model] IJPL-156937 Add moduleId to storage metadata resolving process

The storage metadata resolving process has been enhanced by adding the moduleId to improve the resolution accuracy. It's needed to find plugin classloaders with moduleId

GitOrigin-RevId: 4950b8b1a1fe0bf0bb3e5d26725c13f6c64477da

platform/platform-impl/src/com/intellij/workspaceModel/ide/impl/WorkspaceModelCacheSerializer.kt

@@ -117,19 +117,23 @@ class WorkspaceModelCacheSerializer(vfuManager: VirtualFileUrlManager, urlRelati
   }
 
   object PluginAwareEntityTypesResolver : EntityTypesResolver {
-    override fun getPluginId(clazz: Class<*>): String? {
-      return (clazz.classLoader as? PluginAwareClassLoader)?.pluginDescriptor?.pluginId?.idString
+    override fun getPluginIdAndModuleId(clazz: Class<*>): Pair<String?, String?> {
+      val pluginAwareClassLoader = clazz.classLoader as? PluginAwareClassLoader
+      return pluginAwareClassLoader?.pluginId?.idString to pluginAwareClassLoader?.moduleId
     }
 
-    override fun resolveClass(name: String, pluginId: String?): Class<*> {
-      val classLoader = getClassLoader(pluginId) ?:
+    override fun resolveClass(name: String, pluginId: String?, moduleId: String?): Class<*> {
+      val classLoader = getClassLoader(pluginId, moduleId) ?:
         error("Could not resolve class loader for plugin '$pluginId' with type: $name")
 
       if (name.startsWith("[")) return Class.forName(name, true, classLoader)
       return classLoader.loadClass(name)
     }
 
-    override fun getClassLoader(pluginId: String?): ClassLoader? {
+    override fun getClassLoader(pluginId: String?, moduleId: String?): ClassLoader? {
+      if (moduleId != null) {
+        return PluginManagerCore.getPluginSet().findEnabledModule(moduleId)!!.classLoader
+      }
       val id = pluginId?.let { PluginId.getId(it) }
       if (id != null && !PluginManagerCore.isPluginInstalled(id)) {
          return null

platform/workspace/storage/src/com/intellij/platform/workspace/storage/entityStorageSerializers.kt

@@ -15,8 +15,8 @@ public interface EntityStorageSerializer {
 
 @ApiStatus.Internal
 public interface EntityTypesResolver {
-  public fun getPluginId(clazz: Class<*>): String?
-  public fun resolveClass(name: String, pluginId: String?): Class<*>
+  public fun getPluginIdAndModuleId(clazz: Class<*>): Pair<String?, String?>
+  public fun resolveClass(name: String, pluginId: String?, moduleId: String?): Class<*>
 
   /**
    * Method is used to register collections from the kotlin plugin in kryo.
@@ -27,7 +27,7 @@ public interface EntityTypesResolver {
    * In addition, plugin developers are not recommended to use a different version of kotlin-stdlib than the version for intellij-platform.
    */
   @Obsolete
-  public fun getClassLoader(pluginId: String?): ClassLoader?
+  public fun getClassLoader(pluginId: String?, moduleId: String?): ClassLoader?
 }
 
 public sealed class SerializationResult {

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

@@ -20,7 +20,7 @@ internal fun compareWithCurrentEntitiesMetadata(cacheMetadata: CacheMetadata,
 
   cacheMetadata.forEach { (id, cacheTypeMetadata) ->
     val typeFqn = cacheTypeMetadata.metadata.fqName
-    val metadataStorage = metadataResolver.resolveMetadataStorage(typesResolver, id.metadataStorageFqn, id.pluginId)
+    val metadataStorage = metadataResolver.resolveMetadataStorage(typesResolver, id.metadataStorageFqn, id.pluginId, id.moduleId)
     val currentTypeMetadataHash = metadataResolver.resolveTypeMetadataHashOrNull(metadataStorage, typeFqn)
                                   ?: return NotEqual("Failed to load existing metadata for type $typeFqn")
 
@@ -65,11 +65,13 @@ internal class CacheMetadata(
   private val metadataById: LinkedHashMap<Id, List<SerializableTypeMetadata>>
 ): Iterable<Pair<CacheMetadata.Id, SerializableTypeMetadata>> {
 
-  internal data class Id(val pluginId: PluginId, val metadataStorageFqn: String)
+  internal data class Id(val pluginId: PluginId, val moduleId: ModuleId, val metadataStorageFqn: String)
 
   internal class SerializableTypeMetadata(val metadata: StorageTypeMetadata, val metadataHash: MetadataHash)
 
-  fun getMetadataWithPluginId(): Iterable<Pair<PluginId, StorageTypeMetadata>> = this.map { it.first.pluginId to it.second.metadata }
+  fun getMetadataWithPluginId(): Iterable<Triple<PluginId, ModuleId, StorageTypeMetadata>> = this.map {
+    Triple(it.first.pluginId,  it.first.moduleId, it.second.metadata)
+  }
 
   override fun iterator(): Iterator<Pair<Id, SerializableTypeMetadata>> {
     return CacheMetadataIterator(metadataById.iterator())
@@ -80,11 +82,11 @@ internal class CacheMetadata(
     private val metadataById: MutableMap<Id, MutableMap<String, StorageTypeMetadata>> = linkedMapOf()
 
     fun add(clazz: Class<*>) {
-      val pluginId: PluginId = typesResolver.getPluginId(clazz)
-      val metadataStorage = TypeMetadataResolver.getInstance().resolveMetadataStorage(typesResolver, clazz.name, pluginId)
+      val (pluginId: PluginId, moduleId: ModuleId) = typesResolver.getPluginIdAndModuleId(clazz)
+      val metadataStorage = TypeMetadataResolver.getInstance().resolveMetadataStorage(typesResolver, clazz.name, pluginId, moduleId)
       val typeMetadata = TypeMetadataResolver.getInstance().resolveTypeMetadata(metadataStorage, clazz.name)
 
-      val metadataByFqn = metadataById.getOrPut(Id(pluginId, metadataStorage::class.java.name)) { linkedMapOf() }
+      val metadataByFqn = metadataById.getOrPut(Id(pluginId, moduleId, metadataStorage::class.java.name)) { linkedMapOf() }
       typeMetadata.collectTypesByFqn(metadataByFqn, metadataStorage)
     }
 
@@ -93,7 +95,7 @@ internal class CacheMetadata(
     fun toImmutable(typesResolver: EntityTypesResolver): CacheMetadata {
       val map = LinkedHashMap<Id, List<SerializableTypeMetadata>>()
       metadataById.forEach { (id, metadataByFqn) ->
-        val metadataStorage = TypeMetadataResolver.getInstance().resolveMetadataStorage(typesResolver, id.metadataStorageFqn, id.pluginId)
+        val metadataStorage = TypeMetadataResolver.getInstance().resolveMetadataStorage(typesResolver, id.metadataStorageFqn, id.pluginId, id.moduleId)
         val serializableTypesMetadata = metadataByFqn.values.map {
           val typeMetadataHash = TypeMetadataResolver.getInstance().resolveTypeMetadataHash(metadataStorage, it.fqName)
           SerializableTypeMetadata(it, typeMetadataHash)

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

@@ -260,7 +260,7 @@ public class EntityStorageSerializerImpl(
     val nonObjectCount = input.readVarInt(true)
     repeat(nonObjectCount) {
       val objectClass = kryo.readClassAndObject(input) as TypeInfo
-      val resolvedClass = typesResolver.resolveClass(objectClass.fqName, objectClass.pluginId)
+      val resolvedClass = typesResolver.resolveClass(objectClass.fqName, objectClass.pluginId, objectClass.pluginId)
       classCache.putIfAbsent(objectClass, resolvedClass.toClassId())
       kryo.register(resolvedClass)
     }

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

@@ -3,12 +3,14 @@ package com.intellij.platform.workspace.storage.impl.serialization
 
 import com.intellij.platform.workspace.storage.EntityTypesResolver
 
-
 internal typealias PluginId = String?
+internal typealias ModuleId = String?
 
-internal data class TypeInfo(val fqName: String, val pluginId: PluginId)
+internal data class TypeInfo(val fqName: String, val pluginId: PluginId, val moduleId: ModuleId)
 
 internal data class SerializableEntityId(val arrayId: Int, val type: TypeInfo)
 
-internal fun getTypeInfo(clazz: Class<*>, interner: StorageInterner, typesResolver: EntityTypesResolver): TypeInfo =
-  interner.intern(TypeInfo(clazz.name, typesResolver.getPluginId(clazz)))
+internal fun getTypeInfo(clazz: Class<*>, interner: StorageInterner, typesResolver: EntityTypesResolver): TypeInfo {
+  val (pluginId, moduleId) = typesResolver.getPluginIdAndModuleId(clazz)
+  return interner.intern(TypeInfo(clazz.name, pluginId, moduleId))
+}

platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/serialization/registration/EntitiesRegistrar.kt

@@ -5,6 +5,7 @@ import com.esotericsoftware.kryo.kryo5.Kryo
 import com.intellij.platform.workspace.storage.EntityTypesResolver
 import com.intellij.platform.workspace.storage.impl.containers.Object2IntWithDefaultMap
 import com.intellij.platform.workspace.storage.impl.serialization.CacheMetadata
+import com.intellij.platform.workspace.storage.impl.serialization.ModuleId
 import com.intellij.platform.workspace.storage.impl.serialization.PluginId
 import com.intellij.platform.workspace.storage.impl.serialization.TypeInfo
 import com.intellij.platform.workspace.storage.impl.toClassId
@@ -19,16 +20,16 @@ internal fun registerEntitiesClasses(kryo: Kryo, cacheMetadata: CacheMetadata,
 
 private class EntitiesRegistrar(
   private val typesResolver: EntityTypesResolver,
-  private val typesMetadata: Iterable<Pair<PluginId, StorageTypeMetadata>>,
+  private val typesMetadata: Iterable<Triple<PluginId, ModuleId, StorageTypeMetadata>>,
   private val classCache: Object2IntWithDefaultMap<TypeInfo>
 ): StorageRegistrar {
 
   override fun registerClasses(kryo: Kryo) {
-    typesMetadata.forEach { (pluginId, typeMetadata) ->
+    typesMetadata.forEach { (pluginId, moduleId, typeMetadata) ->
       // TODO("Test it. Custom classes can have another plugin id")
       val clazz = when (typeMetadata) {
-        is EntityMetadata -> resolveClass(typeMetadata.entityDataFqName, pluginId)
-        is FinalClassMetadata -> resolveClass(typeMetadata.fqName, pluginId)
+        is EntityMetadata -> resolveClass(typeMetadata.entityDataFqName, pluginId, moduleId)
+        is FinalClassMetadata -> resolveClass(typeMetadata.fqName, pluginId, moduleId)
         else -> null // we don't need to register abstract types
       }
       if (clazz != null) {
@@ -42,9 +43,9 @@ private class EntitiesRegistrar(
     }
   }
 
-  private fun resolveClass(fqName: String, pluginId: PluginId): Class<*> {
-    val resolvedClass = typesResolver.resolveClass(fqName, pluginId)
-    classCache.putIfAbsent(TypeInfo(fqName, pluginId), resolvedClass.toClassId())
+  private fun resolveClass(fqName: String, pluginId: PluginId, moduleId: ModuleId): Class<*> {
+    val resolvedClass = typesResolver.resolveClass(fqName, pluginId, moduleId)
+    classCache.putIfAbsent(TypeInfo(fqName, pluginId, moduleId), resolvedClass.toClassId())
     return resolvedClass
   }
 }

platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/serialization/registration/StorageRegistrar.kt

@@ -166,7 +166,7 @@ internal class StorageClassesRegistrar(
   }
 
   private fun registerKotlinCollectionsInKotlinPlugin(kryo: Kryo) {
-    val classLoader = typesResolver.getClassLoader(kotlinPluginId)
+    val classLoader = typesResolver.getClassLoader(kotlinPluginId, null)
     if (classLoader != null) {
       kotlinCollectionsToRegistrar.forEach {
         val classInKotlinPlugin = classLoader.loadClass(it.name)

platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/serialization/serializer/StorageSerializerUtil.kt

@@ -85,7 +85,7 @@ internal class StorageSerializerUtil(
           val typeInfo = kryo.readObject(input, TypeInfo::class.java)
           val entityFamily = kryo.readObject(input, ImmutableEntityFamily::class.java)
 
-          val classId = classCache.getOrPut(typeInfo) { typesResolver.resolveClass(typeInfo.fqName, typeInfo.pluginId).toClassId() }
+          val classId = classCache.getOrPut(typeInfo) { typesResolver.resolveClass(typeInfo.fqName, typeInfo.pluginId, typeInfo.moduleId).toClassId() }
           mutableBarrel.fillEmptyFamilies(classId)
           mutableBarrel.entityFamilies[classId] = entityFamily
         }
@@ -112,10 +112,10 @@ internal class StorageSerializerUtil(
       val childClazzInfo = kryo.readClassAndObject(input) as TypeInfo
 
       val parentClass = classCache.computeIfAbsent(parentClazzInfo, ToIntFunction {
-        (typesResolver.resolveClass(parentClazzInfo.fqName, parentClazzInfo.pluginId) as Class<WorkspaceEntity>).toClassId()
+        (typesResolver.resolveClass(parentClazzInfo.fqName, parentClazzInfo.pluginId, parentClazzInfo.moduleId) as Class<WorkspaceEntity>).toClassId()
       })
       val childClass = classCache.computeIfAbsent(childClazzInfo, ToIntFunction {
-        (typesResolver.resolveClass(childClazzInfo.fqName, childClazzInfo.pluginId) as Class<WorkspaceEntity>).toClassId()
+        (typesResolver.resolveClass(childClazzInfo.fqName, childClazzInfo.pluginId, childClazzInfo.moduleId) as Class<WorkspaceEntity>).toClassId()
       })
 
       val connectionType = ConnectionId.ConnectionType.valueOf(input.readString())
@@ -136,7 +136,7 @@ internal class StorageSerializerUtil(
       val arrayId = input.readInt()
       val clazzInfo = kryo.readClassAndObject(input) as TypeInfo
       val clazz = classCache.computeIfAbsent(clazzInfo, ToIntFunction {
-        typesResolver.resolveClass(clazzInfo.fqName, clazzInfo.pluginId).toClassId()
+        typesResolver.resolveClass(clazzInfo.fqName, clazzInfo.pluginId, clazzInfo.moduleId).toClassId()
       })
       return createEntityId(arrayId, clazz)
     }
@@ -313,7 +313,7 @@ internal class StorageSerializerUtil(
 
   private fun SerializableEntityId.toEntityId(classCache: Object2IntWithDefaultMap<TypeInfo>): EntityId {
     val classId = classCache.computeIfAbsent(type, ToIntFunction {
-      typesResolver.resolveClass(name = this.type.fqName, pluginId = this.type.pluginId).toClassId()
+      typesResolver.resolveClass(name = this.type.fqName, pluginId = this.type.pluginId, moduleId = this.type.moduleId).toClassId()
     })
     return createEntityId(arrayId = this.arrayId, clazz = classId)
   }

platform/workspace/storage/src/com/intellij/platform/workspace/storage/metadata/resolver/MetadataStorageResolver.kt

@@ -2,6 +2,7 @@
 package com.intellij.platform.workspace.storage.metadata.resolver
 
 import com.intellij.platform.workspace.storage.EntityTypesResolver
+import com.intellij.platform.workspace.storage.impl.serialization.ModuleId
 import com.intellij.platform.workspace.storage.impl.serialization.PluginId
 import com.intellij.platform.workspace.storage.metadata.MetadataStorage
 import com.intellij.platform.workspace.storage.metadata.MetadataStorageBridge
@@ -14,12 +15,12 @@ internal object MetadataStorageResolver {
   private const val GENERATED_METADATA_STORAGE_IMPL_NAME = "MetadataStorageImpl"
 
   internal fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String,
-                                     pluginId: PluginId): MetadataStorage {
+                                     pluginId: PluginId, moduleId: ModuleId): MetadataStorage {
     val packageName = extractPackageName(typeFqn)
     val metadataStorage = metadataStorageCache.getOrPut(pluginId to packageName) {
       val metadataStorageClass: Class<*>
       try {
-        metadataStorageClass = typesResolver.resolveClass(metadataStorageFqn(packageName), pluginId)
+        metadataStorageClass = typesResolver.resolveClass(metadataStorageFqn(packageName), pluginId, moduleId)
       } catch (e : ClassNotFoundException) {
         throw MissingMetadataStorage(metadataStorageFqn(packageName), typeFqn)
       }

platform/workspace/storage/src/com/intellij/platform/workspace/storage/metadata/resolver/TypeMetadataResolver.kt

@@ -2,6 +2,7 @@
 package com.intellij.platform.workspace.storage.metadata.resolver
 
 import com.intellij.platform.workspace.storage.EntityTypesResolver
+import com.intellij.platform.workspace.storage.impl.serialization.ModuleId
 import com.intellij.platform.workspace.storage.impl.serialization.PluginId
 import com.intellij.platform.workspace.storage.metadata.MetadataHash
 import com.intellij.platform.workspace.storage.metadata.MetadataStorage
@@ -23,7 +24,7 @@ internal interface TypeMetadataResolver {
 
   fun resolveTypeMetadataHashOrNull(metadataStorage: MetadataStorage, typeFqn: String): MetadataHash?
 
-  fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String, pluginId: PluginId): MetadataStorage
+  fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String, pluginId: PluginId, moduleId: ModuleId): MetadataStorage
 
   companion object {
     internal fun getInstance(): TypeMetadataResolver = INSTANCE
@@ -51,6 +52,6 @@ internal object TypeMetadataResolverImpl: TypeMetadataResolver {
   override fun resolveTypeMetadataHashOrNull(metadataStorage: MetadataStorage, typeFqn: String): MetadataHash? =
     metadataStorage.getMetadataHashByTypeFqnOrNull(typeFqn)
 
-  override fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String, pluginId: PluginId): MetadataStorage =
-    MetadataStorageResolver.resolveMetadataStorage(typesResolver, typeFqn, pluginId)
+  override fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String, pluginId: PluginId, moduleId: ModuleId): MetadataStorage =
+    MetadataStorageResolver.resolveMetadataStorage(typesResolver, typeFqn, pluginId, moduleId)
 }

platform/workspace/storage/tests/testSrc/com/intellij/platform/workspace/storage/tests/metadata/serialization/MetadataSerializationRoundTripChecker.kt

@@ -20,15 +20,15 @@ internal var deserialization = false
 object MetadataDiffTestResolver: EntityTypesResolver {
   private val pluginPrefix = "PLUGIN___"
 
-  override fun getPluginId(clazz: Class<*>): String = pluginPrefix + clazz.name
+  override fun getPluginIdAndModuleId(clazz: Class<*>): Pair<String, String?> = pluginPrefix + clazz.name to null
 
-  override fun resolveClass(name: String, pluginId: String?): Class<*> {
+  override fun resolveClass(name: String, pluginId: String?, moduleId: String?): Class<*> {
     return resolveClass(
       if (deserialization) name.replaceCacheVersion() else name
     )
   }
 
-  override fun getClassLoader(pluginId: String?): ClassLoader? {
+  override fun getClassLoader(pluginId: String?, moduleId: String?): ClassLoader? {
     return javaClass.classLoader
   }
 

platform/workspace/storage/tests/testSrc/com/intellij/platform/workspace/storage/tests/metadata/serialization/service/TestTypeMetadataResolver.kt

@@ -2,6 +2,7 @@
 package com.intellij.platform.workspace.storage.tests.metadata.serialization.service
 
 import com.intellij.platform.workspace.storage.EntityTypesResolver
+import com.intellij.platform.workspace.storage.impl.serialization.ModuleId
 import com.intellij.platform.workspace.storage.impl.serialization.PluginId
 import com.intellij.platform.workspace.storage.metadata.MetadataHash
 import com.intellij.platform.workspace.storage.metadata.MetadataStorage
@@ -25,8 +26,8 @@ internal class TestTypeMetadataResolver(
   override fun resolveTypeMetadataHashOrNull(metadataStorage: MetadataStorage, typeFqn: String): MetadataHash? =
     typeMetadataResolver.resolveTypeMetadataHashOrNull(metadataStorage, processTypeFqn(typeFqn))
 
-  override fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String, pluginId: PluginId): MetadataStorage =
-    typeMetadataResolver.resolveMetadataStorage(typesResolver, processTypeFqn(typeFqn), pluginId)
+  override fun resolveMetadataStorage(typesResolver: EntityTypesResolver, typeFqn: String, pluginId: PluginId, moduleId: ModuleId): MetadataStorage =
+    typeMetadataResolver.resolveMetadataStorage(typesResolver, processTypeFqn(typeFqn), pluginId, moduleId)
 
   private fun processTypeFqn(typeFqn: String): String = if (deserialization) typeFqn.replaceCacheVersion() else typeFqn
 }