[workspace model] Not expose PId outside of the store via external index

GitOrigin-RevId: a568520d76ab4126290c02c6b939e5264a88e880

platform/workspaceModel-core-tests/test/com/intellij/workspace/api/pstorage/ExternalEntityIndexTest.kt

@@ -16,19 +16,19 @@ class ExternalEntityIndexTest {
 
     val index = builder.getOrCreateExternalIndex<Int>(INDEX_ID)
     val entity = builder.addPSourceEntity("hello", PSampleEntitySource("source"))
-    index.index(entity.id, 1)
-    index.index(entity.id, 2)
-    assertEquals(2, index.getDataById(entity.id))
-    index.remove(entity.id)
-    assertNull(index.getDataById(entity.id))
-    index.update(entity.id, 3)
-    assertEquals(3, index.getDataById(entity.id))
+    index.index(entity, 1)
+    index.index(entity, 2)
+    assertEquals(2, index.getDataById(entity))
+    index.remove(entity)
+    assertNull(index.getDataById(entity))
+    index.update(entity, 3)
+    assertEquals(3, index.getDataById(entity))
 
     val storage = builder.toStorage()
     val newIndex = storage.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(newIndex)
-    assertEquals(3, newIndex!!.getDataById(entity.id))
-    assertEquals(entity.id, newIndex.getIds(3)?.get(0))
+    assertEquals(3, newIndex!!.getDataById(entity))
+    assertEquals(entity, newIndex.getIds(3)?.get(0))
   }
 
   @Test
@@ -37,26 +37,26 @@ class ExternalEntityIndexTest {
 
     val index = builder.getOrCreateExternalIndex<Int>(INDEX_ID)
     val entity = builder.addPSourceEntity("hello", PSampleEntitySource("source"))
-    index.index(entity.id, 1)
-    assertEquals(1, index.getDataById(entity.id))
+    index.index(entity, 1)
+    assertEquals(1, index.getDataById(entity))
 
     val diff = PEntityStorageBuilder.from(builder.toStorage())
     val externalIndex = diff.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(externalIndex)
     assertNotEquals(index, externalIndex)
-    assertEquals(1, externalIndex!!.getDataById(entity.id))
-    externalIndex.update(entity.id, 2)
-    assertEquals(1, index.getDataById(entity.id))
-    assertEquals(2, externalIndex.getDataById(entity.id))
+    assertEquals(1, externalIndex!!.getDataById(entity))
+    externalIndex.update(entity, 2)
+    assertEquals(1, index.getDataById(entity))
+    assertEquals(2, externalIndex.getDataById(entity))
 
     builder.addDiff(diff)
-    assertEquals(2, index.getDataById(entity.id))
+    assertEquals(2, index.getDataById(entity))
 
     val storage = builder.toStorage()
     val newIndex = storage.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(newIndex)
     assertNotEquals(index, newIndex)
-    assertEquals(2, newIndex!!.getDataById(entity.id))
+    assertEquals(2, newIndex!!.getDataById(entity))
   }
 
   @Test
@@ -65,26 +65,26 @@ class ExternalEntityIndexTest {
 
     val index = builder.getOrCreateExternalIndex<Int>(INDEX_ID)
     val entity = builder.addPSourceEntity("hello", PSampleEntitySource("source"))
-    index.index(entity.id, 1)
-    assertEquals(1, index.getDataById(entity.id))
+    index.index(entity, 1)
+    assertEquals(1, index.getDataById(entity))
 
     val diff = PEntityStorageBuilder.from(builder.toStorage())
     val externalIndex = diff.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(externalIndex)
     assertNotEquals(index, externalIndex)
-    assertEquals(1, externalIndex!!.getDataById(entity.id))
-    externalIndex.remove(entity.id)
-    assertEquals(1, index.getDataById(entity.id))
-    assertNull(externalIndex.getDataById(entity.id))
+    assertEquals(1, externalIndex!!.getDataById(entity))
+    externalIndex.remove(entity)
+    assertEquals(1, index.getDataById(entity))
+    assertNull(externalIndex.getDataById(entity))
 
     builder.addDiff(diff)
-    assertNull(index.getDataById(entity.id))
+    assertNull(index.getDataById(entity))
 
     val storage = builder.toStorage()
     val newIndex = storage.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(newIndex)
     assertNotEquals(index, newIndex)
-    assertNull(newIndex!!.getDataById(entity.id))
+    assertNull(newIndex!!.getDataById(entity))
   }
 
   @Test
