@takara-ai/miru-code 0.1.0

package/src/index/bm25.ts

@@ -0,0 +1,238 @@
+/** Okapi BM25 with inverted postings; optional worker-parallel scoring. */
+
+import { resolveWorkerConcurrency } from "../concurrency.ts";
+import type {
+  Bm25ScoreJob,
+  Bm25ScoreResult,
+  PostingsList,
+  SerializablePostings,
+} from "./bm25-types.ts";
+
+const K1 = 1.5;
+const B = 0.75;
+
+/** Use workers when doc count exceeds this (overhead not worth it below). */
+const PARALLEL_SCORE_MIN_DOCS = 256;
+
+function docRanges(numDocs: number, parts: number): [number, number][] {
+  const n = Math.max(1, Math.min(parts, numDocs));
+  const ranges: [number, number][] = [];
+  const chunk = Math.ceil(numDocs / n);
+  for (let start = 0; start < numDocs; start += chunk) {
+    ranges.push([start, Math.min(numDocs, start + chunk)]);
+  }
+  return ranges;
+}
+
+function runScoreWorker(job: Bm25ScoreJob): Promise<Bm25ScoreResult> {
+  return new Promise((resolve, reject) => {
+    const worker = new Worker(new URL("./bm25-worker.ts", import.meta.url).href, {
+      type: "module",
+    });
+    worker.onmessage = (event: MessageEvent<Bm25ScoreResult>) => {
+      worker.terminate();
+      resolve(event.data);
+    };
+    worker.onerror = (err) => {
+      worker.terminate();
+      reject(err);
+    };
+    worker.postMessage(job);
+  });
+}
+
+export class BM25Index {
+  private docFreq = new Map<string, number>();
+  private docLengths: number[] = [];
+  private postings = new Map<string, PostingsList>();
+  private avgDocLength = 0;
+  private numDocs = 0;
+
+  index(tokenizedDocs: string[][]): void {
+    this.numDocs = tokenizedDocs.length;
+    this.docFreq.clear();
+    this.docLengths = [];
+    this.postings.clear();
+
+    let totalLen = 0;
+    for (let docIndex = 0; docIndex < tokenizedDocs.length; docIndex++) {
+      const doc = tokenizedDocs[docIndex];
+      if (!doc) {
+        continue;
+      }
+      this.docLengths.push(doc.length);
+      totalLen += doc.length;
+
+      const tf = new Map<string, number>();
+      for (const term of doc) {
+        tf.set(term, (tf.get(term) ?? 0) + 1);
+      }
+      for (const [term, count] of tf) {
+        this.docFreq.set(term, (this.docFreq.get(term) ?? 0) + 1);
+        let list = this.postings.get(term);
+        if (!list) {
+          list = [];
+          this.postings.set(term, list);
+        }
+        list.push([docIndex, count]);
+      }
+    }
+    this.avgDocLength = this.numDocs > 0 ? totalLen / this.numDocs : 0;
+  }
+
+  private queryTerms(queryTokens: string[]): { term: string; idf: number }[] {
+    const unique = [...new Set(queryTokens)];
+    const terms: { term: string; idf: number }[] = [];
+    for (const term of unique) {
+      const df = this.docFreq.get(term) ?? 0;
+      if (df === 0) {
+        continue;
+      }
+      terms.push({
+        term,
+        idf: Math.log(1 + (this.numDocs - df + 0.5) / (df + 0.5)),
+      });
+    }
+    return terms;
+  }
+
+  private postingsForTerms(terms: { term: string; idf: number }[]): SerializablePostings {
+    const out: SerializablePostings = {};
+    for (const { term } of terms) {
+      const list = this.postings.get(term);
+      if (list) {
+        out[term] = list;
+      }
+    }
+    return out;
+  }
+
+  getScores(queryTokens: string[], weightMask?: boolean[]): number[] {
+    const scores = new Array<number>(this.numDocs).fill(0);
+    if (this.numDocs === 0 || queryTokens.length === 0) {
+      return scores;
+    }
+
+    const terms = this.queryTerms(queryTokens);
+    const avg = this.avgDocLength || 1;
+
+    for (const { term, idf } of terms) {
+      const list = this.postings.get(term);
+      if (!list) {
+        continue;
+      }
+      for (const [docIndex, tf] of list) {
+        if (weightMask && !weightMask[docIndex]) {
+          continue;
+        }
+        const dl = this.docLengths[docIndex];
+        if (dl === undefined) {
+          continue;
+        }
+        const denom = tf + K1 * (1 - B + (B * dl) / avg);
+        scores[docIndex] = (scores[docIndex] ?? 0) + idf * ((tf * (K1 + 1)) / denom);
+      }
+    }
+
+    return scores;
+  }
+
+  async getScoresAsync(queryTokens: string[], weightMask?: boolean[]): Promise<number[]> {
+    const scores = new Array<number>(this.numDocs).fill(0);
+    if (this.numDocs === 0 || queryTokens.length === 0) {
+      return scores;
+    }
+
+    const terms = this.queryTerms(queryTokens);
+    if (terms.length === 0) {
+      return scores;
+    }
+
+    const concurrency = resolveWorkerConcurrency();
+    if (this.numDocs < PARALLEL_SCORE_MIN_DOCS || concurrency <= 1) {
+      return this.getScores(queryTokens, weightMask);
+    }
+
+    const postings = this.postingsForTerms(terms);
+    const ranges = docRanges(this.numDocs, concurrency);
+    const baseJob = {
+      docLengths: this.docLengths,
+      avgDocLength: this.avgDocLength,
+      terms,
+      postings,
+      weightMask,
+    };
+
+    const partials = await Promise.all(
+      ranges.map(([startDoc, endDoc]) => runScoreWorker({ ...baseJob, startDoc, endDoc })),
+    );
+
+    for (const { startDoc, scores: partial } of partials) {
+      for (let i = 0; i < partial.length; i++) {
+        const value = partial[i];
+        if (value !== undefined) {
+          scores[startDoc + i] = value;
+        }
+      }
+    }
+
+    return scores;
+  }
+
+  toJSON(): {
+    postings: SerializablePostings;
+    docLengths: number[];
+    avgDocLength: number;
+    numDocs: number;
+  } {
+    const postings: SerializablePostings = {};
+    for (const [term, list] of this.postings.entries()) {
+      postings[term] = list;
+    }
+    return {
+      postings,
+      docLengths: this.docLengths,
+      avgDocLength: this.avgDocLength,
+      numDocs: this.numDocs,
+    };
+  }
+
+  static fromJSON(
+    data:
+      | {
+          // Legacy cache format
+          docs: string[][];
+          docLengths: number[];
+          docFreq: [string, number][];
+        }
+      | {
+          // Current cache format
+          postings: SerializablePostings;
+          docLengths: number[];
+          avgDocLength: number;
+          numDocs: number;
+        },
+  ): BM25Index {
+    const idx = new BM25Index();
+
+    // New cache: restore postings directly (no need to re-index tokenized docs).
+    if ("postings" in data) {
+      idx.postings.clear();
+      idx.docFreq.clear();
+      idx.docLengths = data.docLengths;
+      idx.avgDocLength = data.avgDocLength;
+      idx.numDocs = data.numDocs;
+
+      for (const [term, list] of Object.entries(data.postings)) {
+        idx.postings.set(term, list);
+        idx.docFreq.set(term, list.length);
+      }
+
+      return idx;
+    }
+
+    // Legacy cache: re-index from stored tokenized documents.
+    idx.index(data.docs);
+    return idx;
+  }
+}

