feat(editor): preserve caret when formatting SQL

Closes #4632

apps/desktop/src/components/editor/QueryEditor.vue

@@ -75,6 +75,7 @@ import { restoreSqlFromSourcePaste } from "@/lib/sql/sqlSourcePaste";
 import { enabledSqlParameterSyntaxes, resolveSqlVariableSyntaxToggles } from "@/lib/sql/sqlVariableSyntax";
 
 import { createQueryEditorExecutionViewportOwnership, isQueryEditorPositionVisible } from "@/lib/editor/queryEditorExecutionViewport";
+import { mapQueryEditorFormatSelection } from "@/lib/editor/queryEditorFormatSelection";
 import { joinQueryEditorLines } from "@/lib/editor/queryEditorJoinLines";
 
 import { resolveSqlSingleQuoteKeyAction } from "@/lib/sql/sqlQuoteCaret";
@@ -1670,9 +1671,11 @@ async function formatCurrentSql() {
   const to = formatsSelection ? selection.to : originalState.doc.length;
   const source = originalState.sliceDoc(from, to);
   if (!source.trim()) return;
+  const formatDialect = props.formatDialect ?? props.dialect ?? "generic";
 
   try {
     let formatted: string;
+    let formattedAsSql = false;
     if (props.databaseType === "mongodb") {
       formatted = formatMongoShellText(source, settingsStore.editorSettings.sqlFormatter);
     } else {
@@ -1695,7 +1698,8 @@ async function formatCurrentSql() {
           toast(t("toolbar.formatAutoDetectFailed"), 3000);
           return;
         } else {
-          formatted = await formatSqlForEditing(source, props.formatDialect ?? props.dialect ?? "generic", settingsStore.editorSettings.sqlFormatter);
+          formattedAsSql = true;
+          formatted = await formatSqlForEditing(source, formatDialect, settingsStore.editorSettings.sqlFormatter);
         }
       }
     }
@@ -1703,9 +1707,13 @@ async function formatCurrentSql() {
       return;
     }
     if (formatted === source) return;
+    // The semantic mapper's contract is SQL-only. Mongo shell, Elasticsearch,
+    // JSON, and XML keep the replacement selection behavior used before the
+    // SQL caret-preservation fix.
+    const replacementSelection = formattedAsSql ? mapQueryEditorFormatSelection(source, formatted, { anchor: selection.anchor - from, head: selection.head - from }, formatDialect) : formatsSelection ? { anchor: 0, head: formatted.length } : { anchor: formatted.length, head: formatted.length };
     currentView.dispatch({
       changes: { from, to, insert: formatted },
-      selection: formatsSelection ? { anchor: from, head: from + formatted.length } : { anchor: from + formatted.length },
+      selection: { anchor: from + replacementSelection.anchor, head: from + replacementSelection.head },
     });
   } catch (e: any) {
     emit("formatError", String(e?.message || e));

apps/desktop/src/components/editor/__tests__/QueryEditor.formatCompressRequestReplay.spec.ts

@@ -1,6 +1,8 @@
 // @vitest-environment happy-dom
 
 import { createApp, defineComponent, h, nextTick, reactive } from "vue";
+import { undo } from "@codemirror/commands";
+import { EditorView } from "@codemirror/view";
 import { createPinia, setActivePinia } from "pinia";
 import { createI18n } from "vue-i18n";
 import { afterEach, describe, expect, it, vi } from "vitest";
@@ -8,6 +10,7 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import QueryEditor from "@/components/editor/QueryEditor.vue";
 import { compressSqlText, formatSqlForEditing } from "@/lib/sql/sqlFormatter";
 import { DEFAULT_SQL_FORMATTER_SETTINGS } from "@/lib/sql/sqlFormatterConfig";
+import type { DatabaseType } from "@/types/database";
 
 // The request ids are per-kind global counters in App.vue, so ids keep
 // increasing for the lifetime of the process. The module-scope replay cursor in
@@ -33,7 +36,7 @@ type EditorState = {
 
 const WAIT = { timeout: 5000, interval: 20 };
 
-function mountEditor(initial: { modelValue: string; tabId: string; formatRequestId?: number; compressRequestId?: number }) {
+function mountEditor(initial: { modelValue: string; tabId: string; databaseType?: DatabaseType; formatRequestId?: number; compressRequestId?: number }) {
   const pinia = createPinia();
   setActivePinia(pinia);
   const state = reactive<EditorState>({
@@ -48,7 +51,7 @@ function mountEditor(initial: { modelValue: string; tabId: string; formatRequest
       h(QueryEditor, {
         modelValue: state.modelValue,
         tabId: state.tabId,
-        databaseType: "mysql",
+        databaseType: initial.databaseType ?? "mysql",
         dialect: "mysql",
         formatDialect: "mysql",
         autoFocus: false,
@@ -158,4 +161,111 @@ describe("QueryEditor format/compress request replay", () => {
 
     expect(second.emits).toEqual([]);
   });
+
+  it("keeps the whole-document format caret on its original SQL token", async () => {
+    const { state, emits, host } = mountEditor({ modelValue: COMPRESSED_SQL, tabId: "tab-a" });
+    await waitForEditor(host);
+    const editorView = EditorView.findFromDOM(host.querySelector(".cm-editor") as HTMLElement)!;
+    const caret = COMPRESSED_SQL.lastIndexOf("tenant_id") + 4;
+    editorView.dispatch({ selection: { anchor: caret } });
+
+    state.formatRequestId = 4;
+    await vi.waitFor(() => expect(emits.at(-1)).toBe(FORMATTED_SQL), WAIT);
+
+    expect(editorView.state.selection.main.head).toBe(FORMATTED_SQL.lastIndexOf("tenant_id") + 4);
+    expect(editorView.state.selection.main.head).toBeLessThan(editorView.state.doc.length);
+    expect(undo(editorView)).toBe(true);
+    expect(editorView.state.doc.toString()).toBe(COMPRESSED_SQL);
+    expect(editorView.state.selection.main.head).toBe(caret);
+    expect(undo(editorView)).toBe(false);
+  });
+
+  it("does not move the caret or dispatch a change when a parse failure returns the source", async () => {
+    const source = "SELECT 1 .";
+    const { state, emits, host } = mountEditor({ modelValue: source, tabId: "tab-a" });
+    await waitForEditor(host);
+    const editorView = EditorView.findFromDOM(host.querySelector(".cm-editor") as HTMLElement)!;
+    const caret = source.indexOf("1");
+    editorView.dispatch({ selection: { anchor: caret } });
+
+    state.formatRequestId = 5;
+    await settle();
+    await settle();
+
+    expect(emits).toEqual([]);
+    expect(editorView.state.selection.main.head).toBe(caret);
+  });
+
+  it.each([
+    ["forward", 0, COMPRESSED_SQL.length, 0, FORMATTED_SQL.length, 6],
+    ["reverse", COMPRESSED_SQL.length, 0, FORMATTED_SQL.length, 0, 7],
+  ] as const)("keeps a %s formatted SQL selection complete and in its original direction", async (_direction, anchor, head, expectedAnchor, expectedHead, requestId) => {
+    const { state, emits, host } = mountEditor({ modelValue: COMPRESSED_SQL, tabId: "tab-a" });
+    await waitForEditor(host);
+    const editorView = EditorView.findFromDOM(host.querySelector(".cm-editor") as HTMLElement)!;
+    editorView.dispatch({ selection: { anchor, head } });
+
+    state.formatRequestId = requestId;
+    await vi.waitFor(() => expect(emits.at(-1)).toBe(FORMATTED_SQL), WAIT);
+
+    expect(editorView.state.selection.main.anchor).toBe(expectedAnchor);
+    expect(editorView.state.selection.main.head).toBe(expectedHead);
+  });
+
+  it.each([
+    {
+      name: "Mongo shell",
+      databaseType: "mongodb",
+      source: 'db.items.find({name:"Ada",age:42});',
+      formatted: 'db.items.find({\n  name: "Ada",\n  age: 42\n});',
+      requestId: 8,
+    },
+    {
+      name: "Elasticsearch request",
+      databaseType: "elasticsearch",
+      source: 'GET /items/_search {"query":{"match_all":{}}}',
+      formatted: 'GET /items/_search\n{\n  "query": {\n    "match_all": {}\n  }\n}',
+      requestId: 9,
+    },
+    {
+      name: "JSON",
+      databaseType: "mysql",
+      source: '{"name":"Ada","items":[1,2]}',
+      formatted: '{\n  "name": "Ada",\n  "items": [\n    1,\n    2\n  ]\n}',
+      requestId: 10,
+    },
+    {
+      name: "XML",
+      databaseType: "mysql",
+      source: '<root><item id="1"/><item id="2"/></root>',
+      formatted: '<root>\n  <item id="1"/>\n  <item id="2"/>\n</root>',
+      requestId: 11,
+    },
+  ] as const)("keeps the existing whole-document caret behavior for $name formatting", async ({ databaseType, source, formatted, requestId }) => {
+    const { state, emits, host } = mountEditor({ modelValue: source, tabId: "tab-a", databaseType });
+    await waitForEditor(host);
+    const editorView = EditorView.findFromDOM(host.querySelector(".cm-editor") as HTMLElement)!;
+    editorView.dispatch({ selection: { anchor: Math.floor(source.length / 2) } });
+
+    state.formatRequestId = requestId;
+    await vi.waitFor(() => expect(emits.at(-1)).toBe(formatted), WAIT);
+
+    expect(editorView.state.selection.main.anchor).toBe(formatted.length);
+    expect(editorView.state.selection.main.head).toBe(formatted.length);
+  });
+
+  it("keeps the existing forward full-range selection for non-SQL formatting", async () => {
+    const source = '{"name":"Ada","items":[1,2]}';
+    const formatted = '{\n  "name": "Ada",\n  "items": [\n    1,\n    2\n  ]\n}';
+    const { state, emits, host } = mountEditor({ modelValue: source, tabId: "tab-a" });
+    await waitForEditor(host);
+    const editorView = EditorView.findFromDOM(host.querySelector(".cm-editor") as HTMLElement)!;
+    editorView.dispatch({ selection: { anchor: source.length, head: 0 } });
+
+    state.formatRequestId = 12;
+    await vi.waitFor(() => expect(emits.at(-1)).toBe(formatted), WAIT);
+
+    expect(editorView.state.selection.main.anchor).toBe(0);
+    expect(editorView.state.selection.main.head).toBe(formatted.length);
+  });
 });

apps/desktop/src/lib/__tests__/editor/queryEditorFormatSelection.spec.ts

@@ -0,0 +1,110 @@
+import { describe, expect, it } from "vitest";
+import { mapQueryEditorFormatSelection } from "@/lib/editor/queryEditorFormatSelection";
+import { formatSqlForEditing } from "@/lib/sql/sqlFormatter";
+
+function occurrence(text: string, needle: string, ordinal: number): number {
+  let index = -1;
+  for (let current = 0; current <= ordinal; current += 1) {
+    index = text.indexOf(needle, index + 1);
+    if (index < 0) throw new Error(`Missing occurrence ${ordinal} of ${needle}`);
+  }
+  return index;
+}
+
+describe("mapQueryEditorFormatSelection", () => {
+  it("keeps the caret in the same keyword and token after formatting inserts earlier newlines", async () => {
+    const source = "select id, name, email, status, created_at from users where status = 'active' and tenant_id = 42";
+    const formatted = await formatSqlForEditing(source, "mysql");
+    const keywordCaret = source.indexOf("where") + 2;
+    const tokenCaret = source.indexOf("status", source.indexOf("where")) + 3;
+
+    expect(formatted.slice(0, formatted.indexOf("WHERE"))).toContain("\n");
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: keywordCaret, head: keywordCaret }, "mysql")).toEqual({
+      anchor: formatted.indexOf("WHERE") + 2,
+      head: formatted.indexOf("WHERE") + 2,
+    });
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: tokenCaret, head: tokenCaret }, "mysql")).toEqual({
+      anchor: formatted.indexOf("status", formatted.indexOf("WHERE")) + 3,
+      head: formatted.indexOf("status", formatted.indexOf("WHERE")) + 3,
+    });
+  });
+
+  it("maps repeated tokens by their stable ordinal instead of searching for the first text match", async () => {
+    const source = "select status, status from users where status = 'status' order by status;";
+    const formatted = await formatSqlForEditing(source, "mysql");
+    const caret = occurrence(source, "status", 2) + 4;
+    const expected = occurrence(formatted, "status", 2) + 4;
+
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: caret, head: caret }, "mysql")).toEqual({ anchor: expected, head: expected });
+  });
+
+  it.each([
+    ["string", "'-- repeated text'", 0, 5],
+    ["block comment", "/* repeated text */", 0, 6],
+    ["line comment", "-- repeated text", 1, 4],
+    ["trailing string", "'repeated text'", 0, 4],
+  ] as const)("maps a caret inside a %s token without confusing its contents for SQL", (_name, token, ordinal, offset) => {
+    const source = "select '-- repeated text' as note, value /* repeated text */ from logs -- repeated text\nwhere value = 'repeated text';";
+    const formatted = "SELECT\n  '-- repeated text' AS note,\n  value /* repeated text */\nFROM logs -- repeated text\nWHERE value = 'repeated text';";
+    const caret = occurrence(source, token, ordinal) + offset;
+    const expected = occurrence(formatted, token, ordinal) + offset;
+
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: caret, head: caret }, "mysql")).toEqual({ anchor: expected, head: expected });
+  });
+
+  it("maps a caret proportionally within changed whitespace between stable tokens", () => {
+    const source = "SELECT id      FROM users";
+    const formatted = "SELECT id\nFROM users";
+    const caret = source.indexOf("id") + 5;
+    const mapped = mapQueryEditorFormatSelection(source, formatted, { anchor: caret, head: caret }, "mysql").head;
+
+    expect(mapped).toBeGreaterThanOrEqual(formatted.indexOf("id") + 2);
+    expect(mapped).toBeLessThanOrEqual(formatted.indexOf("FROM"));
+  });
+
+  it("keeps unchanged formatter output and its selection untouched", () => {
+    const source = "SELECT id FROM users;";
+    expect(mapQueryEditorFormatSelection(source, source, { anchor: 18, head: 7 }, "mysql")).toEqual({ anchor: 18, head: 7 });
+  });
+
+  it("preserves the selected formatted range and its direction", async () => {
+    const source = "select id, name from users where active = 1";
+    const formatted = await formatSqlForEditing(source, "mysql");
+
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: source.length, head: 0 }, "mysql")).toEqual({
+      anchor: formatted.length,
+      head: 0,
+    });
+  });
+
+  it("stays in the corresponding statement when changed tokens require the relative fallback", () => {
+    const source = "select one; select old_value from old_table;";
+    const formatted = "SELECT one;\n\nEXECUTE completely_different_expression;";
+    const caret = source.indexOf("old_value") + 4;
+    const mapped = mapQueryEditorFormatSelection(source, formatted, { anchor: caret, head: caret }, "mysql").head;
+    const secondStatement = formatted.indexOf("EXECUTE");
+
+    expect(mapped).toBeGreaterThanOrEqual(secondStatement);
+    expect(mapped).toBeLessThan(formatted.lastIndexOf(";"));
+  });
+
+  it("uses the formatted document end when the corresponding statement is unavailable", () => {
+    const source = "select one; select two;";
+    const formatted = "SELECT one;";
+    const caret = source.indexOf("two") + 1;
+
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: caret, head: caret }, "mysql")).toEqual({
+      anchor: formatted.length,
+      head: formatted.length,
+    });
+  });
+
+  it("leaves the caret untouched when a parse failure makes the formatter return the source", async () => {
+    const source = "SELECT 1 .";
+    const formatted = await formatSqlForEditing(source, "generic");
+    const caret = source.indexOf("1");
+
+    expect(formatted).toBe(source);
+    expect(mapQueryEditorFormatSelection(source, formatted, { anchor: caret, head: caret }, "generic")).toEqual({ anchor: caret, head: caret });
+  });
+});

