fix: correct lazy lookup optimization

BartWaardenburg · BartWaardenburg · commit 1ab1c17c59b9 · 2026-07-13T13:55:01.000+02:00
diff --git a/benchmarks/real-world.mjs b/benchmarks/real-world.mjs
@@ -1,5 +1,6 @@
 import { createBench, latencyMeanMs, latencyP99Ms, throughputHz } from "./codspeed.mjs";
 import { createDeterministicLookups, setFailureExitCode } from "./workload.mjs";
+import assert from "node:assert/strict";
 import { readFileSync, existsSync } from "node:fs";
 import { dirname, join } from "node:path";
 import { fileURLToPath } from "node:url";
@@ -27,6 +28,49 @@ const createOrderedLookups = (count, maxLine, descending) =>
     };
   });
 
+const createLookupPatterns = (maxLine) => {
+  const midLine = Math.floor(maxLine / 2);
+  return [
+    { name: "ascending", lookups: createOrderedLookups(LOOKUP_COUNT, maxLine, false) },
+    { name: "descending", lookups: createOrderedLookups(LOOKUP_COUNT, maxLine, true) },
+    {
+      name: "repeated",
+      lookups: Array.from({ length: LOOKUP_COUNT }, () => ({ line: midLine, column: 20 })),
+    },
+    {
+      name: "randomized",
+      lookups: createDeterministicLookups(LOOKUP_COUNT, maxLine, LOOKUP_MAX_COLUMN, LOOKUP_SEED),
+    },
+  ];
+};
+
+const assertLazyMatchesEager = (eager, lazy, lookups, context) => {
+  for (const { line, column } of lookups) {
+    assert.deepEqual(
+      lazy.originalPositionFor(line, column),
+      eager.originalPositionFor(line, column),
+      `${context} at ${line}:${column}`,
+    );
+  }
+};
+
+const consumeLookups = (map, lookups) => {
+  let checksum = 2_166_136_261;
+
+  for (const { line, column } of lookups) {
+    const result = map.originalPositionFor(line, column);
+    const value =
+      result === null
+        ? 0
+        : result.line ^ result.column ^ (result.source?.length ?? 0) ^ (result.name?.length ?? 0);
+    checksum = Math.imul(checksum ^ value, 16_777_619) >>> 0;
+  }
+
+  return checksum;
+};
+
+let lookupChecksum = 0;
+
 // ── Load fixtures ────────────────────────────────────────────────
 
 const FIXTURES = [
@@ -181,39 +225,78 @@ for (const { name, json, size } of maps) {
   );
 }
 
-// ── Fast-lazy lookup order ───────────────────────────────────────
+// ── Fast-lazy cold map lookup order ─────────────────────────────
 
-console.log("\n--- Fast-Lazy Lookup Order ---\n");
+console.log("\n--- Fast-Lazy Cold Map: Construct and First Lookup Pass ---\n");
 
 for (const { name, json, size, lines } of maps) {
   console.log(`### ${name}\n`);
 
-  const midLine = Math.floor(lines / 2);
-  const patterns = [
-    { name: "ascending", lookups: createOrderedLookups(LOOKUP_COUNT, lines, false) },
-    { name: "descending", lookups: createOrderedLookups(LOOKUP_COUNT, lines, true) },
-    {
-      name: "repeated",
-      lookups: Array.from({ length: LOOKUP_COUNT }, () => ({ line: midLine, column: 20 })),
-    },
-    {
-      name: "randomized",
-      lookups: createDeterministicLookups(LOOKUP_COUNT, lines, LOOKUP_MAX_COLUMN, LOOKUP_SEED),
-    },
-  ].map((pattern) => ({ ...pattern, map: new FastSourceMap(json) }));
+  const eager = new SourceMap(json);
+  const patterns = createLookupPatterns(lines);
+
+  for (const pattern of patterns) {
+    const lazy = new FastSourceMap(json);
+    assertLazyMatchesEager(eager, lazy, pattern.lookups, `${name} ${pattern.name}`);
+    lazy.free();
+  }
+  eager.free();
+
+  const isLargeMap = size > 1024 * 1024;
+  const bench = createBench({
+    warmupIterations: isLargeMap ? 2 : 5,
+    iterations: isLargeMap ? 10 : 50,
+  });
+  const prefix = `real_world_lazy_lookup_cold_1000x[${name}]`;
+
+  for (const pattern of patterns) {
+    bench.add(`${prefix} ${pattern.name}`, () => {
+      const lazy = new FastSourceMap(json);
+      lookupChecksum = (lookupChecksum + consumeLookups(lazy, pattern.lookups)) >>> 0;
+      lazy.free();
+    });
+  }
+
+  await bench.run();
+
+  console.table(
+    bench.tasks.map((task) => ({
+      Name: task.name,
+      "ops/sec": Math.round(throughputHz(task)).toLocaleString(),
+      "avg (μs)": (latencyMeanMs(task) * 1000).toFixed(1),
+      "per lookup (ns)": Math.round(
+        (latencyMeanMs(task) * 1_000_000) / LOOKUP_COUNT,
+      ).toLocaleString(),
+    })),
+  );
+}
+
+// ── Fast-lazy warm cache lookup order ────────────────────────────
+
+console.log("\n--- Fast-Lazy Warm Cache: Reuse Decoded Lines ---\n");
+
+for (const { name, json, size, lines } of maps) {
+  console.log(`### ${name}\n`);
+
+  const patterns = createLookupPatterns(lines).map((pattern) => ({
+    ...pattern,
+    map: new FastSourceMap(json),
+  }));
+
+  for (const pattern of patterns) {
+    lookupChecksum = (lookupChecksum + consumeLookups(pattern.map, pattern.lookups)) >>> 0;
+  }
 
   const isLargeMap = size > 1024 * 1024;
   const bench = createBench({
     warmupIterations: isLargeMap ? 5 : 20,
     iterations: isLargeMap ? 50 : 200,
   });
-  const prefix = `real_world_lazy_lookup_1000x[${name}]`;
+  const prefix = `real_world_lazy_lookup_warm_1000x[${name}]`;
 
   for (const pattern of patterns) {
     bench.add(`${prefix} ${pattern.name}`, () => {
-      for (const { line, column } of pattern.lookups) {
-        pattern.map.originalPositionFor(line, column);
-      }
+      lookupChecksum = (lookupChecksum + consumeLookups(pattern.map, pattern.lookups)) >>> 0;
     });
   }
 
@@ -235,6 +318,8 @@ for (const { name, json, size, lines } of maps) {
   }
 }
 