package/src/index/create.ts

@@ -0,0 +1,195 @@
+import { relative, resolve } from "node:path";
+import { chunkSource } from "../chunking/chunking.ts";
+import { resolveWorkerConcurrency } from "../concurrency.ts";
+import type { EmbeddingBackend } from "../embeddings/openai.ts";
+import { envOptionalInt } from "../env.ts";
+import { tokenize } from "../tokens.ts";
+import type { Chunk, ContentType } from "../types.ts";
+import { BM25Index } from "./bm25.ts";
+import { walkFiles } from "./file-walker.ts";
+import { detectLanguage, getExtensions, getFileStatus, readFileText } from "./files.ts";
+import type { SemanticIndex } from "./semantic-index.ts";
+import { enrichForBm25 } from "./sparse.ts";
+import { buildSemanticIndex } from "./vector-storage.ts";
+
+const DEFAULT_PIPELINE_EMBED_BATCH = 64;
+const DEFAULT_PIPELINE_EMBED_INFLIGHT = 8;
+
+function resolvePipelineEmbedBatch(): number {
+  return (
+    envOptionalInt(["MIRU_PIPELINE_EMBED_BATCH", "SEMBLE_PIPELINE_EMBED_BATCH"], 1) ??
+    DEFAULT_PIPELINE_EMBED_BATCH
+  );
+}
+
+function resolvePipelineEmbedInflight(): number {
+  return (
+    envOptionalInt(["MIRU_PIPELINE_EMBED_INFLIGHT", "SEMBLE_PIPELINE_EMBED_INFLIGHT"], 1) ??
+    DEFAULT_PIPELINE_EMBED_INFLIGHT
+  );
+}
+
+export async function createIndexFromPath(
+  path: string,
+  embeddings: EmbeddingBackend,
+  content: ContentType[] = ["code"],
+  displayRoot?: string,
+): Promise<{ bm25: BM25Index; semantic: SemanticIndex; chunks: Chunk[] }> {
+  const profile = process.env.MIRU_PROFILE === "1" || process.env.SEMBLE_PROFILE === "1";
+  const started = performance.now();
+  let fileProcessMs = 0;
+  let embedBackpressureWaitMs = 0;
+  let fileEnumerated = 0;
+  let fileTasks = 0;
+  let emptyFileTasks = 0;
+
+  if (profile && "resetStats" in embeddings && typeof embeddings.resetStats === "function") {
+    embeddings.resetStats();
+  }
+
+  const resolved = resolve(path);
+  const root = displayRoot ? resolve(displayRoot) : resolved;
+  const extensions = getExtensions(content);
+  const fileConcurrency = resolveWorkerConcurrency();
+  const embedBatchSize = resolvePipelineEmbedBatch();
+  const maxEmbedInflight = resolvePipelineEmbedInflight();
+  const readyBySeq = new Map<number, Chunk[]>();
+  const pendingEmbedChunks: Chunk[] = [];
+  const emittedChunkBatches: Chunk[][] = [];
+  const embedInFlight = new Set<Promise<Float32Array[]>>();
+  const embedPromises: Promise<Float32Array[]>[] = [];
+  let fileSeq = 0;
+  let nextEmitSeq = 0;
+
+  const processFile = async (filePath: string): Promise<Chunk[]> => {
+    const t = performance.now();
+    const language = detectLanguage(filePath);
+    try {
+      const status = await getFileStatus(filePath);
+      if (status !== "valid") {
+        emptyFileTasks++;
+        return [] as Chunk[];
+      }
+      const source = await readFileText(filePath);
+      const chunkPath = displayRoot ? relative(root, filePath).replace(/\\/g, "/") : filePath;
+      return chunkSource(source, chunkPath, language);
+    } catch {
+      emptyFileTasks++;
+      return [];
+    } finally {
+      fileProcessMs += performance.now() - t;
+    }
+  };
+
+  const flushReadySequentialChunks = (): void => {
+    while (readyBySeq.has(nextEmitSeq)) {
+      const chunks = readyBySeq.get(nextEmitSeq);
+      readyBySeq.delete(nextEmitSeq);
+      if (chunks) {
+        pendingEmbedChunks.push(...chunks);
+      }
+      nextEmitSeq++;
+    }
+  };
+
+  const maybeScheduleEmbed = async (force = false): Promise<void> => {
+    while (
+      pendingEmbedChunks.length > 0 &&
+      (force || pendingEmbedChunks.length >= embedBatchSize)
+    ) {
+      if (embedInFlight.size >= maxEmbedInflight) {
+        const waitStart = performance.now();
+        await Promise.race(embedInFlight);
+        embedBackpressureWaitMs += performance.now() - waitStart;
+        continue;
+      }
+
+      const take = force
+        ? pendingEmbedChunks.length
+        : Math.min(embedBatchSize, pendingEmbedChunks.length);
+      const batch = pendingEmbedChunks.splice(0, take);
+      emittedChunkBatches.push(batch);
+      const promise = embeddings.embedDocuments(batch.map((chunk) => chunk.content));
+      embedPromises.push(promise);
+      embedInFlight.add(promise);
+      promise.finally(() => {
+        embedInFlight.delete(promise);
+      });
+    }
+  };
+
+  const runningFiles = new Set<Promise<void>>();
+
+  const startFileTask = (filePath: string, seq: number): void => {
+    const task = (async () => {
+      const chunks = await processFile(filePath);
+      readyBySeq.set(seq, chunks);
+      flushReadySequentialChunks();
+      await maybeScheduleEmbed(false);
+    })().finally(() => {
+      runningFiles.delete(task);
+    });
+    runningFiles.add(task);
+  };
+
+  for await (const filePath of walkFiles(resolved, extensions)) {
+    fileEnumerated++;
+    while (runningFiles.size >= fileConcurrency) {
+      await Promise.race(runningFiles);
+    }
+    fileTasks++;
+    startFileTask(filePath, fileSeq++);
+  }
+
+  while (runningFiles.size > 0) {
+    await Promise.race(runningFiles);
+  }
+
+  flushReadySequentialChunks();
+  await maybeScheduleEmbed(true);
+  await Promise.all(embedPromises);
+
+  const chunks = emittedChunkBatches.flat();
+
+  if (chunks.length === 0) {
+    throw new Error(`No supported files found under ${path}.`);
+  }
+
+  const vectors = (await Promise.all(embedPromises)).flat();
+  const bm25 = new BM25Index();
+  bm25.index(chunks.map((c) => tokenize(enrichForBm25(c))));
+  const semantic = buildSemanticIndex(vectors);
+
+  if (profile) {
+    const finished = performance.now();
+    const embedStats =
+      "getStats" in embeddings && typeof embeddings.getStats === "function"
+        ? embeddings.getStats()
+        : null;
+    console.error(
+      JSON.stringify({
+        profile: "index_build",
+        path: resolved,
+        elapsed_ms: finished - started,
+        file_enumerated: fileEnumerated,
+        file_tasks: fileTasks,
+        empty_or_skipped_files: emptyFileTasks,
+        chunks: chunks.length,
+        vectors: vectors.length,
+        pipeline: {
+          file_concurrency: fileConcurrency,
+          embed_batch_size: embedBatchSize,
+          embed_inflight: maxEmbedInflight,
+          emitted_batches: emittedChunkBatches.length,
+        },
+        stage_ms: {
+          file_process_total: fileProcessMs,
+          embed_backpressure_wait: embedBackpressureWaitMs,
+        },
+        embedding_transport: embedStats,
+      }),
+    );
+  }
+
+  return { bm25, semantic, chunks };
+}