@@ -93,32 +93,32 @@ class ExternalEntityIndexTest {
 
     val index = builder.getOrCreateExternalIndex<Int>(INDEX_ID)
     val entity = builder.addPSourceEntity("hello", PSampleEntitySource("source"))
-    index.index(entity.id, 1)
-    assertEquals(1, index.getDataById(entity.id))
+    index.index(entity, 1)
+    assertEquals(1, index.getDataById(entity))
 
     val diff = PEntityStorageBuilder.from(builder.toStorage())
     val newEntity = builder.addPSourceEntity("world", PSampleEntitySource("source"))
     val externalIndex = diff.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(externalIndex)
     assertNotEquals(index, externalIndex)
-    assertEquals(1, externalIndex!!.getDataById(entity.id))
-    externalIndex.update(newEntity.id, 2)
-    assertEquals(1, index.getDataById(entity.id))
-    assertEquals(1, externalIndex.getDataById(entity.id))
-    assertEquals(2, externalIndex.getDataById(newEntity.id))
+    assertEquals(1, externalIndex!!.getDataById(entity))
+    externalIndex.update(newEntity, 2)
+    assertEquals(1, index.getDataById(entity))
+    assertEquals(1, externalIndex.getDataById(entity))
+    assertEquals(2, externalIndex.getDataById(newEntity))
 
     builder.addDiff(diff)
-    assertEquals(1, index.getDataById(entity.id))
-    assertEquals(2, index.getDataById(newEntity.id))
+    assertEquals(1, index.getDataById(entity))
+    assertEquals(2, index.getDataById(newEntity))
 
     val storage = builder.toStorage()
     val newIndex = storage.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(newIndex)
     assertNotEquals(index, newIndex)
-    assertEquals(1, newIndex!!.getDataById(entity.id))
-    assertEquals(2, newIndex.getDataById(newEntity.id))
-    assertEquals(entity.id, newIndex.getIds(1)?.get(0))
-    assertEquals(newEntity.id, newIndex.getIds(2)?.get(0))
+    assertEquals(1, newIndex!!.getDataById(entity))
+    assertEquals(2, newIndex.getDataById(newEntity))
+    assertEquals(entity, newIndex.getIds(1)?.get(0))
+    assertEquals(newEntity, newIndex.getIds(2)?.get(0))
   }
 
   @Test
@@ -127,16 +127,16 @@ class ExternalEntityIndexTest {
 
     val index = builder.getOrCreateExternalIndex<Int>(INDEX_ID)
     val entity = builder.addPSourceEntity("hello", PSampleEntitySource("source"))
-    index.index(entity.id, 1)
-    assertEquals(1, index.getDataById(entity.id))
+    index.index(entity, 1)
+    assertEquals(1, index.getDataById(entity))
 
     val diff = PEntityStorageBuilder.from(builder.toStorage())
     diff.removeExternalIndex<Int>(INDEX_ID)
     assertNull(diff.getExternalIndex<Int>(INDEX_ID))
-    assertEquals(1, index.getDataById(entity.id))
+    assertEquals(1, index.getDataById(entity))
 
     builder.addDiff(diff)
-    assertEquals(1, index.getDataById(entity.id))
+    assertEquals(1, index.getDataById(entity))
     assertNull(builder.getExternalIndex<Int>(INDEX_ID))
 
     val storage = builder.toStorage()
@@ -152,21 +152,21 @@ class ExternalEntityIndexTest {
 
     val diff = PEntityStorageBuilder.from(builder.toStorage())
     val index = diff.getOrCreateExternalIndex<Int>(INDEX_ID)
-    index.index(entity.id, 1)
+    index.index(entity, 1)
     assertNull(builder.getExternalIndex<Int>(INDEX_ID))
-    assertEquals(1, index.getDataById(entity.id))
+    assertEquals(1, index.getDataById(entity))
 
     builder.addDiff(diff)
-    assertEquals(1, index.getDataById(entity.id))
+    assertEquals(1, index.getDataById(entity))
     val entityIndex = builder.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(entityIndex)
     assertNotEquals(index, entityIndex)
-    assertEquals(1, entityIndex!!.getDataById(entity.id))
+    assertEquals(1, entityIndex!!.getDataById(entity))
 
     val storage = builder.toStorage()
     val newIndex = storage.getExternalIndex<Int>(INDEX_ID)
     assertNotNull(newIndex)
     assertNotEquals(entityIndex, newIndex)
-    assertEquals(1, newIndex!!.getDataById(entity.id))
+    assertEquals(1, newIndex!!.getDataById(entity))
   }
 }

platform/workspaceModel-core/src/pstorage/PEntityStorage.kt

