From bc460c94e5ade23b07d5983eb442cdbd16f573c0 Mon Sep 17 00:00:00 2001
From: Liudmila Kornilova <liudmila.kornilova@jetbrains.com>
Date: Mon, 2 Mar 2026 18:22:50 +0100
Subject: [PATCH] IJPL-237264 Optimize replaceBySource for root entities
 without symbolic id

(cherry picked from commit ef3ca691f15c4ca6ea69e39fedd54a767df57545)

IJ-MR-194156

GitOrigin-RevId: 9736bc2670726b72b44cdbfb6d416e8cf63456df
---
 ...WorkspaceModelBenchmarksPerformanceTest.kt | 19 +++++++++++++++++++
 .../impl/ImmutableEntityStorageImpl.kt        |  4 +++-
 .../storage/impl/ReplaceBySourceAsTree.kt     | 15 +++++++--------
 3 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/WorkspaceModelBenchmarksPerformanceTest.kt b/platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/WorkspaceModelBenchmarksPerformanceTest.kt
index d2da4beee8e1..e998cf9bc452 100644
--- a/platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/WorkspaceModelBenchmarksPerformanceTest.kt
+++ b/platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/WorkspaceModelBenchmarksPerformanceTest.kt
@@ -48,6 +48,7 @@ import com.intellij.platform.workspace.storage.testEntities.entities.OoChildWith
 import com.intellij.platform.workspace.storage.testEntities.entities.OoParentEntity
 import com.intellij.platform.workspace.storage.testEntities.entities.ParentEntity
 import com.intellij.platform.workspace.storage.testEntities.entities.ParentMultipleEntity
+import com.intellij.platform.workspace.storage.testEntities.entities.WithSoftLinkEntity
 import com.intellij.platform.workspace.storage.testEntities.entities.modifyNamedChildEntity
 import com.intellij.platform.workspace.storage.testEntities.entities.modifyNamedEntity
 import com.intellij.platform.workspace.storage.testEntities.entities.modifyOoChildWithNullableParentEntity
@@ -229,6 +230,24 @@ class WorkspaceModelBenchmarksPerformanceTest {
       .attempts(1).start()
   }
 
+  @Test
+  fun replaceBySourceALotOfRootEntitiesWithoutSymbolicId(testInfo: TestInfo) {
+    val source: MutableEntityStorage = MutableEntityStorage.create()
+    val target: MutableEntityStorage = MutableEntityStorage.create()
+    val size = 30_000
+
+    repeat(size) {
+      source addEntity WithSoftLinkEntity(NameId("$it"), MySource) // could be any entity without symbolic id
+      target addEntity WithSoftLinkEntity(NameId("$it"), MySource)
+    }
+
+    Benchmark.newBenchmark(testInfo.displayName) {
+      target.replaceBySource({ true }, source.toSnapshot())
+    }
+      .warmupIterations(0)
+      .attempts(1).start()
+  }
+
   @Test
   fun serializeCommunityProject(testInfo: TestInfo) {
     val storageBuilder = MutableEntityStorage.create()
diff --git a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ImmutableEntityStorageImpl.kt b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ImmutableEntityStorageImpl.kt
index 8b6fc29340b3..1cb7689b6f0d 100644
--- a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ImmutableEntityStorageImpl.kt
+++ b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ImmutableEntityStorageImpl.kt
@@ -203,7 +203,9 @@ internal class MutableEntityStorageImpl(
 
   override fun <E : WorkspaceEntity> entities(entityClass: Class<E>): Sequence<E> = getEntitiesTimeMs.addMeasuredTime {
     @Suppress("UNCHECKED_CAST")
-    entitiesByType[entityClass.toClassId()]?.all()?.map { it.createEntity(this) } as? Sequence<E> ?: emptySequence()
+    entitiesByType[entityClass.toClassId()]?.all()?.map {
+      it.createEntity(this) // we have to wrap entity data with the link to the entity storage to be able to resolve parent/children
+    } as? Sequence<E> ?: emptySequence()
   }
 
   override fun <E : WorkspaceEntityWithSymbolicId, R : WorkspaceEntity> referrers(
diff --git a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ReplaceBySourceAsTree.kt b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ReplaceBySourceAsTree.kt
index 7b6b0b3f9127..90e892ff9683 100644
--- a/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ReplaceBySourceAsTree.kt
+++ b/platform/workspace/storage/src/com/intellij/platform/workspace/storage/impl/ReplaceBySourceAsTree.kt
@@ -914,14 +914,13 @@ internal class ReplaceBySourceAsTree {
       }
       else {
         val oppositeEntityData = oppositeStorage.entityDataByIdOrDie(rootEntity.id)
-        goalStorage.entities(rootEntity.id.clazz.findWorkspaceEntity())
-          .filter {
-            val itId = (it as WorkspaceEntityBase).id
-            if (goalState[itId] != null) return@filter false
-            goalStorage.entityDataByIdOrDie(itId).equalsByKey(oppositeEntityData) && goalStorage.refs.getParentRefsOfChild(itId.asChild())
-              .isEmpty()
-          }
-          .firstOrNull()
+        val entityClass = rootEntity.id.clazz.findWorkspaceEntity()
+        val data = goalStorage.entitiesByType[entityClass.toClassId()]?.all()?.firstOrNull { entityData ->
+          if (!entityData.equalsByKey(oppositeEntityData)) return@firstOrNull false
+          val itId = entityData.createEntityId()
+          goalState[itId] == null && goalStorage.refs.getParentRefsOfChild(itId.asChild()).isEmpty()
+        }
+        data?.createEntity(goalStorage)
       }
     }
   }