package/src/index/dense.ts

@@ -0,0 +1,105 @@
+/** Cosine-distance vector index with optional chunk selector. */
+import { mkdir } from "node:fs/promises";
+import type { SemanticIndex } from "./semantic-index.ts";
+import { selectTopKByDistance } from "./topk.ts";
+
+export class VectorIndex implements SemanticIndex {
+  private vectors: Float32Array[];
+
+  constructor(vectors: Float32Array[]) {
+    this.vectors = vectors;
+  }
+
+  get size(): number {
+    return this.vectors.length;
+  }
+
+  get dimensions(): number {
+    return this.vectors[0]?.length ?? 0;
+  }
+
+  getVectors(): readonly Float32Array[] {
+    return this.vectors;
+  }
+
+  memoryBytes(): number {
+    return this.vectors.length * this.dimensions * 4;
+  }
+
+  query(
+    queryVector: Float32Array,
+    k: number,
+    selector?: number[],
+  ): { indices: number[]; distances: number[] } {
+    if (k < 1) {
+      throw new Error(`k should be >= 1, is now ${k}`);
+    }
+
+    const numVectors = this.vectors.length;
+    if (numVectors === 0) {
+      return { indices: [], distances: [] };
+    }
+
+    const indices = selector ?? Array.from({ length: numVectors }, (_, i) => i);
+    const effectiveK = Math.min(k, indices.length);
+    if (effectiveK === 0) {
+      return { indices: [], distances: [] };
+    }
+
+    const top = selectTopKByDistance(
+      indices.flatMap((idx) => {
+        const vec = this.vectors[idx];
+        if (!vec) {
+          return [];
+        }
+        return [{ index: idx, distance: cosineDistance(queryVector, vec) }];
+      }),
+      effectiveK,
+    );
+
+    return {
+      indices: top.map((t) => t.index),
+      distances: top.map((t) => t.distance),
+    };
+  }
+
+  async save(dir: string): Promise<void> {
+    await mkdir(dir, { recursive: true });
+
+    const dim = this.dimensions;
+    const count = this.size;
+    const buffer = new Float32Array(count * dim);
+    for (let i = 0; i < count; i++) {
+      const vec = this.vectors[i];
+      if (vec) {
+        buffer.set(vec, i * dim);
+      }
+    }
+    await Bun.write(`${dir}/vectors.bin`, buffer);
+    await Bun.write(
+      `${dir}/meta.json`,
+      JSON.stringify({ count, dimensions: dim, storage: "float32" }),
+    );
+  }
+
+  static async load(dir: string): Promise<VectorIndex> {
+    const meta = JSON.parse(await Bun.file(`${dir}/meta.json`).text()) as {
+      count: number;
+      dimensions: number;
+    };
+    const raw = new Float32Array(await Bun.file(`${dir}/vectors.bin`).arrayBuffer());
+    const vectors: Float32Array[] = [];
+    for (let i = 0; i < meta.count; i++) {
+      vectors.push(raw.subarray(i * meta.dimensions, (i + 1) * meta.dimensions));
+    }
+    return new VectorIndex(vectors);
+  }
+}
+
+function cosineDistance(a: Float32Array, b: Float32Array): number {
+  let dot = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += (a[i] ?? 0) * (b[i] ?? 0);
+  }
+  return 1 - dot;
+}