+console.log(`Lookup checksum: ${lookupChecksum}`);
+
 // ── Single lookup ────────────────────────────────────────────────
 
 console.log("\n--- Single Lookup ---\n");
diff --git a/crates/sourcemap/benches/parse.rs b/crates/sourcemap/benches/parse.rs
@@ -642,6 +642,13 @@ fn randomized_lazy_lookup_input() -> LazyLookupInput {
     lazy_lookup_input(lookups)
 }
 
+fn warm_lazy_lookup_input(input: LazyLookupInput) -> LazyLookupInput {
+    for &(line, column) in &input.lookups {
+        black_box(input.sm.original_position_for(black_box(line), black_box(column)));
+    }
+    input
+}
+
 fn run_lazy_lookups(input: &mut LazyLookupInput) {
     for &(line, column) in &input.lookups {
         black_box(input.sm.original_position_for(black_box(line), black_box(column)));
@@ -667,28 +674,52 @@ fn bench_lookup(criterion: &mut Criterion) {
 fn bench_lazy_lookup(criterion: &mut Criterion) {
     bench_with_input(
         criterion,
-        "lazy_lookup_1000x_ascending",
+        "lazy_lookup_cold_1000x_ascending",
         ascending_lazy_lookup_input,
         run_lazy_lookups,
     );
     bench_with_input(
         criterion,
-        "lazy_lookup_1000x_descending",
+        "lazy_lookup_cold_1000x_descending",
         descending_lazy_lookup_input,
         run_lazy_lookups,
     );
     bench_with_input(
         criterion,
-        "lazy_lookup_1000x_repeated",
+        "lazy_lookup_cold_1000x_repeated",
         repeated_lazy_lookup_input,
         run_lazy_lookups,
     );
     bench_with_input(
         criterion,
-        "lazy_lookup_1000x_randomized",
+        "lazy_lookup_cold_1000x_randomized",
         randomized_lazy_lookup_input,
         run_lazy_lookups,
     );
+    bench_with_input(
+        criterion,
+        "lazy_lookup_warm_1000x_ascending",
+        || warm_lazy_lookup_input(ascending_lazy_lookup_input()),
+        run_lazy_lookups,
+    );
+    bench_with_input(
+        criterion,
+        "lazy_lookup_warm_1000x_descending",
+        || warm_lazy_lookup_input(descending_lazy_lookup_input()),
+        run_lazy_lookups,
+    );
+    bench_with_input(
+        criterion,
+        "lazy_lookup_warm_1000x_repeated",
+        || warm_lazy_lookup_input(repeated_lazy_lookup_input()),
+        run_lazy_lookups,
+    );
+    bench_with_input(
+        criterion,
+        "lazy_lookup_warm_1000x_randomized",
+        || warm_lazy_lookup_input(randomized_lazy_lookup_input()),
+        run_lazy_lookups,
+    );
 }
 
 fn bench_vlq_decode(criterion: &mut Criterion) {
diff --git a/crates/sourcemap/src/lib.rs b/crates/sourcemap/src/lib.rs
@@ -23,7 +23,7 @@
 //! assert_eq!(pos.column, 0);
 //! ```
 
-use std::cell::{OnceCell, RefCell};
+use std::cell::{Cell, OnceCell, RefCell};
 use std::collections::HashMap;
 use std::fmt;
 use std::io;
@@ -32,9 +32,6 @@ use serde::Deserialize;
 use srcmap_codec::{DecodeError, vlq_encode_unsigned};
 use srcmap_scopes::{Binding, CallSite, GeneratedRange, OriginalScope, Position, ScopeInfo};
 
-#[cfg(test)]
-use std::cell::Cell;
-
 pub mod js_identifiers;
 pub mod offset_lookup;
 pub mod source_view;
@@ -1990,11 +1987,10 @@ pub struct LazySourceMap {
     /// If true, decode_line must decode sequentially from the start.
     fast_scan: bool,
 
-    /// Cumulative VLQ states indexed by the next undecoded line.
-    state_checkpoints: RefCell<Vec<VlqState>>,
-
-    #[cfg(test)]
-    decode_work: Cell<u32>,
+    /// Highest line fully decoded so far (for progressive decode in fast-scan mode).
+    /// VLQ state at the end of this line is stored in `decode_state`.
+    decode_watermark: Cell<u32>,
+    decode_state: Cell<VlqState>,
 }
 
 impl LazySourceMap {
@@ -2032,13 +2028,8 @@ impl LazySourceMap {
             decoded_lines: RefCell::new(HashMap::new()),
             source_map,
             fast_scan,
-            state_checkpoints: RefCell::new(if fast_scan {
-                vec![VlqState::default()]
-            } else {
-                Vec::new()
-            }),
-            #[cfg(test)]
-            decode_work: Cell::new(0),
+            decode_watermark: Cell::new(0),
+            decode_state: Cell::new(VlqState::default()),
         }
     }
 
@@ -2331,16 +2322,13 @@ impl LazySourceMap {
         }
 
         if self.fast_scan {
-            let mut checkpoints = self.state_checkpoints.borrow_mut();
-            let start = line_idx.min(checkpoints.len() - 1);
-            let mut state = checkpoints[start];
+            let watermark = self.decode_watermark.get();
+            let start = if line >= watermark { watermark } else { 0 };
+            let mut state =
+                if line >= watermark { self.decode_state.get() } else { VlqState::default() };
 
-            for current_line_idx in start..=line_idx {
-                #[cfg(test)]
-                self.decode_work.set(self.decode_work.get() + 1);
-
-                let current_line = current_line_idx as u32;
-                let info = &self.line_info[current_line_idx];
+            for current_line in start..=line {
+                let info = &self.line_info[current_line as usize];
                 if self.decoded_lines.borrow().contains_key(&current_line) {
                     // Already cached — just walk VLQ bytes to compute end-state
                     let bytes = self.raw_mappings.as_bytes();
@@ -2350,13 +2338,11 @@ impl LazySourceMap {
                     state = new_state;
                     self.decoded_lines.borrow_mut().insert(current_line, mappings);
                 }
+            }
 
-                let next_line_idx = current_line_idx + 1;
-                if next_line_idx == checkpoints.len() {
-                    checkpoints.push(state);
-                } else {
-                    checkpoints[next_line_idx] = state;
-                }
+            if line + 1 > self.decode_watermark.get() {
+                self.decode_watermark.set(line + 1);
+                self.decode_state.set(state);
             }
 
             return Ok(());
@@ -7302,22 +7288,6 @@ mod tests {
         assert_eq!(loc0.line, 0);
     }
 
-    #[test]
-    fn lazy_sourcemap_backward_cache_miss_resumes_from_checkpoint() {
-        let mappings = std::iter::repeat_n("AACA", 64).collect::<Vec<_>>().join(";");
-        let json =
-            format!(r#"{{"version":3,"sources":["a.js"],"names":[],"mappings":"{mappings}"}}"#);
-        let sm = LazySourceMap::from_json_fast(&json).unwrap();
-
-        sm.decode_line(63).unwrap();
-        sm.decoded_lines.borrow_mut().remove(&15);
-        sm.decode_work.set(0);
-
-        sm.decode_line(15).unwrap();
-
-        assert_eq!(sm.decode_work.get(), 1, "must resume at the line-15 checkpoint");
-    }
-
     #[test]
     fn lazy_sourcemap_fast_scan_matches_prescan_in_descending_and_randomized_order() {
         let mappings = std::iter::repeat_n("AACA,CAAC", 64).collect::<Vec<_>>().join(";");
