[Workspace Model] Add an API to get an entity reference from the entity directly

GitOrigin-RevId: c5d1cb47142970969c45d1474ed7b99da3773ede

platform/workspaceModel/storage/src/com/intellij/workspaceModel/storage/WorkspaceEntityStorage.kt

@@ -35,6 +35,8 @@ interface WorkspaceEntity {
    * internal IDs of the corresponding entities.
    */
   fun hasEqualProperties(e: WorkspaceEntity): Boolean
+
+  fun <E : WorkspaceEntity> createReference(): EntityReference<E>
 }
 
 /**

platform/workspaceModel/storage/src/com/intellij/workspaceModel/storage/impl/Entities.kt

@@ -123,6 +123,10 @@ abstract class WorkspaceEntityBase : ReferableWorkspaceEntity, Any() {
     }
   }
 
+  override fun <E : WorkspaceEntity> createReference(): EntityReference<E> {
+    return EntityReferenceImpl(this.id)
+  }
+
   override fun toString(): String = "$id"
 
   override fun equals(other: Any?): Boolean {

platform/workspaceModel/storage/tests/testSrc/com/intellij/workspaceModel/storage/ReferencesInStorageTest.kt

@@ -1,16 +1,9 @@
 // Copyright 2000-2020 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.intellij.workspaceModel.storage
 
+import com.intellij.workspaceModel.storage.entities.*
 import com.intellij.workspaceModel.storage.impl.WorkspaceEntityStorageBuilderImpl
 import com.intellij.workspaceModel.storage.impl.url.VirtualFileUrlManagerImpl
-import com.intellij.workspaceModel.storage.entities.*
-import com.intellij.workspaceModel.storage.entities.ModifiableChildEntity
-import com.intellij.workspaceModel.storage.entities.ModifiableChildWithOptionalParentEntity
-import com.intellij.workspaceModel.storage.entities.ModifiableParentEntity
-import com.intellij.workspaceModel.storage.entities.ChildEntity
-import com.intellij.workspaceModel.storage.entities.ChildWithOptionalParentEntity
-import com.intellij.workspaceModel.storage.entities.DataClass
-import com.intellij.workspaceModel.storage.entities.ParentEntity
 import com.intellij.workspaceModel.storage.url.VirtualFileUrlManager
 import org.junit.Assert
 import org.junit.Assert.assertEquals
@@ -168,6 +161,18 @@ class ReferencesInStorageTest {
     assertEquals(emptyList<ParentEntity>(), builder.entities(ParentEntity::class.java).toList())
   }
 
+  @Test
+  fun `remove parent entity referenced via two paths via entity ref`() {
+    val builder = WorkspaceEntityStorageBuilderImpl.create()
+    val parent = builder.addParentEntity()
+    builder.addChildEntity(parent, "child", DataClass("data", parent.createReference()))
+    builder.assertConsistency()
+    builder.removeEntity(parent)
+    builder.assertConsistency()
+    assertEquals(emptyList<ChildEntity>(), builder.entities(ChildEntity::class.java).toList())
+    assertEquals(emptyList<ParentEntity>(), builder.entities(ParentEntity::class.java).toList())
+  }
+
   @Test
   fun `modify parent property`() {
     val builder = WorkspaceEntityStorageBuilderImpl.create()
@@ -264,6 +269,26 @@ class ReferencesInStorageTest {
     assertEquals(oldParent, child.dataClass!!.parent.resolve(builder))
   }
 
+  @Test
+  fun `modify reference to parent via data class via entity ref`() {
+    val builder = WorkspaceEntityStorageBuilderImpl.create()
+    val parent1 = builder.addParentEntity("parent1")
+    val oldParent = builder.addParentEntity("parent2")
+    val child = builder.addChildEntity(parent1, "child", DataClass("data", oldParent.createReference()))
+    val newParent = builder.addParentEntity("new")
+    builder.assertConsistency()
+    val newChild = builder.modifyEntity(ModifiableChildEntity::class.java, child) {
+      dataClass = DataClass("data2", newParent.createReference())
+    }
+    builder.assertConsistency()
+    assertEquals("child", newChild.childProperty)
+    assertEquals("data2", newChild.dataClass!!.stringProperty)
+    assertEquals(setOf(oldParent, newParent, parent1), builder.entities(ParentEntity::class.java).toSet())
+    assertEquals(newChild, builder.singleChild())
+    assertEquals(newParent, newChild.dataClass.parent.resolve(builder))
+    assertEquals(oldParent, child.dataClass!!.parent.resolve(builder))
+  }
+
   @Test
   fun `builder from storage`() {
     val storage = WorkspaceEntityStorageBuilderImpl.create().apply {