apps/desktop/src/lib/editor/queryEditorFormatSelection.ts

@@ -0,0 +1,231 @@
+import { tokenizeSqlSemantic } from "@/lib/sql/semantic/tokens";
+import type { SqlSemanticToken } from "@/lib/sql/semantic/types";
+
+export interface QueryEditorFormatSelection {
+  anchor: number;
+  head: number;
+}
+
+interface StatementSpan {
+  from: number;
+  to: number;
+  tokens: SqlSemanticToken[];
+}
+
+interface StatementMapping {
+  source: StatementSpan;
+  formatted: StatementSpan;
+  tokenMatches: Map<number, number>;
+}
+
+function clampPosition(position: number, length: number): number {
+  return Math.max(0, Math.min(length, position));
+}
+
+function tokenKey(value: SqlSemanticToken): string {
+  return `${value.kind}\0${value.normalized}`;
+}
+
+function tokenizeForFormatting(text: string, dialectId: string): SqlSemanticToken[] {
+  const mysqlFamily = dialectId === "mysql" || dialectId === "doris";
+  return tokenizeSqlSemantic(text, dialectId, {
+    mysqlBackslashEscape: mysqlFamily,
+    mysqlDashCommentRequiresWhitespace: mysqlFamily,
+  });
+}
+
+function statementSpans(text: string, tokens: readonly SqlSemanticToken[]): StatementSpan[] {
+  const statements: StatementSpan[] = [];
+  let from = 0;
+  let tokenFrom = 0;
+
+  for (let index = 0; index < tokens.length; index += 1) {
+    const current = tokens[index]!;
+    if (current.kind !== "punctuation" || current.text !== ";" || current.depth !== 0) continue;
+    statements.push({ from, to: current.span.end, tokens: tokens.slice(tokenFrom, index + 1) });
+    from = current.span.end;
+    tokenFrom = index + 1;
+  }
+
+  statements.push({ from, to: text.length, tokens: tokens.slice(tokenFrom) });
+  return statements;
+}
+
+function lowerBound(values: readonly number[], target: number): number {
+  let low = 0;
+  let high = values.length;
+  while (low < high) {
+    const middle = (low + high) >>> 1;
+    if (values[middle]! < target) low = middle + 1;
+    else high = middle;
+  }
+  return low;
+}
+
+/**
+ * Keeps only a monotonic chain of unique token matches. This is a small
+ * patience-diff style fallback: formatting never intentionally reorders SQL
+ * tokens, so a crossing match is less trustworthy than leaving that region to
+ * the statement-relative fallback.
+ */
+function monotonicTokenMatches(candidates: ReadonlyArray<readonly [number, number]>): Map<number, number> {
+  if (candidates.length === 0) return new Map();
+
+  const tails: number[] = [];
+  const tailCandidates: number[] = [];
+  const previous = Array.from({ length: candidates.length }, () => -1);
+
+  for (let index = 0; index < candidates.length; index += 1) {
+    const formattedIndex = candidates[index]![1];
+    const insertion = lowerBound(tails, formattedIndex);
+    tails[insertion] = formattedIndex;
+    previous[index] = insertion > 0 ? tailCandidates[insertion - 1]! : -1;
+    tailCandidates[insertion] = index;
+  }
+
+  const matches = new Map<number, number>();
+  let candidateIndex = tailCandidates[tails.length - 1] ?? -1;
+  while (candidateIndex >= 0) {
+    const [sourceIndex, formattedIndex] = candidates[candidateIndex]!;
+    matches.set(sourceIndex, formattedIndex);
+    candidateIndex = previous[candidateIndex]!;
+  }
+  return matches;
+}
+
+function buildTokenMatches(source: readonly SqlSemanticToken[], formatted: readonly SqlSemanticToken[]): Map<number, number> {
+  if (source.length === formatted.length && source.every((value, index) => tokenKey(value) === tokenKey(formatted[index]!))) {
+    return new Map(source.map((_, index) => [index, index]));
+  }
+
+  // When the streams differ, repeated text is ambiguous. Only unique token
+  // values are considered as anchors, then crossing anchors are discarded.
+  // Ordinary formatting takes the exact-stream path above, so repeated SQL
+  // identifiers still map by their stable token ordinal.
+  const sourceIndexes = new Map<string, number[]>();
+  const formattedIndexes = new Map<string, number[]>();
+  for (let index = 0; index < source.length; index += 1) {
+    const key = tokenKey(source[index]!);
+    const indexes = sourceIndexes.get(key);
+    if (indexes) indexes.push(index);
+    else sourceIndexes.set(key, [index]);
+  }
+  for (let index = 0; index < formatted.length; index += 1) {
+    const key = tokenKey(formatted[index]!);
+    const indexes = formattedIndexes.get(key);
+    if (indexes) indexes.push(index);
+    else formattedIndexes.set(key, [index]);
+  }
+
+  const candidates: Array<readonly [number, number]> = [];
+  for (let sourceIndex = 0; sourceIndex < source.length; sourceIndex += 1) {
+    const sourceOccurrences = sourceIndexes.get(tokenKey(source[sourceIndex]!));
+    const formattedOccurrences = formattedIndexes.get(tokenKey(source[sourceIndex]!));
+    if (sourceOccurrences?.length === 1 && formattedOccurrences?.length === 1) {
+      candidates.push([sourceIndex, formattedOccurrences[0]!]);
+    }
+  }
+  return monotonicTokenMatches(candidates);
+}
+
+function scalePosition(position: number, sourceFrom: number, sourceTo: number, formattedFrom: number, formattedTo: number): number {
+  if (sourceTo <= sourceFrom) return formattedFrom;
+  const boundedPosition = Math.max(sourceFrom, Math.min(sourceTo, position));
+  const ratio = (boundedPosition - sourceFrom) / (sourceTo - sourceFrom);
+  return Math.max(formattedFrom, Math.min(formattedTo, Math.round(formattedFrom + ratio * (formattedTo - formattedFrom))));
+}
+
+function positionInMappedToken(position: number, source: SqlSemanticToken, formatted: SqlSemanticToken): number {
+  return scalePosition(position, source.span.start, source.span.end, formatted.span.start, formatted.span.end);
+}
+
+function directTokenAtPosition(tokens: readonly SqlSemanticToken[], position: number): number | undefined {
+  const interior = tokens.findIndex((value) => value.span.start < position && position < value.span.end);
+  if (interior >= 0) return interior;
+
+  // At an adjacent-token boundary, keep the caret after the preceding token.
+  // This matches how a user normally reaches that position by typing.
+  for (let index = tokens.length - 1; index >= 0; index -= 1) {
+    if (tokens[index]!.span.end === position) return index;
+  }
+  const starting = tokens.findIndex((value) => value.span.start === position);
+  return starting >= 0 ? starting : undefined;
+}
+
+function mapWithinStatement(position: number, mapping: StatementMapping): number {
+  const directSourceIndex = directTokenAtPosition(mapping.source.tokens, position);
+  if (directSourceIndex !== undefined) {
+    const formattedIndex = mapping.tokenMatches.get(directSourceIndex);
+    if (formattedIndex !== undefined) {
+      return positionInMappedToken(position, mapping.source.tokens[directSourceIndex]!, mapping.formatted.tokens[formattedIndex]!);
+    }
+  }
+
+  let left: { source: number; formatted: number } | undefined;
+  let right: { source: number; formatted: number } | undefined;
+  for (let sourceIndex = 0; sourceIndex < mapping.source.tokens.length; sourceIndex += 1) {
+    const formattedIndex = mapping.tokenMatches.get(sourceIndex);
+    if (formattedIndex === undefined) continue;
+    const sourceToken = mapping.source.tokens[sourceIndex]!;
+    const formattedToken = mapping.formatted.tokens[formattedIndex]!;
+    if (sourceToken.span.end <= position) left = { source: sourceToken.span.end, formatted: formattedToken.span.end };
+    if (!right && sourceToken.span.start >= position) right = { source: sourceToken.span.start, formatted: formattedToken.span.start };
+  }
+
+  const sourceFrom = left?.source ?? mapping.source.from;
+  const sourceTo = right?.source ?? mapping.source.to;
+  const formattedFrom = left?.formatted ?? mapping.formatted.from;
+  const formattedTo = right?.formatted ?? mapping.formatted.to;
+  if (sourceFrom <= position && position <= sourceTo && formattedFrom <= formattedTo) {
+    return scalePosition(position, sourceFrom, sourceTo, formattedFrom, formattedTo);
+  }
+
+  return scalePosition(position, mapping.source.from, mapping.source.to, mapping.formatted.from, mapping.formatted.to);
+}
+
+function statementIndexAtPosition(statements: readonly StatementSpan[], position: number): number {
+  const index = statements.findIndex((statement) => position <= statement.to);
+  return index >= 0 ? index : statements.length - 1;
+}
+
+function mapPosition(position: number, source: string, formatted: string, sourceStatements: readonly StatementSpan[], mappings: ReadonlyArray<StatementMapping | undefined>): number {
+  const clamped = clampPosition(position, source.length);
+  if (clamped === 0) return 0;
+  if (clamped === source.length) return formatted.length;
+
+  const statementIndex = statementIndexAtPosition(sourceStatements, clamped);
+  const mapping = mappings[statementIndex];
+  return mapping ? mapWithinStatement(clamped, mapping) : formatted.length;
+}
+
+/**
+ * Maps a caret or selection through SQL formatting without diffing every
+ * character. Token-equivalent statements map by token ordinal; changed token
+ * streams use unique monotonic anchors and finally a relative position inside
+ * the same semicolon-delimited statement.
+ */
+export function mapQueryEditorFormatSelection(source: string, formatted: string, selection: QueryEditorFormatSelection, dialectId = "generic"): QueryEditorFormatSelection {
+  if (source === formatted) {
+    return {
+      anchor: clampPosition(selection.anchor, source.length),
+      head: clampPosition(selection.head, source.length),
+    };
+  }
+
+  const sourceStatements = statementSpans(source, tokenizeForFormatting(source, dialectId));
+  const formattedStatements = statementSpans(formatted, tokenizeForFormatting(formatted, dialectId));
+  const mappings = sourceStatements.map((sourceStatement, index): StatementMapping | undefined => {
+    const formattedStatement = formattedStatements[index];
+    if (!formattedStatement) return undefined;
+    return {
+      source: sourceStatement,
+      formatted: formattedStatement,
+      tokenMatches: buildTokenMatches(sourceStatement.tokens, formattedStatement.tokens),
+    };
+  });
+
+  return {
+    anchor: mapPosition(selection.anchor, source, formatted, sourceStatements, mappings),
+    head: mapPosition(selection.head, source, formatted, sourceStatements, mappings),
+  };
+}