IJPL-159611 Optimize querying of document guarded blocks

This commit makes `getOffsetGuard` and `getRangeGuard` faster by replacing O(n) algorithm with O(log(n)) algorithm. The previous approach used linear search to find the appropriate range marker that was a bottleneck on a workload with thousands of guarded blocks. Also, the guarded blocks were kept on hard references in an array. Preserving that required to adjust `keepIntervalOnWeakReference` method in the range marker tree to avoid collecting the guarded blocks by gc GitOrigin-RevId: f83e6d643846e805ee0dad864a1b09cb526ec323
2026-09-27 10:03:11 +07:00 · 2024-08-07 19:52:26 +00:00
parent e18e2e9dce
commit 7ba1e0310d
8 changed files with 235 additions and 32 deletions
@@ -272,6 +272,10 @@ public interface Document extends UserDataHolder {
  /**
   * Marks a range of text in the document as read-only (attempts to modify text in the
   * range cause {@link ReadOnlyFragmentModificationException} to be thrown).
+   * This range marker, unlike the ones created via {@link #createRangeMarker},
+   * is not automatically removed from the document if there are no more references to it.
+   * Therefore, there is no need for the caller to keep a reference to the created marker.
+   * However, it is the caller's responsibility to remove created marker via {@link #removeGuardedBlock}.
   *
   * @param startOffset the start offset of the text range to mark as read-only.
   * @param endOffset   the end offset of the text range to mark as read-only.
@@ -40,20 +40,23 @@ import java.lang.ref.SoftReference;
 import java.lang.ref.WeakReference;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;

 import static com.intellij.reference.SoftReference.dereference;

 public final class DocumentImpl extends UserDataHolderBase implements DocumentEx {
  private static final Logger LOG = Logger.getInstance(DocumentImpl.class);
  private static final int STRIP_TRAILING_SPACES_BULK_MODE_LINES_LIMIT = 1000;
+  private static final List<RangeMarker> GUARDED_IN_PROGRESS = new ArrayList<>(0);

  private final LockFreeCOWSortedArray<DocumentListener> myDocumentListeners =
    new LockFreeCOWSortedArray<>(PrioritizedDocumentListener.COMPARATOR, DocumentListener.ARRAY_FACTORY);
  private final RangeMarkerTree<RangeMarkerEx> myRangeMarkers = new RangeMarkerTree<>(this);
-  private final RangeMarkerTree<RangeMarkerEx> myPersistentRangeMarkers = new RangeMarkerTree<>(this);
-  private final List<RangeMarker> myGuardedBlocks = new ArrayList<>();
+  private final RangeMarkerTree<RangeMarkerEx> myPersistentRangeMarkers = new PersistentRangeMarkerTree(this);
+  private final AtomicReference<List<RangeMarker>> myGuardedBlocks = new AtomicReference<>();
  private ReadonlyFragmentModificationHandler myReadonlyFragmentModificationHandler;

  private final Object myLineSetLock = ObjectUtils.sentinel("line set lock");
@@ -438,48 +441,71 @@ public final class DocumentImpl extends UserDataHolderBase implements DocumentEx
  @Override
  public @NotNull RangeMarker createGuardedBlock(int startOffset, int endOffset) {
    LOG.assertTrue(startOffset <= endOffset, "Should be startOffset <= endOffset");
-    RangeMarker block = createRangeMarker(startOffset, endOffset, true);
-    myGuardedBlocks.add(block);
+    GuardedBlock block = new GuardedBlock(this, startOffset, endOffset);
+    myGuardedBlocks.set(null);
    return block;
  }

  @Override
  public void removeGuardedBlock(@NotNull RangeMarker block) {
-    myGuardedBlocks.remove(block);
+    if (!GuardedBlock.isGuarded(block)) {
+      throw new IllegalArgumentException("range markers is not a guarded block");
+    }
+    block.dispose();
+    myGuardedBlocks.set(null);
  }

  @Override
  public @NotNull List<RangeMarker> getGuardedBlocks() {
-    return myGuardedBlocks;
+    List<RangeMarker> cachedBlocks = myGuardedBlocks.get();
+    if (cachedBlocks != null && cachedBlocks != GUARDED_IN_PROGRESS) {
+      return cachedBlocks;
+    }
+    if (myGuardedBlocks.compareAndSet(null, GUARDED_IN_PROGRESS)) {
+      List<RangeMarker> blocks = collectGuardedBlocks();
+      if (!myGuardedBlocks.compareAndSet(GUARDED_IN_PROGRESS, blocks)) {
+        // another thread created or removed a block, force recalculation
+        myGuardedBlocks.set(null);
+      }
+      return blocks;
+    }
+    // another thread is already collecting the result, return without commiting
+    return collectGuardedBlocks();
+  }
+
+  private @NotNull List<RangeMarker> collectGuardedBlocks() {
+    List<RangeMarker> blocks = new ArrayList<>();
+    myPersistentRangeMarkers.processAll(GuardedBlock.processor(block -> {
+      blocks.add(block);
+      return true;
+    }));
+    // prevent the users from being misled that modifying this list affects actual guarded blocks
+    return Collections.unmodifiableList(blocks);
  }

  @Override
  public RangeMarker getOffsetGuard(int offset) {
-    List<RangeMarker> guardedBlocks = myGuardedBlocks;
-    // Way too much garbage would be produced otherwise in AbstractList.iterator()
-    //noinspection ForLoopReplaceableByForEach
-    for (int i = 0; i < guardedBlocks.size(); i++) {
-      RangeMarker block = guardedBlocks.get(i);
-      if (offsetInRange(offset, block.getStartOffset(), block.getEndOffset())) return block;
-    }
-
-    return null;
+    Ref<RangeMarker> blockRef = new Ref<>();
+    myPersistentRangeMarkers.processContaining(offset, GuardedBlock.processor(block -> {
+      blockRef.set(block);
+      return false;
+    }));
+    return blockRef.get();
  }

  @Override
  public RangeMarker getRangeGuard(int start, int end) {
-    List<RangeMarker> guardedBlocks = myGuardedBlocks;
-    //noinspection ForLoopReplaceableByForEach
-    for (int i = 0; i < guardedBlocks.size(); i++) {
-      RangeMarker block = guardedBlocks.get(i);
+    Ref<RangeMarker> blockRef = new Ref<>();
+    myPersistentRangeMarkers.processOverlappingWith(start, end, GuardedBlock.processor(block -> {
      if (rangesIntersect(start, end, true, true,
                          block.getStartOffset(), block.getEndOffset(),
                          block.isGreedyToLeft(), block.isGreedyToRight())) {
-        return block;
+        blockRef.set(block);
+        return false;
      }
-    }
-
-    return null;
+      return true;
+    }));
+    return blockRef.get();
  }

  @Override
@@ -493,10 +519,6 @@ public final class DocumentImpl extends UserDataHolderBase implements DocumentEx
    myCheckGuardedBlocks--;
  }

-  private static boolean offsetInRange(int offset, int start, int end) {
-    return start <= offset && offset < end;
-  }
-
  private static boolean rangesIntersect(int start0, int end0, boolean start0Inclusive, boolean end0Inclusive,
                                         int start1, int end1, boolean start1Inclusive, boolean end1Inclusive) {
    if (start0 > start1 || start0 == start1 && !start0Inclusive) {
@@ -1226,6 +1248,23 @@ public final class DocumentImpl extends UserDataHolderBase implements DocumentEx
    if (myDoingBulkUpdate) throw new UnexpectedBulkUpdateStateException(myBulkUpdateEnteringTrace);
  }

+  /**
+   * RangeMarkerTree that keeps all intervals on weak references except the guarded blocks.
+   * This class must be static because it should not capture 'this' reference to the document.
+   * Otherwise, there will be a chain of hard references {@code file -> tree -> document} and gc won't collect the document
+   */
+  private static final class PersistentRangeMarkerTree extends RangeMarkerTree<RangeMarkerEx> {
+    PersistentRangeMarkerTree(@NotNull Document document) {
+      super(document);
+    }
+
+    @Override
+    protected boolean keepIntervalOnWeakReference(@NotNull RangeMarkerEx interval) {
+      // prevent guarded blocks to be collected by gc
+      return !GuardedBlock.isGuarded(interval);
+    }
+  }
+
  private static final class UnexpectedBulkUpdateStateException extends RuntimeException implements ExceptionWithAttachments {
    private final Attachment[] myAttachments;

@@ -0,0 +1,33 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.intellij.openapi.editor.impl;
+
+import com.intellij.openapi.editor.RangeMarker;
+import com.intellij.openapi.editor.ex.DocumentEx;
+import com.intellij.openapi.editor.ex.RangeMarkerEx;
+import com.intellij.util.Processor;
+import org.jetbrains.annotations.NotNull;
+
+final class GuardedBlock extends PersistentRangeMarker {
+
+  static @NotNull Processor<RangeMarkerEx> processor(Processor<? super RangeMarkerEx> processor) {
+    return rangeMarker -> {
+      if (isGuarded(rangeMarker)) {
+        return processor.process(rangeMarker);
+      }
+      return true;
+    };
+  }
+
+  static boolean isGuarded(@NotNull RangeMarker rangeMarker) {
+    return rangeMarker instanceof GuardedBlock;
+  }
+
+  GuardedBlock(DocumentEx document, int startOffset, int endOffset) {
+    super(document, startOffset, endOffset, true);
+  }
+
+  @Override
+  public @NotNull String toString() {
+    return super.toString().replace("PersistentRangeMarker", "GuardedBlock");
+  }
+}
@@ -353,7 +353,7 @@ abstract class IntervalTreeImpl<T> extends RedBlackTree<T> implements IntervalTr
  }

  private @NotNull Supplier<? extends T> createGetter(@NotNull T interval) {
-    return keepIntervalsOnWeakReferences()
+    return keepIntervalOnWeakReference(interval)
           ? new WeakReferencedGetter<>(interval, myReferenceQueue)
           : new StaticSupplier<>(interval);
  }
@@ -405,7 +405,7 @@ abstract class IntervalTreeImpl<T> extends RedBlackTree<T> implements IntervalTr
    }
  }

