refactor(vnext): share embedded region lookup

Author: Light2Dark

src/vnext/__tests__/source.test.ts

@@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest";
 import {
   createIdentitySqlSource,
   createMaskedSqlSource,
+  findSqlEmbeddedRegionAtOrAfter,
   mapAnalysisRangeToOriginal,
   mapOriginalRangeToAnalysis,
   MAX_SQL_EMBEDDED_REGIONS,
@@ -188,6 +189,20 @@ describe("SQL source snapshots", () => {
     expect(source.analysisText).toBe(" b ");
   });
 
+  it("finds the first embedded region ending after a position", () => {
+    const source = createMaskedSqlSource("abcdef", [
+      { from: 1, language: "python", to: 2 },
+      { from: 3, language: "jinja", to: 5 },
+    ]);
+
+    expect(findSqlEmbeddedRegionAtOrAfter(source, 0)).toBe(0);
+    expect(findSqlEmbeddedRegionAtOrAfter(source, 1)).toBe(0);
+    expect(findSqlEmbeddedRegionAtOrAfter(source, 2)).toBe(1);
+    expect(findSqlEmbeddedRegionAtOrAfter(source, 4)).toBe(1);
+    expect(findSqlEmbeddedRegionAtOrAfter(source, 5)).toBe(2);
+    expect(findSqlEmbeddedRegionAtOrAfter(source, 6)).toBe(2);
+  });
+
   it.each([
     null,
     {},

src/vnext/bounded-sql-lexer.ts

@@ -9,7 +9,10 @@ import {
   sqlIdentifierStartLengthAt,
   type SqlLexicalProfile,
 } from "./lexical.js";
-import type { SqlSourceSnapshot } from "./source.js";
+import {
+  findSqlEmbeddedRegionAtOrAfter,
+  type SqlSourceSnapshot,
+} from "./source.js";
 
 export const MAX_BOUNDED_SQL_LEXEMES = 16_384;
 
@@ -34,24 +37,6 @@ export type BoundedSqlLexerResource =
   | "dollar-quote-delimiter"
   | "lexical-token";
 
-function findSqlRegionAtOrAfter(
-  source: SqlSourceSnapshot,
-  position: number,
-): number {
-  let low = 0;
-  let high = source.embeddedRegions.length;
-  while (low < high) {
-    const middle = low + Math.floor((high - low) / 2);
-    const region = source.embeddedRegions[middle];
-    if (!region || region.to <= position) {
-      low = middle + 1;
-    } else {
-      high = middle;
-    }
-  }
-  return low;
-}
-
 export class BoundedSqlLexer {
   readonly #profile: SqlLexicalProfile;
   readonly #source: SqlSourceSnapshot;
@@ -72,7 +57,7 @@ export class BoundedSqlLexer {
     this.#cursor = from;
     this.#to = to;
     this.#profile = profile;
-    this.#regionIndex = findSqlRegionAtOrAfter(source, from);
+    this.#regionIndex = findSqlEmbeddedRegionAtOrAfter(source, from);
   }
 
   next(): BoundedSqlLexeme | null {

src/vnext/query-site.ts

@@ -5,7 +5,10 @@ import {
   type BoundedSqlLexerResource,
 } from "./bounded-sql-lexer.js";
 import type { SqlLexicalProfile } from "./lexical.js";
-import type { SqlSourceSnapshot } from "./source.js";
+import {
+  findSqlEmbeddedRegionAtOrAfter,
+  type SqlSourceSnapshot,
+} from "./source.js";
 import type {
   ExactSqlStatementSlot,
   SqlStatementSlot,
@@ -200,30 +203,14 @@ function createRange(from: number, to: number): SqlQuerySiteRange {
   return Object.freeze(range);
 }
 
-function findRegionAtOrAfter(
-  source: SqlSourceSnapshot,
-  position: number,
-): number {
-  let low = 0;
-  let high = source.embeddedRegions.length;
-  while (low < high) {
-    const middle = low + Math.floor((high - low) / 2);
-    const region = source.embeddedRegions[middle];
-    if (!region || region.to <= position) {
-      low = middle + 1;
-    } else {
-      high = middle;
-    }
-  }
-  return low;
-}
-
 function regionContains(
   source: SqlSourceSnapshot,
   position: number,
 ): boolean {
   const region =
-    source.embeddedRegions[findRegionAtOrAfter(source, position)];
+    source.embeddedRegions[
+      findSqlEmbeddedRegionAtOrAfter(source, position)
+    ];
   return Boolean(
     region && region.from <= position && position < region.to,
   );
@@ -672,7 +659,10 @@ function intersectsRegion(
   from: number,
   to: number,
 ): boolean {
-  const region = source.embeddedRegions[findRegionAtOrAfter(source, from)];
+  const region =
+    source.embeddedRegions[
+      findSqlEmbeddedRegionAtOrAfter(source, from)
+    ];
   return Boolean(region && region.from < to && from < region.to);
 }
 

src/vnext/source.ts

@@ -38,6 +38,24 @@ export interface SqlSourceSnapshot {
   readonly originalText: string;
 }
 
+export function findSqlEmbeddedRegionAtOrAfter(
+  source: SqlSourceSnapshot,
+  position: number,
+): number {
+  let low = 0;
+  let high = source.embeddedRegions.length;
+  while (low < high) {
+    const middle = low + Math.floor((high - low) / 2);
+    const region = source.embeddedRegions[middle];
+    if (!region || region.to <= position) {
+      low = middle + 1;
+    } else {
+      high = middle;
+    }
+  }
+  return low;
+}
+
 interface MissingDataProperty {
   readonly found: false;
 }