package/src/index/file-walker.ts

@@ -0,0 +1,113 @@
+import type { Dirent } from "node:fs";
+import { readdir } from "node:fs/promises";
+import { join, relative } from "node:path";
+import ignore, { type Ignore } from "ignore";
+
+const DEFAULT_IGNORED_DIRS = new Set([
+  ".git",
+  ".hg",
+  ".svn",
+  "__pycache__",
+  "node_modules",
+  ".venv",
+  "venv",
+  ".tox",
+  ".mypy_cache",
+  ".pytest_cache",
+  ".ruff_cache",
+  ".cache",
+  ".miru",
+  ".next",
+  "dist",
+  "build",
+  ".eggs",
+]);
+
+interface IgnoreSpec {
+  base: string;
+  spec: Ignore;
+}
+
+async function loadIgnoreForDir(directory: string): Promise<IgnoreSpec | null> {
+  const paths = [join(directory, ".gitignore"), join(directory, ".miruignore")];
+  const lines: string[] = [];
+
+  for (const p of paths) {
+    const file = Bun.file(p);
+    if (await file.exists()) {
+      lines.push(...(await file.text()).split("\n"));
+    }
+  }
+
+  if (lines.length === 0) {
+    return null;
+  }
+  return {
+    base: directory,
+    spec: ignore().add(lines),
+  };
+}
+
+function isIgnoredBySpecs(fullPath: string, isDir: boolean, specs: IgnoreSpec[]): boolean {
+  for (const spec of specs) {
+    const relToSpec = relative(spec.base, fullPath).replace(/\\/g, "/");
+    if (!relToSpec || relToSpec.startsWith("..")) {
+      continue;
+    }
+    if (spec.spec.ignores(isDir ? `${relToSpec}/` : relToSpec)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+export async function* walkFiles(root: string, extensions: string[]): AsyncGenerator<string> {
+  const extSet = new Set(extensions.map((e) => e.toLowerCase()));
+  const loadedSpecs = new Map<string, IgnoreSpec | null>();
+
+  async function* walk(dir: string, inheritedSpecs: IgnoreSpec[]): AsyncGenerator<string> {
+    let entries: Dirent[];
+    try {
+      entries = await readdir(dir, { withFileTypes: true });
+    } catch {
+      return;
+    }
+
+    let spec = loadedSpecs.get(dir);
+    if (spec === undefined) {
+      spec = (await loadIgnoreForDir(dir)) ?? null;
+      loadedSpecs.set(dir, spec);
+    }
+    const activeSpecs = spec ? [...inheritedSpecs, spec] : inheritedSpecs;
+
+    for (const entry of entries) {
+      const fullPath = join(dir, entry.name);
+
+      if (entry.isDirectory()) {
+        if (DEFAULT_IGNORED_DIRS.has(entry.name)) {
+          continue;
+        }
+        if (isIgnoredBySpecs(fullPath, true, activeSpecs)) {
+          continue;
+        }
+        yield* walk(fullPath, activeSpecs);
+        continue;
+      }
+
+      if (!entry.isFile()) {
+        continue;
+      }
+
+      if (isIgnoredBySpecs(fullPath, false, activeSpecs)) {
+        continue;
+      }
+
+      const ext = entry.name.includes(".") ? `.${entry.name.split(".").pop()?.toLowerCase()}` : "";
+      if (entry.name.toLowerCase() === "dockerfile" || extSet.has(ext)) {
+        yield fullPath;
+      }
+    }
+  }
+
+  yield* walk(root, []);
+}