[Workspace Model] [IJPL-345] Make MutableEntityStorage#collectChanges internal and move it into instrumentation level

GitOrigin-RevId: 3222772bca9d22737e8c4e9fffbaadb5145d1fd9

java/compiler/impl/src/com/intellij/packaging/impl/artifacts/workspacemodel/ArtifactManagerBridge.kt

@@ -28,6 +28,8 @@ import com.intellij.platform.diagnostic.telemetry.Compiler
 import com.intellij.platform.diagnostic.telemetry.TelemetryManager
 import com.intellij.platform.diagnostic.telemetry.helpers.addMeasuredTimeMillis
 import com.intellij.platform.workspace.storage.*
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.util.concurrency.annotations.RequiresReadLock
 import com.intellij.util.concurrency.annotations.RequiresWriteLock
 import com.intellij.util.containers.BidirectionalMap
@@ -156,6 +158,7 @@ class ArtifactManagerBridge(private val project: Project) : ArtifactManager(), D
     return modificationTracker
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   @RequiresWriteLock
   fun commit(artifactModel: ArtifactModifiableModelBridge) = commitMs.addMeasuredTimeMillis {
     // XXX @RequiresReadLock annotation doesn't work for kt now
@@ -164,7 +167,7 @@ class ArtifactManagerBridge(private val project: Project) : ArtifactManager(), D
     updateCustomElements(artifactModel.diff)
 
     val current = project.workspaceModel.currentSnapshot
-    val changes = artifactModel.diff.collectChanges()[ArtifactEntity::class.java] ?: emptyList()
+    val changes = (artifactModel.diff as MutableEntityStorageInstrumentation).collectChanges()[ArtifactEntity::class.java] ?: emptyList()
 
     val removed = mutableSetOf<ArtifactBridge>()
     val added = mutableListOf<ArtifactBridge>()

java/compiler/impl/src/com/intellij/packaging/impl/artifacts/workspacemodel/ArtifactModifiableModelBridge.kt

@@ -22,6 +22,8 @@ import com.intellij.platform.diagnostic.telemetry.helpers.addMeasuredTimeMillis
 import com.intellij.platform.workspace.storage.EntityChange
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.impl.VersionedEntityStorageOnBuilder
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
 import com.intellij.util.EventDispatcher
 import com.intellij.util.concurrency.annotations.RequiresWriteLock
@@ -213,6 +215,7 @@ class ArtifactModifiableModelBridge(
     manager.commit(this)
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   override fun dispose() = disposeMs.addMeasuredTimeMillis {
     val artifacts: MutableList<Artifact> = ArrayList()
 
@@ -227,7 +230,7 @@ class ArtifactModifiableModelBridge(
     (ArtifactPointerManager.getInstance(project) as ArtifactPointerManagerImpl).disposePointers(artifacts)
 
     val current = WorkspaceModel.getInstance(project).currentSnapshot
-    val changes = diff.collectChanges()[ArtifactEntity::class.java] ?: emptyList()
+    val changes = (diff as MutableEntityStorageInstrumentation).collectChanges()[ArtifactEntity::class.java] ?: emptyList()
 
     val added = mutableListOf<ArtifactBridge>()
     val changed = mutableListOf<ArtifactBridge>()

platform/backend/workspace/src/BuilderSnapshot.kt

@@ -40,8 +40,9 @@ public class BuilderSnapshot @ApiStatus.Internal constructor(private val version
    * Execution of this function may take considerable time, so it's better to invoke it from a background thread, and only pass its result
    * to [WorkspaceModel.replaceProjectModel] under write-lock.
    */
+  @OptIn(EntityStorageInstrumentationApi::class)
   public fun getStorageReplacement(): StorageReplacement {
-    val changes = builder.collectChanges()
+    val changes = (builder as MutableEntityStorageInstrumentation).collectChanges()
     return StorageReplacement(version, builder, changes)
   }
 }

platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/GlobalWorkspaceModel.kt

@@ -14,11 +14,12 @@ import com.intellij.platform.backend.workspace.GlobalWorkspaceModelCache
 import com.intellij.platform.backend.workspace.WorkspaceModel
 import com.intellij.platform.diagnostic.telemetry.helpers.addMeasuredTimeMillis
 import com.intellij.platform.workspace.jps.JpsGlobalFileEntitySource
-import com.intellij.platform.workspace.jps.JpsProjectFileEntitySource
 import com.intellij.platform.workspace.jps.entities.*
 import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.VersionedEntityStorageImpl
 import com.intellij.platform.workspace.storage.impl.assertConsistency
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
 import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
 import com.intellij.util.concurrency.ThreadingAssertions
@@ -37,6 +38,7 @@ import org.jetbrains.annotations.NonNls
 import java.util.concurrent.atomic.AtomicLong
 import kotlin.system.measureTimeMillis
 
+@OptIn(EntityStorageInstrumentationApi::class)
 @ApiStatus.Internal
 class GlobalWorkspaceModel : Disposable {
   /**
@@ -86,11 +88,12 @@ class GlobalWorkspaceModel : Disposable {
     entityStorage = VersionedEntityStorageImpl(EntityStorageSnapshot.empty())
 
     val callback = JpsGlobalModelSynchronizer.getInstance().loadInitialState(mutableEntityStorage, entityStorage, loadedFromCache)
-    val changes = mutableEntityStorage.collectChanges()
+    val changes = (mutableEntityStorage as MutableEntityStorageInstrumentation).collectChanges()
     entityStorage.replace(mutableEntityStorage.toSnapshot(), changes, {}, {})
     callback.invoke()
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   fun updateModel(description: @NonNls String, updater: (MutableEntityStorage) -> Unit) {
     ThreadingAssertions.assertWriteAccess()
     if (modelVersionUpdate.get() == entityStorage.pointer.version) {
@@ -110,7 +113,7 @@ class GlobalWorkspaceModel : Disposable {
       }
       val changes: Map<Class<*>, List<EntityChange<*>>>
       collectChangesTimeMillis = measureTimeMillis {
-        changes = builder.collectChanges()
+        changes = (builder as MutableEntityStorageInstrumentation).collectChanges()
       }
       initializingTimeMillis = measureTimeMillis {
         this.initializeBridges(changes, builder)

platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/ModulePreUpdateHandler.kt

@@ -9,11 +9,14 @@ import com.intellij.platform.workspace.jps.entities.ModuleId
 import com.intellij.platform.workspace.storage.EntityChange
 import com.intellij.platform.workspace.storage.EntityStorage
 import com.intellij.platform.workspace.storage.MutableEntityStorage
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 
 class ModulePreUpdateHandler : WorkspaceModelPreUpdateHandler {
+  @OptIn(EntityStorageInstrumentationApi::class)
   override fun update(before: EntityStorage, builder: MutableEntityStorage): Boolean {
     // TODO: 21.12.2020 We need an api to find removed modules faster
-    val changes = builder.collectChanges()
+    val changes = (builder as MutableEntityStorageInstrumentation).collectChanges()
 
     val removedModuleSymbolicIds = LinkedHashSet<ModuleId>()
     changes[ModuleEntity::class.java]?.asSequence()?.forEach { change ->

platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/WorkspaceModelImpl.kt

@@ -17,6 +17,8 @@ import com.intellij.platform.diagnostic.telemetry.helpers.addMeasuredTimeMillis
 import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.VersionedEntityStorageImpl
 import com.intellij.platform.workspace.storage.impl.assertConsistency
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.serviceContainer.AlreadyDisposedException
 import com.intellij.workspaceModel.core.fileIndex.EntityStorageKind
 import com.intellij.workspaceModel.core.fileIndex.WorkspaceFileIndex
@@ -121,6 +123,7 @@ open class WorkspaceModelImpl(private val project: Project, private val cs: Coro
     loadedFromCache = false
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   @Synchronized
   final override fun updateProjectModel(description: @NonNls String, updater: (MutableEntityStorage) -> Unit) {
     ApplicationManager.getApplication().assertWriteAccessAllowed()
@@ -144,7 +147,7 @@ open class WorkspaceModelImpl(private val project: Project, private val cs: Coro
 
       val changes: Map<Class<*>, List<EntityChange<*>>>
       collectChangesTimeMillis = measureTimeMillis {
-        changes = builder.collectChanges()
+        changes = (builder as MutableEntityStorageInstrumentation).collectChanges()
       }
       initializingTimeMillis = measureTimeMillis {
         this.initializeBridges(changes, builder)
@@ -201,6 +204,7 @@ open class WorkspaceModelImpl(private val project: Project, private val cs: Coro
    *
    * **N.B** For more information on why this and other methods were marked by Synchronized see IDEA-313151
    */
+  @OptIn(EntityStorageInstrumentationApi::class)
   @ApiStatus.Obsolete
   @Synchronized
   fun updateProjectModelSilent(description: @NonNls String, updater: (MutableEntityStorage) -> Unit) {
@@ -212,7 +216,7 @@ open class WorkspaceModelImpl(private val project: Project, private val cs: Coro
 
     val generalTime = measureTimeMillis {
       val before = entityStorage.current
-      val builder = MutableEntityStorage.from(entityStorage.current)
+      val builder = MutableEntityStorage.from(entityStorage.current) as MutableEntityStorageInstrumentation
       updateTimeMillis = measureTimeMillis {
         updater(builder)
       }
@@ -262,13 +266,14 @@ open class WorkspaceModelImpl(private val project: Project, private val cs: Coro
     }
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   override fun updateUnloadedEntities(description: @NonNls String, updater: (MutableEntityStorage) -> Unit) {
     ApplicationManager.getApplication().assertWriteAccessAllowed()
     if (project.isDisposed) return
 
     val time = measureTimeMillis {
       val before = currentSnapshotOfUnloadedEntities
-      val builder = MutableEntityStorage.from(before)
+      val builder = MutableEntityStorage.from(before) as MutableEntityStorageInstrumentation
       updater(builder)
       startPreUpdateHandlers(before, builder)
       val changes = builder.collectChanges()

platform/projectModel-impl/src/com/intellij/workspaceModel/ide/impl/orphanage.kt

@@ -11,6 +11,8 @@ import com.intellij.platform.workspace.jps.OrphanageWorkerEntitySource
 import com.intellij.platform.workspace.jps.entities.*
 import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.VersionedEntityStorageImpl
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
 import com.intellij.util.concurrency.annotations.RequiresWriteLock
 import com.intellij.workspaceModel.ide.EntitiesOrphanage
@@ -23,6 +25,7 @@ class EntitiesOrphanageImpl(private val project: Project) : EntitiesOrphanage {
   override val currentSnapshot: EntityStorageSnapshot
     get() = entityStorage.current
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   @RequiresWriteLock
   override fun update(updater: (MutableEntityStorage) -> Unit) {
     ApplicationManager.getApplication().assertWriteAccessAllowed()
@@ -32,7 +35,7 @@ class EntitiesOrphanageImpl(private val project: Project) : EntitiesOrphanage {
 
     updater(builder)
 
-    val changes = builder.collectChanges()
+    val changes = (builder as MutableEntityStorageInstrumentation).collectChanges()
 
     checkIfParentsAlreadyExist(changes, builder)
 

platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/ImlReplaceBySourceTest.kt

@@ -11,6 +11,8 @@ import com.intellij.platform.workspace.jps.serialization.impl.JpsProjectEntities
 import com.intellij.platform.workspace.storage.EntityChange
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.impl.url.toVirtualFileUrl
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.tests.checkConsistency
 import com.intellij.platform.workspace.storage.toBuilder
 import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
@@ -26,6 +28,7 @@ import org.jetbrains.jps.model.serialization.PathMacroUtil
 import org.junit.*
 import java.io.File
 
+@OptIn(EntityStorageInstrumentationApi::class)
 class ImlReplaceBySourceTest {
   @Rule
   @JvmField
@@ -71,7 +74,7 @@ class ImlReplaceBySourceTest {
     val projectDir = temp.root.toVirtualFileUrl(virtualFileManager)
     val configLocation = JpsProjectConfigLocation.DirectoryBased(projectDir, projectDir.append(PathMacroUtil.DIRECTORY_STORE_NAME))
 
-    var builder = MutableEntityStorage.create()
+    var builder = MutableEntityStorage.create() as MutableEntityStorageInstrumentation
     JpsProjectEntitiesLoader.loadModule(moduleFile.toPath(), configLocation, builder, TestErrorReporter,
                                         SerializationContextForTests(virtualFileManager, CachingJpsFileContentReader(configLocation)))
 
@@ -97,7 +100,7 @@ class ImlReplaceBySourceTest {
 
     val before = builder.toSnapshot()
 
-    builder = before.toBuilder()
+    builder = before.toBuilder() as MutableEntityStorageInstrumentation
     builder.replaceBySource({ true }, replaceWith.toSnapshot())
 
     val changes = builder.collectChanges().values.flatten()
@@ -120,7 +123,7 @@ class ImlReplaceBySourceTest {
   }
 
   private fun replaceBySourceFullReplace(projectFile: File) {
-    var storageBuilder1 = MutableEntityStorage.create()
+    var storageBuilder1 = MutableEntityStorage.create() as MutableEntityStorageInstrumentation
     val data = com.intellij.workspaceModel.ide.impl.jps.serialization.loadProject(projectFile.asConfigLocation(virtualFileManager),
                                                                                   storageBuilder1, storageBuilder1, virtualFileManager)
 
@@ -132,7 +135,7 @@ class ImlReplaceBySourceTest {
     }
 
     val before = storageBuilder1.toSnapshot()
-    storageBuilder1 = before.toBuilder()
+    storageBuilder1 = before.toBuilder() as MutableEntityStorageInstrumentation
     storageBuilder1.checkConsistency()
     storageBuilder1.replaceBySource(sourceFilter = { true }, replaceWith = storageBuilder2.toSnapshot())
     storageBuilder1.checkConsistency()

platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/WorkspaceModelBenchmarksPerformanceTest.kt

@@ -21,6 +21,8 @@ import com.intellij.platform.workspace.storage.WorkspaceEntity
 import com.intellij.platform.workspace.storage.impl.cache.CacheResetTracker
 import com.intellij.platform.workspace.storage.impl.serialization.EntityStorageSerializerImpl
 import com.intellij.platform.workspace.storage.impl.url.toVirtualFileUrl
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.query.entities
 import com.intellij.platform.workspace.storage.query.flatMap
 import com.intellij.platform.workspace.storage.query.groupBy
@@ -456,10 +458,11 @@ class WorkspaceModelBenchmarksPerformanceTest {
     //    val metrics = mapOf("duration" to duration, "duration_replace" to applyStorageTime)
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   @Test
   fun `collect changes`(testInfo: TestInfo) {
     val builders = List(1000) {
-      val builder = MutableEntityStorage.create()
+      val builder = MutableEntityStorage.create() as MutableEntityStorageInstrumentation
 
       // Set initial state
       repeat(1000) {
@@ -493,7 +496,7 @@ class WorkspaceModelBenchmarksPerformanceTest {
             this.parentEntity = null
           }
         }
-      }
+      } as MutableEntityStorageInstrumentation
     }
 
     PlatformTestUtil.startPerformanceTest(testInfo.displayName, 100500) {

platform/workspace/jps/tests/testSrc/com/intellij/workspaceModel/ide/impl/jps/serialization/jpsTestUtils.kt

@@ -1,4 +1,6 @@
 // Copyright 2000-2022 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+@file:OptIn(EntityStorageInstrumentationApi::class)
+
 package com.intellij.workspaceModel.ide.impl.jps.serialization
 
 import com.intellij.configurationStore.StoreUtil.saveDocumentsAndProjectsAndApp
@@ -27,6 +29,8 @@ import com.intellij.platform.workspace.jps.entities.LibraryEntity
 import com.intellij.platform.workspace.jps.serialization.impl.*
 import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.url.toVirtualFileUrl
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
 import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
 import com.intellij.testFramework.UsefulTestCase
@@ -371,8 +375,8 @@ internal fun checkSaveProjectAfterChange(originalProjectFile: File,
                                          forceAllFilesRewrite: Boolean = false) {
   val projectData = copyAndLoadProject(originalProjectFile, virtualFileManager, unloadedModuleNameHolder, checkConsistencyAfterLoading,
                                        externalStorageConfigurationManager)
-  val builder = MutableEntityStorage.from(projectData.storage)
-  val unloadedEntitiesBuilder = MutableEntityStorage.from(projectData.unloadedEntitiesStorage)
+  val builder = MutableEntityStorage.from(projectData.storage) as MutableEntityStorageInstrumentation
+  val unloadedEntitiesBuilder = MutableEntityStorage.from(projectData.unloadedEntitiesStorage) as MutableEntityStorageInstrumentation
   change(builder, projectData.orphanage.toBuilder(), unloadedEntitiesBuilder, projectData.configLocation)
   val changesList = builder.collectChanges().values + unloadedEntitiesBuilder.collectChanges().values
   val changedSources = changesList.flatMapTo(HashSet()) { changes ->

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

@@ -179,18 +179,6 @@ public interface MutableEntityStorage : EntityStorage {
    */
   public fun replaceBySource(sourceFilter: (EntitySource) -> Boolean, replaceWith: EntityStorage)
 
-  /**
-   * Return changes in entities recorded in this instance. [original] parameter is used to get the old instances of modified
-   * and removed entities.
-   * 
-   * This function isn't supported to be used by client code directly. In order to subscribe to changes in entities inside the IDE process,
-   * use [WorkspaceModelTopics][com.intellij.platform.backend.workspace.WorkspaceModelTopics].
-   *
-   * To understand how the changes are collected see the KDoc for [com.intellij.platform.backend.workspace.WorkspaceModelChangeListener]
-   */
-  @ApiStatus.Internal
-  public fun collectChanges(): Map<Class<*>, List<EntityChange<*>>>
-
   /**
    * Merges changes from [diff] to this storage. 
    * It's supposed that [diff] was created either via [MutableEntityStorage.create] function, or via [MutableEntityStorage.from] with the 

platform/workspace/storage/src/com/intellij/platform/workspace/storage/instrumentation/EntityStorageInstrumentation.kt

@@ -54,6 +54,17 @@ public interface MutableEntityStorageInstrumentation : MutableEntityStorage, Ent
   @ApiStatus.Internal
   public fun hasSameEntities(): Boolean
 
+  /**
+   * Return changes in entities recorded in this instance.
+   *
+   * This function isn't supported to be used by client code directly. In order to subscribe to changes in entities inside the IDE process,
+   * use [WorkspaceModelChangeListener][com.intellij.platform.backend.workspace.WorkspaceModelChangeListener].
+   *
+   * To understand how the changes are collected see the KDoc for [com.intellij.platform.backend.workspace.WorkspaceModelChangeListener]
+   */
+  @ApiStatus.Internal
+  public fun collectChanges(): Map<Class<*>, List<EntityChange<*>>>
+
   /**
    * Replaces existing children of a given parent with a new list of children.
    *

platform/workspace/storage/tests/testSrc/com/intellij/platform/workspace/storage/tests/CollectChangesInBuilderTest.kt

@@ -5,6 +5,8 @@ import com.intellij.platform.workspace.storage.EntityChange
 import com.intellij.platform.workspace.storage.EntityStorageSnapshot
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.impl.url.VirtualFileUrlManagerImpl
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.testEntities.entities.*
 import com.intellij.platform.workspace.storage.toBuilder
 import org.junit.jupiter.api.BeforeEach
@@ -328,16 +330,18 @@ class CollectChangesInBuilderTest {
     assertIs<EntityChange.Replaced<*>>(changes[OptionalOneToOneChildEntity::class.java]?.single())
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   private fun assertChangelogSize(size: Int,
                                   myBuilder: MutableEntityStorage = builder,
                                   original: EntityStorageSnapshot = initialStorage): Map<Class<*>, List<EntityChange<*>>> {
-    val changes = myBuilder.collectChanges()
+    val changes = (myBuilder as MutableEntityStorageInstrumentation).collectChanges()
     assertEquals(size, changes.values.flatten().size)
     return changes
   }
 
+  @OptIn(EntityStorageInstrumentationApi::class)
   private fun collectSampleEntityChanges(): List<EntityChange<SampleEntity>> {
-    val changes = builder.collectChanges()
+    val changes = (builder as MutableEntityStorageInstrumentation).collectChanges()
     if (changes.isEmpty()) return emptyList()
     return changes.entries.single().value as List<EntityChange<SampleEntity>>
   }

platform/workspace/storage/tests/testSrc/com/intellij/platform/workspace/storage/tests/WorkspaceEntityEqualityTest.kt

@@ -3,6 +3,8 @@ package com.intellij.platform.workspace.storage.tests
 
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.impl.url.VirtualFileUrlManagerImpl
+import com.intellij.platform.workspace.storage.instrumentation.EntityStorageInstrumentationApi
+import com.intellij.platform.workspace.storage.instrumentation.MutableEntityStorageInstrumentation
 import com.intellij.platform.workspace.storage.testEntities.entities.SampleEntity
 import com.intellij.platform.workspace.storage.testEntities.entities.SampleEntitySource
 import com.intellij.platform.workspace.storage.testEntities.entities.modifyEntity
@@ -12,15 +14,16 @@ import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 import kotlin.test.*
 
+@OptIn(EntityStorageInstrumentationApi::class)
 class WorkspaceEntityEqualityTest {
 
-  private lateinit var builderOne: MutableEntityStorage
-  private lateinit var builderTwo: MutableEntityStorage
+  private lateinit var builderOne: MutableEntityStorageInstrumentation
+  private lateinit var builderTwo: MutableEntityStorageInstrumentation
 
   @BeforeEach
   fun setUp() {
-    builderOne = MutableEntityStorage.create()
-    builderTwo = MutableEntityStorage.create()
+    builderOne = MutableEntityStorage.create() as MutableEntityStorageInstrumentation
+    builderTwo = MutableEntityStorage.create() as MutableEntityStorageInstrumentation
   }
 
   @Test
@@ -124,7 +127,7 @@ class WorkspaceEntityEqualityTest {
     assertEquals(entityInEvent, entityInSnapshot)
     assertEquals(entityInSnapshot, entityInEvent)
 
-    val newBuilder = MutableEntityStorage.from(snapshot)
+    val newBuilder = MutableEntityStorage.from(snapshot) as MutableEntityStorageInstrumentation
     newBuilder.modifyEntity(entityInSnapshot) {
       stringProperty = "Data"
     }