perf(cache): batch writes and validate hydration

src/chunking/chunkFile.ts

@@ -39,7 +39,12 @@ export interface ChunkFileOptions {
  * @param opts Chunking options
  * @returns Array of semantic chunks
  */
-export function chunkFile(opts: ChunkFileOptions): Chunk[] {
+export type ChunkFileWithSymbolsResult = {
+  chunks: Chunk[];
+  symbolChunks: Chunk[];
+};
+
+export function chunkFileWithSymbols(opts: ChunkFileOptions): ChunkFileWithSymbolsResult {
   const { language, source, filePath, minTokens = 150, maxTokens = 400, tokenizer = countWhitespaceTokens } = opts;
   const matches = getChunkMatches(language, source, filePath);
   const newlineOffsets = collectNewlineOffsets(source);
@@ -79,7 +84,43 @@ export function chunkFile(opts: ChunkFileOptions): Chunk[] {
   }
 
   preliminaryChunks.sort((left, right) => left.sourceStart - right.sourceStart || right.sourceEnd - left.sourceEnd);
-  const mergedChunks = mergeSmallChunks(preliminaryChunks, minTokens, maxTokens, tokenizer);
+
+  if (minTokens <= 1) {
+    const completeChunks = fillGapsWithMiscChunks(
+      preliminaryChunks,
+      source,
+      language.id,
+      filePath,
+      tokenizer,
+      1,
+      maxTokens,
+      newlineOffsets,
+    );
+    const chunks = withStableChunkIds(
+      completeChunks.map(({ sourceStart: _sourceStart, sourceEnd: _sourceEnd, ...chunk }) => chunk),
+      language.id,
+      filePath,
+    );
+    return { chunks, symbolChunks: chunks };
+  }
+
+  const completeSymbolChunks = fillGapsWithMiscChunks(
+    preliminaryChunks,
+    source,
+    language.id,
+    filePath,
+    tokenizer,
+    1,
+    maxTokens,
+    newlineOffsets,
+  );
+  const symbolChunks = withStableChunkIds(
+    completeSymbolChunks.map(({ sourceStart: _sourceStart, sourceEnd: _sourceEnd, ...chunk }) => chunk),
+    language.id,
+    filePath,
+  );
+
+  const mergedChunks = mergeSmallChunks([...preliminaryChunks], minTokens, maxTokens, tokenizer);
   const completeChunks = fillGapsWithMiscChunks(
     mergedChunks,
     source,
@@ -90,12 +131,17 @@ export function chunkFile(opts: ChunkFileOptions): Chunk[] {
     maxTokens,
     newlineOffsets,
   );
-
-  return withStableChunkIds(
+  const chunks = withStableChunkIds(
     completeChunks.map(({ sourceStart: _sourceStart, sourceEnd: _sourceEnd, ...chunk }) => chunk),
     language.id,
     filePath,
   );
+
+  return { chunks, symbolChunks };
+}
+
+export function chunkFile(opts: ChunkFileOptions): Chunk[] {
+  return chunkFileWithSymbols(opts).chunks;
 }
 
 function appendBlockChunks(

src/duplicates/unitCache.ts

@@ -360,43 +360,80 @@ export function tryLoadDuplicateUnitsFromCache(
   }
 }
 
-export function writeDuplicateUnitsToCache(
+export type PendingDuplicateUnitCacheWrite = {
+  file: string;
+  variant: string;
+  units: DuplicateInternalUnit[];
+};
+
+export function writeDuplicateUnitsBatchToCache(
   index: ProjectIndex,
-  file: string,
-  variant: string,
-  units: DuplicateInternalUnit[],
+  writes: readonly PendingDuplicateUnitCacheWrite[],
 ): void {
-  const sig = duplicateUnitCacheSignature(index, file);
-  if (!sig) return;
-  const key = duplicateUnitCacheKey(file, variant);
+  if (!writes.length) return;
+  const root = index.projectRoot ?? "";
   if (index.cacheMode === "memory") {
-    writeDuplicateUnitMemoryCache(key, { sig, units });
+    for (const write of writes) {
+      const sig = duplicateUnitCacheSignature(index, write.file);
+      if (!sig) continue;
+      writeDuplicateUnitMemoryCache(duplicateUnitCacheKey(write.file, write.variant), {
+        sig,
+        units: write.units,
+      });
+    }
     return;
   }
-  if (index.cacheMode === "disk") {
-    try {
-      const entry = duplicateUnitDiskCache(index);
-      const root = index.projectRoot ?? "";
+  if (index.cacheMode !== "disk") return;
+  try {
+    const entry = duplicateUnitDiskCache(index);
+    const preparedWrites: Array<{
+      file: string;
+      variant: string;
+      sig: string;
+      payload: Buffer;
+    }> = [];
+    for (const write of writes) {
+      const sig = duplicateUnitCacheSignature(index, write.file);
+      if (!sig) continue;
       const payload = brotliCompressSync(
-        JSON.stringify(transformDuplicateUnits(root, serializeDuplicateUnits(units), true)),
-        {
-          params: { [zlibConstants.BROTLI_PARAM_QUALITY]: 4 },
-        },
+        JSON.stringify(transformDuplicateUnits(root, serializeDuplicateUnits(write.units), true)),
+        { params: { [zlibConstants.BROTLI_PARAM_QUALITY]: 4 } },
       );
-      entry?.statements?.write.run(
-        cacheRelativePath(root, file),
-        variant,
+      preparedWrites.push({
+        file: cacheRelativePath(root, write.file),
+        variant: write.variant,
         sig,
-        DUPLICATE_UNIT_CACHE_VERSION,
         payload,
-        Date.now(),
-      );
-    } catch {
-      // best-effort cache
+      });
     }
+    if (!preparedWrites.length || !entry?.db || !entry.statements) return;
+    const now = Date.now();
+    entry.db.transaction(() => {
+      for (const write of preparedWrites) {
+        entry.statements?.write.run(
+          write.file,
+          write.variant,
+          write.sig,
+          DUPLICATE_UNIT_CACHE_VERSION,
+          write.payload,
+          now,
+        );
+      }
+    })();
+  } catch {
+    // best-effort cache
   }
 }
 
+export function writeDuplicateUnitsToCache(
+  index: ProjectIndex,
+  file: string,
+  variant: string,
+  units: DuplicateInternalUnit[],
+): void {
+  writeDuplicateUnitsBatchToCache(index, [{ file, variant, units }]);
+}
+
 export function deserializeDuplicateUnits(value: unknown): DuplicateInternalUnit[] | null {
   if (!Array.isArray(value) || !value.every(isDuplicateSerializedUnit)) return null;
   return value.map((unit) => ({
@@ -430,7 +467,7 @@ function transformDuplicateUnits(
 ): DuplicateSerializedUnit[] {
   return units.map((unit) => ({
     ...unit,
-    file: transformDuplicatePath(root, unit.file, toRelative),
+    file: cacheRelativePath(root, unit.file),
     absoluteFile: transformDuplicatePath(root, unit.absoluteFile, toRelative),
     handle: transformDuplicatePath(root, unit.handle, toRelative),
     fileHandle: transformDuplicatePath(root, unit.fileHandle, toRelative),

src/duplicates/units.ts

@@ -2,7 +2,7 @@ import crypto from "node:crypto";
 import fsp from "node:fs/promises";
 import path from "node:path";
 import { LANG_CONFIGS } from "../bootstrap/treeSitterLanguages.js";
-import { chunkFile, type Chunk } from "../chunking/chunkFile.js";
+import { chunkFile, chunkFileWithSymbols, type Chunk } from "../chunking/chunkFile.js";
 import { chunkTextFile } from "../chunking/chunkTextFile.js";
 import {
   countDuplicateTokens,
@@ -25,7 +25,13 @@ import { maskJsLikeCommentsStringsAndRegex } from "../util/comments.js";
 import { collectLineStartOffsets } from "../util/lines.js";
 import { assertFilePathWithinRoot, fileIdentityKey, normalizePath, toProjectDisplayPath } from "../util/paths.js";
 import { logWithLevel } from "../logging.js";
-import { duplicateUnitCacheVariant, tryLoadDuplicateUnitsFromCache, writeDuplicateUnitsToCache } from "./unitCache.js";
+import {
+  duplicateUnitCacheVariant,
+  tryLoadDuplicateUnitsFromCache,
+  writeDuplicateUnitsBatchToCache,
+  writeDuplicateUnitsToCache,
+  type PendingDuplicateUnitCacheWrite,
+} from "./unitCache.js";
 import type {
   CollectedDuplicateUnits,
   DuplicateAstContext,
@@ -290,6 +296,29 @@ export function makeDuplicateChunks(
   return chunkTextFile({ source, filePath, languageId, minTokens, maxTokens, tokenizer: countDuplicateTokens });
 }
 
+export function makeDuplicateChunksWithSymbols(
+  filePath: string,
+  languageId: string,
+  textOnly: boolean,
+  source: string,
+  minTokens: number,
+  maxTokens: number,
+): { chunks: Chunk[]; symbolChunks: Chunk[] } {
+  const langConfig = LANG_CONFIGS[chunkLanguageAliases[languageId] ?? languageId];
+  if (langConfig && !textOnly) {
+    return chunkFileWithSymbols({
+      language: langConfig,
+      source,
+      filePath,
+      minTokens,
+      maxTokens,
+      tokenizer: countDuplicateTokens,
+    });
+  }
+  const chunks = chunkTextFile({ source, filePath, languageId, minTokens, maxTokens, tokenizer: countDuplicateTokens });
+  return { chunks, symbolChunks: [] };
+}
+
 export function makeSymbolSourceChunks(
   filePath: string,
   languageId: string,
@@ -506,6 +535,7 @@ export async function collectDuplicateUnits(
   let belowThresholdUnits = 0;
   const belowThresholdUnitsByFile = new Map<string, number>();
 
+  const pendingWrites: PendingDuplicateUnitCacheWrite[] = [];
   for (const file of normalizedFiles) {
     const cachedUnits = tryLoadDuplicateUnitsFromCache(index, file, variant);
     const fileUnits =
@@ -521,7 +551,7 @@ export async function collectDuplicateUnits(
         astContextCache,
       ));
     if (!cachedUnits) {
-      writeDuplicateUnitsToCache(index, file, variant, fileUnits);
+      pendingWrites.push({ file, variant, units: fileUnits });
     }
     for (const unit of fileUnits) {
       if (!shouldKeepUnit(unit, options.includeSmall, options.minTokens)) {
@@ -532,6 +562,9 @@ export async function collectDuplicateUnits(
       units.push(unit);
     }
   }
+  if (pendingWrites.length) {
+    writeDuplicateUnitsBatchToCache(index, pendingWrites);
+  }
 
   units.sort((left, right) => {
     const fileCompare = left.absoluteFile.localeCompare(right.absoluteFile);
@@ -567,8 +600,14 @@ export async function buildDuplicateUnitsForFile(
   }
 
   const astContext = language.textOnly ? undefined : await getDuplicateAstContext(index, file, source, astContextCache);
-  const chunks = makeDuplicateChunks(file, language.id, language.textOnly, source, minTokens, maxTokens);
-  const symbolChunks = makeSymbolSourceChunks(file, language.id, language.textOnly, source, maxTokens);
+  const { chunks, symbolChunks } = makeDuplicateChunksWithSymbols(
+    file,
+    language.id,
+    language.textOnly,
+    source,
+    minTokens,
+    maxTokens,
+  );
   const symbolUnits = (moduleIndex?.locals ?? [])
     .map((symbol) => {
       const chunk = findChunkForSymbol(symbol, symbolChunks);

src/indexer/build-cache.ts

@@ -21,14 +21,19 @@ export {
   fileSignature,
   pruneDiskModuleCache,
   tryLoadFromCache,
+  writeModulesToCache,
   writeToCache,
   type FileSignature,
+  type PendingModuleCacheWrite,
 } from "./build-cache/module-cache.js";
 export {
+  BLOOM_FILTER_SNAPSHOT_FILENAME,
+  BLOOM_FILTER_SNAPSHOT_VERSION,
   createProjectSnapshotIdentity,
   projectSnapshotFilesSignature,
   tryLoadDetailedSymbolGraphSnapshot,
   tryLoadPersistedBloomFilters,
+  tryLoadProjectSnapshotModules,
   tryLoadProjectIndexSnapshot,
   writeDetailedSymbolGraphSnapshot,
   writeProjectIndexSnapshot,