[diffs/editor] Fix cursor move

packages/diffs/src/editor/pieceTable.ts

@@ -82,9 +82,8 @@ export class PieceTable {
   #piecesCache: Piece[] = [];
   #length = 0;
   #lineCount = 0;
-  #lastVisitedLine:
-    | [line: number, includeLineBreak: boolean, text: string]
-    | null = null;
+  #lastVisitedLine: [number, boolean, string] | null = null;
+  #lastVisitedLineLength: [number, boolean, number] | null = null;
 
   constructor(originalText: string) {
     this.#original = new TextBuffer(originalText);
@@ -120,9 +119,47 @@ export class PieceTable {
     }
     const text = this.getTextSlice(offset[0], offset[1], !includeLineBreak);
     this.#lastVisitedLine = [line, includeLineBreak, text];
+    this.#lastVisitedLineLength = [line, includeLineBreak, text.length];
     return text;
   }
 
+  getLineLength(line: number, includeLineBreak = false): number {
+    const lastVisitedLineLength = this.#lastVisitedLineLength;
+    const lastVisitedLine = this.#lastVisitedLine;
+    if (
+      lastVisitedLineLength !== null &&
+      lastVisitedLineLength[0] === line &&
+      lastVisitedLineLength[1] === includeLineBreak
+    ) {
+      return lastVisitedLineLength[2];
+    }
+    if (
+      lastVisitedLine !== null &&
+      lastVisitedLine[0] === line &&
+      lastVisitedLine[1] === includeLineBreak
+    ) {
+      const length = lastVisitedLine[2].length;
+      this.#lastVisitedLineLength = [line, includeLineBreak, length];
+      return length;
+    }
+    const offset = this.#getLineOffset(line);
+    if (offset === undefined) {
+      throw new Error(`Line index out of range: ${line}`);
+    }
+    const [start, end] = offset;
+    let length = end - start;
+    if (!includeLineBreak) {
+      while (
+        length > 0 &&
+        isEOL(this.charAt(start + length - 1).charCodeAt(0))
+      ) {
+        length--;
+      }
+    }
+    this.#lastVisitedLineLength = [line, includeLineBreak, length];
+    return length;
+  }
+
   getTextSlice(start: number, end: number, trimEOF = false): string {
     if (start >= end) {
       return '';
@@ -367,6 +404,7 @@ export class PieceTable {
 
     this.#setPieces(nextPieces);
     this.#lastVisitedLine = null;
+    this.#lastVisitedLineLength = null;
   }
 
   delete(offset: number, length: number): void {
@@ -407,6 +445,7 @@ export class PieceTable {
 
     this.#setPieces(nextPieces);
     this.#lastVisitedLine = null;
+    this.#lastVisitedLineLength = null;
   }
 
   applyEdits(edits: readonly ResolvedTextEdit[]): void {
@@ -487,6 +526,7 @@ export class PieceTable {
 
     this.#setPieces(nextPieces);
     this.#lastVisitedLine = null;
+    this.#lastVisitedLineLength = null;
   }
 
   positionAt(offset: number): Position {
@@ -533,22 +573,6 @@ export class PieceTable {
     return offset[0] + character;
   }
 
-  offsetsAt(positions: readonly Position[]): number[] {
-    const offsets: number[] = Array.from({ length: positions.length });
-    if (positions.length === 0) {
-      return offsets;
-    }
-    if (this.#length === 0) {
-      return offsets.fill(0);
-    }
-
-    for (let i = 0; i < positions.length; i++) {
-      offsets[i] = this.offsetAt(positions[i]);
-    }
-
-    return offsets;
-  }
-
   #findPieceAtOffset(
     offset: number
   ): [node: PieceNode, offsetInPiece: number] | undefined {

packages/diffs/src/editor/selection.ts

@@ -134,8 +134,7 @@ export function mapCursorMove(
         const indent = getLeadingSpaces(textDocument.getLineText(line));
         character = character === indent ? 0 : indent;
       } else {
-        character =
-          shortcut === 'start' ? 0 : textDocument.getLineText(line).length;
+        character = shortcut === 'start' ? 0 : textDocument.getLineLength(line);
       }
       if (selection.direction === DirectionBackward) {
         line = selection.start.line;
@@ -147,6 +146,8 @@ export function mapCursorMove(
     } else if (shortcut === 'down') {
       line = Math.min(Math.max(lineCount - 1, 0), line + 1);
     } else if (isCollapsedSelection(selection)) {
+      const lineLength = textDocument.getLineLength(line);
+      character = Math.min(character, lineLength);
       if (shortcut === 'left') {
         character--;
 
@@ -155,12 +156,12 @@ export function mapCursorMove(
             character = 0;
           } else {
             line = Math.max(0, line - 1);
-            character = textDocument.getLineText(line).length;
+            character = textDocument.getLineLength(line);
           }
         }
       } else {
         character++;
-        if (character > textDocument.getLineText(line).length) {
+        if (character > lineLength) {
           if (line === lineCount - 1) {
             character--;
           } else {
@@ -234,7 +235,9 @@ export function applyTextChangeToSelections<LAnnotation>(
   for (const selection of selections) {
     selectionPositions.push(selection.start, selection.end);
   }
-  const selectionOffsets = textDocument.offsetsAt(selectionPositions);
+  const selectionOffsets = selectionPositions.map((position) =>
+    textDocument.offsetAt(position)
+  );
   const primaryStartOffset = selectionOffsets[(selections.length - 1) * 2];
   const primaryEndOffset = selectionOffsets[(selections.length - 1) * 2 + 1];
   const ordered: Array<{
@@ -393,7 +396,9 @@ export function applyTextReplaceToSelections<LAnnotation>(
   for (const selection of selections) {
     selectionPositions.push(selection.start, selection.end);
   }
-  const selectionOffsets = textDocument.offsetsAt(selectionPositions);
+  const selectionOffsets = selectionPositions.map((position) =>
+    textDocument.offsetAt(position)
+  );
   const ordered: Array<{
     index: number;
     start: number;
@@ -554,7 +559,7 @@ export function applyTransposeToSelections<LAnnotation>(
 
     const { line, character } = selection.start;
     const offset = anchor;
-    const lineLength = textDocument.getLineText(line).length;
+    const lineLength = textDocument.getLineLength(line);
     let edit: ResolvedTextEdit | undefined;
 
     if (character > 0 && character < lineLength) {
@@ -573,7 +578,7 @@ export function applyTransposeToSelections<LAnnotation>(
       nextOffsetPairs.push([offset, offset]);
     } else if (character === 0 && line > 0 && lineLength > 0) {
       const prevLine = line - 1;
-      const prevLength = textDocument.getLineText(prevLine).length;
+      const prevLength = textDocument.getLineLength(prevLine);
       const prevEnd = textDocument.offsetAt({
         line: prevLine,
         character: prevLength,
@@ -696,7 +701,7 @@ export function applyDeleteSoftLineBackwardToSelections<LAnnotation>(
         direction: DirectionNone,
       };
     }
-    const prevLineLength = textDocument.getLineText(line - 1).length;
+    const prevLineLength = textDocument.getLineLength(line - 1);
     return {
       start: { line: line - 1, character: prevLineLength },
       end: { line, character: 0 },
@@ -1042,7 +1047,7 @@ export function getDocumentFullSelection(
   textDocument: TextDocument<unknown>
 ): EditorSelection {
   const lastLine = textDocument.lineCount - 1;
-  const lastCharacter = textDocument.getLineText(lastLine)?.length ?? 0;
+  const lastCharacter = textDocument.getLineLength(lastLine);
   return {
     start: { line: 0, character: 0 },
     end: { line: lastLine, character: lastCharacter },
@@ -1058,7 +1063,7 @@ export function getDocumentBoundarySelection(
   atEnd: boolean
 ): EditorSelection {
   const line = atEnd ? textDocument.lineCount - 1 : 0;
-  const character = atEnd ? (textDocument.getLineText(line)?.length ?? 0) : 0;
+  const character = atEnd ? textDocument.getLineLength(line) : 0;
   const start = { line, character };
   return {
     start: start,
@@ -1214,7 +1219,7 @@ function resolveDeleteWordBackwardRange(
     if (line === 0) {
       return [caret, caret];
     }
-    const prevLineLength = textDocument.getLineText(line - 1).length;
+    const prevLineLength = textDocument.getLineLength(line - 1);
     return [
       { line: line - 1, character: prevLineLength },
       { line, character: 0 },

packages/diffs/src/editor/textDocument.ts

@@ -169,11 +169,7 @@ export class TextDocument<LAnnotation> {
   }
 
   offsetAt(position: Position): number {
-    return this.#pieceTable.offsetAt(position);
-  }
-
-  offsetsAt(positions: readonly Position[]): number[] {
-    return this.#pieceTable.offsetsAt(positions);
+    return this.#pieceTable.offsetAt(this.normalizePosition(position));
   }
 
   getText(range?: Range): string {
@@ -184,6 +180,10 @@ export class TextDocument<LAnnotation> {
     return this.#pieceTable.getLineText(line, includeLineBreak);
   }
 
+  getLineLength(line: number, includeLineBreak?: boolean): number {
+    return this.#pieceTable.getLineLength(line, includeLineBreak);
+  }
+
   charAt(offset: number): string;
   charAt(position: Position): string;
   charAt(positionOrOffset: Position | number): string {
@@ -330,7 +330,7 @@ export class TextDocument<LAnnotation> {
       line,
       character: Math.max(
         0,
-        Math.min(position.character, this.getLineText(line).length)
+        Math.min(position.character, this.getLineLength(line))
       ),
     };
   }

packages/diffs/test/editorPieceTable.test.ts

@@ -138,6 +138,19 @@ describe('PieceTable', () => {
     expect(() => table.getLineText(99)).toThrow('Line index out of range: 99');
   });
 
+  test('getLineLength matches getLineText without slicing', () => {
+    const table = new PieceTable('first\r\nsecond\n');
+
+    expect(table.getLineLength(0)).toBe(table.getLineText(0).length);
+    expect(table.getLineLength(1)).toBe(table.getLineText(1).length);
+    expect(table.getLineLength(2)).toBe(0);
+    expect(table.getLineLength(0, true)).toBe(7);
+    expect(table.getLineLength(1, true)).toBe(7);
+    expect(() => table.getLineLength(99)).toThrow(
+      'Line index out of range: 99'
+    );
+  });
+
   test('maps between offsets and positions', () => {
     const table = new PieceTable('ab\nc');
 

packages/diffs/test/editorSelection.test.ts

@@ -1157,6 +1157,69 @@ describe('mapSelectionMove', () => {
       createSelection(0, 4, 0, 4),
     ]);
   });
+
+  test('moves left from a goal column past a shorter line end', () => {
+    const textDocument = new TextDocument(
+      'inmemory://1',
+      'this is a much longer line\nshort\n'
+    );
+    const onShortLine = mapCursorMove(
+      textDocument,
+      [createSelection(0, 20, 0, 20)],
+      'down'
+    );
+
+    expect(onShortLine).toEqual([createSelection(1, 20, 1, 20)]);
+
+    expect(mapCursorMove(textDocument, onShortLine, 'left')).toEqual([
+      createSelection(1, 4, 1, 4),
+    ]);
+  });
+
+  test('preserves goal column across short and empty lines', () => {
+    const textDocument = new TextDocument(
+      'inmemory://1',
+      'this is a much longer line here\nshort\n\nanother much longer line here\n'
+    );
+    const onShortLine = mapCursorMove(
+      textDocument,
+      [createSelection(0, 20, 0, 20)],
+      'down'
+    );
+    const onEmptyLine = mapCursorMove(textDocument, onShortLine, 'down');
+    const onLongLine = mapCursorMove(textDocument, onEmptyLine, 'down');
+
+    expect(onShortLine).toEqual([createSelection(1, 20, 1, 20)]);
+    expect(onEmptyLine).toEqual([createSelection(2, 20, 2, 20)]);
+    expect(onLongLine).toEqual([createSelection(3, 20, 3, 20)]);
+  });
+
+  test('inserts at the clamped caret after moving onto a shorter line', () => {
+    const textDocument = new TextDocument(
+      'inmemory://1',
+      'this is a much longer line\nshort\nnext\n'
+    );
+    const onShortLine = mapCursorMove(
+      textDocument,
+      [createSelection(0, 20, 0, 20)],
+      'down'
+    );
+    const { nextSelections, change } = applyTextChangeToSelections(
+      textDocument,
+      onShortLine,
+      {
+        start: textDocument.offsetAt(onShortLine[0].start),
+        end: textDocument.offsetAt(onShortLine[0].end),
+        text: 'X',
+      }
+    );
+
+    expect(textDocument.getText()).toBe(
+      'this is a much longer line\nshortX\nnext\n'
+    );
+    expect(nextSelections).toEqual([createSelection(1, 6, 1, 6)]);
+    expect(change).toBeDefined();
+  });
 });
 
 describe('mapSelectionRangeMove', () => {

packages/diffs/test/editorTextDocument.test.ts

@@ -84,6 +84,9 @@ describe('TextDocument', () => {
     expect(d.getLineText(0)).toBe('first');
     expect(d.getLineText(1)).toBe('second');
     expect(d.getLineText(2)).toBe('');
+    expect(d.getLineLength(0)).toBe(5);
+    expect(d.getLineLength(1)).toBe(6);
+    expect(d.getLineLength(2)).toBe(0);
     expect(
       d.getText({
         start: { line: 0, character: 0 },