From 30f0c1475d9acf9d843f70be778f04d18ae5f4c3 Mon Sep 17 00:00:00 2001
From: Peng Lyu <penn.lv@gmail.com>
Date: Fri, 7 Jun 2024 15:12:31 -0700
Subject: [PATCH] TextBuffer API: getNearestChunk to support incremental tree
 sitter parsing. (#214648)

---
 src/vs/editor/common/model.ts                 |  5 +++++
 .../pieceTreeTextBuffer/pieceTreeBase.ts      | 21 +++++++++++++++++++
 .../pieceTreeTextBuffer.ts                    |  4 ++++
 .../pieceTreeTextBuffer.test.ts               | 16 ++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/src/vs/editor/common/model.ts b/src/vs/editor/common/model.ts
index e21aa7d600c7a..134234bbfe59b 100644
--- a/src/vs/editor/common/model.ts
+++ b/src/vs/editor/common/model.ts
@@ -1431,6 +1431,11 @@ export interface IReadonlyTextBuffer {
 	getLineFirstNonWhitespaceColumn(lineNumber: number): number;
 	getLineLastNonWhitespaceColumn(lineNumber: number): number;
 	findMatchesLineByLine(searchRange: Range, searchData: SearchData, captureMatches: boolean, limitResultCount: number): FindMatch[];
+
+	/**
+	 * Get nearest chunk of text after `offset` in the text buffer.
+	 */
+	getNearestChunk(offset: number): string;
 }
 
 /**
diff --git a/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeBase.ts b/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeBase.ts
index b75d0d75a70ef..24f90651f9502 100644
--- a/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeBase.ts
+++ b/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeBase.ts
@@ -666,6 +666,27 @@ export class PieceTreeBase {
 		return this._getCharCode(nodePos);
 	}
 
+	public getNearestChunk(offset: number): string {
+		const nodePos = this.nodeAt(offset);
+		if (nodePos.remainder === nodePos.node.piece.length) {
+			// the offset is at the head of next node.
+			const matchingNode = nodePos.node.next();
+			if (!matchingNode || matchingNode === SENTINEL) {
+				return '';
+			}
+
+			const buffer = this._buffers[matchingNode.piece.bufferIndex];
+			const startOffset = this.offsetInBuffer(matchingNode.piece.bufferIndex, matchingNode.piece.start);
+			return buffer.buffer.substring(startOffset, startOffset + matchingNode.piece.length);
+		} else {
+			const buffer = this._buffers[nodePos.node.piece.bufferIndex];
+			const startOffset = this.offsetInBuffer(nodePos.node.piece.bufferIndex, nodePos.node.piece.start);
+			const targetOffset = startOffset + nodePos.remainder;
+			const targetEnd = startOffset + nodePos.node.piece.length;
+			return buffer.buffer.substring(targetOffset, targetEnd);
+		}
+	}
+
 	public findMatchesInNode(node: TreeNode, searcher: Searcher, startLineNumber: number, startColumn: number, startCursor: BufferCursor, endCursor: BufferCursor, searchData: SearchData, captureMatches: boolean, limitResultCount: number, resultLen: number, result: FindMatch[]) {
 		const buffer = this._buffers[node.piece.bufferIndex];
 		const startOffsetInBuffer = this.offsetInBuffer(node.piece.bufferIndex, node.piece.start);
diff --git a/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.ts b/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.ts
index 12d7e0b098179..a369298c0c97b 100644
--- a/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.ts
+++ b/src/vs/editor/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.ts
@@ -167,6 +167,10 @@ export class PieceTreeTextBuffer extends Disposable implements ITextBuffer {
 		return this.getValueLengthInRange(range, eol);
 	}
 
+	public getNearestChunk(offset: number): string {
+		return this._pieceTree.getNearestChunk(offset);
+	}
+
 	public getLength(): number {
 		return this._pieceTree.getLength();
 	}
diff --git a/src/vs/editor/test/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.test.ts b/src/vs/editor/test/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.test.ts
index 67a70b1bbb2ae..01cc7cb6ef60f 100644
--- a/src/vs/editor/test/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.test.ts
+++ b/src/vs/editor/test/common/model/pieceTreeTextBuffer/pieceTreeTextBuffer.test.ts
@@ -1817,6 +1817,22 @@ suite('buffer api', () => {
 		assert.strictEqual(pieceTable.getLineCharCode(2, 3), 'e'.charCodeAt(0), 'e');
 		assert.strictEqual(pieceTable.getLineCharCode(2, 4), '2'.charCodeAt(0), '2');
 	});
+
+	test('getNearestChunk', () => {
+		const pieceTree = createTextBuffer(['012345678']);
+		ds.add(pieceTree);
+		const pt = pieceTree.getPieceTree();
+
+		pt.insert(3, 'ABC');
+		assert.equal(pt.getLineContent(1), '012ABC345678');
+		assert.equal(pt.getNearestChunk(3), 'ABC');
+		assert.equal(pt.getNearestChunk(6), '345678');
+
+		pt.delete(9, 1);
+		assert.equal(pt.getLineContent(1), '012ABC34578');
+		assert.equal(pt.getNearestChunk(6), '345');
+		assert.equal(pt.getNearestChunk(9), '78');
+	});
 });
 
 suite('search offset cache', () => {