@@ -766,13 +766,16 @@ internal class PEntityStorageBuilder(
 
   @Suppress("UNCHECKED_CAST")
   fun <T> getOrCreateExternalIndex(identifier: String): MutableExternalEntityIndex<T> {
-    return externalIndices.computeIfAbsent(identifier) { MutableExternalEntityIndex<T>() }
-      as MutableExternalEntityIndex<T>
+    val index = externalIndices.computeIfAbsent(identifier) { MutableExternalEntityIndex<T>() } as MutableExternalEntityIndex<T>
+    index.setTypedEntityStorage(this)
+    return index
   }
 
   @Suppress("UNCHECKED_CAST")
   override fun <T> getExternalIndex(identifier: String): MutableExternalEntityIndex<T>? {
-    return externalIndices[identifier] as? MutableExternalEntityIndex<T>
+    val index = externalIndices[identifier] as? MutableExternalEntityIndex<T>
+    index?.setTypedEntityStorage(this)
+    return index
   }
 
   @Suppress("UNCHECKED_CAST")
@@ -977,7 +980,9 @@ internal sealed class AbstractPEntityStorage : TypedEntityStorage {
 
   @Suppress("UNCHECKED_CAST")
   open fun <T> getExternalIndex(identifier: String): ExternalEntityIndex<T>? {
-    return externalIndices[identifier] as? ExternalEntityIndex<T>
+    val index = externalIndices[identifier] as? ExternalEntityIndex<T>
+    index?.setTypedEntityStorage(this)
+    return index
   }
 }
 

platform/workspaceModel-core/src/pstorage/external/ExternalEntityIndex.kt

@@ -3,14 +3,28 @@ package com.intellij.workspace.api.pstorage.external
 
 import com.intellij.util.containers.BidirectionalMap
 import com.intellij.workspace.api.TypedEntity
+import com.intellij.workspace.api.pstorage.AbstractPEntityStorage
 import com.intellij.workspace.api.pstorage.PId
+import com.intellij.workspace.api.pstorage.PTypedEntity
 import com.intellij.workspace.api.pstorage.external.ExternalEntityIndex.MutableExternalEntityIndex.IndexLogRecord.Add
 import com.intellij.workspace.api.pstorage.external.ExternalEntityIndex.MutableExternalEntityIndex.IndexLogRecord.Remove
 import java.util.*
 
 open class ExternalEntityIndex<T> private constructor(internal val index: BidirectionalMap<PId<out TypedEntity>, T>) {
-  internal fun getIds(data: T): List<PId<out TypedEntity>>? = index.getKeysByValue(data)
-  internal fun getDataById(id: PId<out TypedEntity>): T? = index[id]
+  private var entityStorage: AbstractPEntityStorage? = null
+  internal fun getIds(data: T): List<TypedEntity>? {
+    val storage = entityStorage ?: return null
+    return index.getKeysByValue(data)?.toMutableList()?.mapNotNull { storage.entityDataById(it)?.createEntity(storage) }
+  }
+
+  internal fun getDataById(entity: TypedEntity): T? {
+    entity as PTypedEntity
+    return index[entity.id]
+  }
+
+  internal fun setTypedEntityStorage(storage: AbstractPEntityStorage) {
+    entityStorage = storage
+  }
 
   internal fun copyIndex(): BidirectionalMap<PId<out TypedEntity>, T> {
     val copy = BidirectionalMap<PId<out TypedEntity>, T>()
@@ -24,17 +38,32 @@ open class ExternalEntityIndex<T> private constructor(internal val index: Bidire
   ) : ExternalEntityIndex<T>(index) {
     constructor() : this(BidirectionalMap<PId<out TypedEntity>, T>(), mutableListOf())
 
-    internal fun index(id: PId<out TypedEntity>, data: T) {
+    internal fun index(entity: TypedEntity, data: T) {
+      entity as PTypedEntity
+      index(entity.id, data)
+    }
+
+    private fun index(id: PId<out TypedEntity>, data: T) {
       index[id] = data
       indexLog.add(Add(id, data))
     }
 
-    internal fun update(id: PId<out TypedEntity>, newData: T) {
+    internal fun update(entity: TypedEntity, newData: T) {
+      entity as PTypedEntity
+      update(entity.id, newData)
+    }
+
+    private fun update(id: PId<out TypedEntity>, newData: T) {
       remove(id)
       index(id, newData)
     }
 
-    internal fun remove(id: PId<out TypedEntity>) {
+    internal fun remove(entity: TypedEntity) {
+      entity as PTypedEntity
+      remove(entity.id)
+    }
+
+    private fun remove(id: PId<out TypedEntity>) {
       index.remove(id)
       indexLog.add(Remove(id))
     }