-  protected boolean keepIntervalsOnWeakReferences() {
+  protected boolean keepIntervalOnWeakReference(@NotNull T interval) {
    return true;
  }

@@ -19,7 +19,7 @@ final class RangeHighlighterTree extends RangeMarkerTree<RangeHighlighterEx> {
  }

  @Override
-  protected boolean keepIntervalsOnWeakReferences() {
+  protected boolean keepIntervalOnWeakReference(@NotNull RangeHighlighterEx interval) {
    return false;
  }

@@ -13,7 +13,7 @@ class HardReferencingRangeMarkerTree<T extends RangeMarkerImpl> extends RangeMar
  }

  @Override
-  protected boolean keepIntervalsOnWeakReferences() {
+  protected boolean keepIntervalOnWeakReference(@NotNull T interval) {
    return false;
  }

@@ -4,7 +4,13 @@ import com.intellij.codeInsight.lookup.LookupManager;
 import com.intellij.openapi.actionSystem.IdeActions;
 import com.intellij.openapi.editor.Document;
 import com.intellij.openapi.editor.RangeMarker;
+import com.intellij.openapi.editor.ex.DocumentEx;
+import com.intellij.openapi.util.TextRange;
 import com.intellij.testFramework.fixtures.BasePlatformTestCase;
+import com.intellij.util.CommonProcessors;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.List;

 public class GuardBlockTest extends BasePlatformTestCase {
  private RangeMarker createGuard(final int start, final int end) {
@@ -157,4 +163,125 @@ public class GuardBlockTest extends BasePlatformTestCase {
    }
    fail("must be read only at " + 3);
  }
+
+  public void testGuardedBlockListIsEmpty() {
+    myFixture.configureByText("x.txt", "0123456789");
+    assertEmpty("expected no guarded block in the document", getGuardedBlocks());
+  }
+
+  public void testCreateGuardedBlock() {
+    myFixture.configureByText("x.txt", "0123456789");
+    assertEmpty("expected no guarded block in the document", getGuardedBlocks());
+    createGuard(0, 1);
+    assertEquals("expected one guarded block in the document", 1, getGuardedBlocks().size());
+    createGuard(0, 1);
+    assertEquals("expected two guarded blocks in the document", 2, getGuardedBlocks().size());
+  }
+
+  public void testRemoveGuardedBlock() {
+    myFixture.configureByText("x.txt", "0123456789");
+    getDocument().removeGuardedBlock(createGuard(0, 1));
+    assertEmpty("expected the guarded block to be removed", getGuardedBlocks());
+  }
+
+  public void testCreateRemoveCreateGuardedBlock() {
+    myFixture.configureByText("x.txt", "0123456789");
+    getDocument().removeGuardedBlock(createGuard(0, 1));
+    createGuard(0, 1);
+    assertEquals(1, getGuardedBlocks().size());
+  }
+
+  public void testGetGuardedBlockByOffset() {
+    myFixture.configureByText("x.txt", "0123456789");
+    DocumentEx document = getDocument();
+    RangeMarker guard = createGuard(1, 2);
+    assertNull("no guarded block expected covering the offset", document.getOffsetGuard(0));
+    assertNull("no guarded block expected covering the offset", document.getOffsetGuard(2));
+    assertEquals("guarded block expected covering the offset", guard, document.getOffsetGuard(1));
+  }
+
+  public void testGetRangeGuardInRange() {
+    myFixture.configureByText("x.txt", "0123456789");
+    DocumentEx document = getDocument();
+    RangeMarker guard = createGuard(2, 5);
+    assertEquals("intersection expected ( [ ] )", guard, document.getRangeGuard(0, 6));
+    assertEquals("intersection expected ( [ ) ]", guard, document.getRangeGuard(0, 3));
+    assertEquals("intersection expected [ ( ] )", guard, document.getRangeGuard(4, 6));
+    assertEquals("intersection expected [ () ]", guard, document.getRangeGuard(3, 4));
+    assertEquals("intersection expected [()  ]", guard, document.getRangeGuard(2, 3));
+    assertEquals("intersection expected [  ()]", guard, document.getRangeGuard(4, 5));
+  }
+
+  public void testGetRangeGuardNotInRange() {
+    myFixture.configureByText("x.txt", "0123456789");
+    DocumentEx document = getDocument();
+    createGuard(2, 5);
+    assertNull("no intersection expected () [ ]", document.getRangeGuard(0, 1));
+    assertNull("no intersection expected [ ] ()", document.getRangeGuard(6, 7));
+  }
+
+  public void testGetRangeGuardInRangeGreedy() {
+    myFixture.configureByText("x.txt", "0123456789");
+    DocumentEx document = getDocument();
+    RangeMarker guard = createGuard(2, 5);
+    assertNull("no intersection expected ( | ]", document.getRangeGuard(0, 2));
+    assertNull("no intersection expected [ | )", document.getRangeGuard(5, 6));
+    guard.setGreedyToLeft(true);
+    guard.setGreedyToRight(true);
+    assertEquals("intersection expected ( | ]", guard, document.getRangeGuard(0, 2));
+    assertEquals("intersection expected [ | )", guard, document.getRangeGuard(5, 6));
+  }
+
+  public void testGuardIsAvailableViaDocumentRangeMarkerTree() {
+    myFixture.configureByText("x.txt", "0123456789");
+    createGuard(0, 1);
+    var collector = new CommonProcessors.CollectProcessor<>();
+    getDocument().processRangeMarkers(collector);
+    assertEquals("guarded block expected to be available via range marker tree", 1, collector.getResults().size());
+  }
+
+  public void testGuardedBlockApiIsNotBypassed() {
+    myFixture.configureByText("x.txt", "0123456789");
+    RangeMarker guard = createGuard(0, 1);
+    assertThrows(Exception.class, () -> getGuardedBlocks().add(guard));
+    assertThrows(Exception.class, () -> getGuardedBlocks().clear());
+    assertThrows(Exception.class, () -> getGuardedBlocks().remove(guard));
+    assertEquals("guarded block list must not be modified directly", 1, getGuardedBlocks().size());
+    assertNotEmpty(getGuardedBlocks());
+  }
+
+  @SuppressWarnings("CallToSystemGC")
+  public void testGuardIsNotCollectedByGC() throws InterruptedException {
+    myFixture.configureByText("x.txt", "0123456789");
+    createGuard(0, 1);
+    System.gc(); Thread.sleep(200);
+    System.gc(); Thread.sleep(200);
+    List<RangeMarker> blocks = getGuardedBlocks();
+    RangeMarker guard = blocks.get(0);
+
+    assertTrue("guarded block expected to be valid", guard.isValid());
+    assertEquals("guarded block bounds expected to be preserved", new TextRange(0, 1), guard.getTextRange());
+    assertEquals("expected only one guarded block in the document", 1, blocks.size());
+  }
+
+  public void testRemovedGuardIsDisposed() {
+    myFixture.configureByText("x.txt", "0123456789");
+    RangeMarker guard = createGuard(0, 1);
+    getDocument().removeGuardedBlock(guard);
+
+    assertEmpty("expected no guarded block in document", getDocument().getGuardedBlocks());
+    assertFalse("removed guard expected to be not valid", guard.isValid());
+
+    var collector = new CommonProcessors.CollectProcessor<>();
+    getDocument().processRangeMarkers(collector);
+    assertEmpty("removed guard expected to be unregistered from range tree", collector.getResults());
+  }
+
+  private @NotNull List<RangeMarker> getGuardedBlocks() {
+    return getDocument().getGuardedBlocks();
+  }
+
+  private @NotNull DocumentEx getDocument() {
+    return ((DocumentEx)myFixture.getEditor().getDocument());
+  }
 }
@@ -95,7 +95,7 @@ internal class TerminalPromptModelImpl(
  @RequiresEdt
  private fun doUpdatePrompt(renderingInfo: TerminalPromptRenderingInfo) {
    DocumentUtil.writeInRunUndoTransparentAction {
-      document.guardedBlocks.clear()
+      document.guardedBlocks.forEach { document.removeGuardedBlock(it) }
      document.replaceString(0, commandStartOffset, renderingInfo.text)
      document.createGuardedBlock(0, renderingInfo.text.length)
    }