dsh-recall/lib/chunk.js

/**
 * Splitting source into retrievable units.
 *
 * This is where retrieval quality is decided, and the rule that matters is that
 * **a chunk must be a thing a person would name**. Splitting on a fixed line
 * count produces units that begin mid-function and end mid-conditional; their
 * embedding sits between two topics and matches neither, so the search returns
 * plausible-looking noise. Splitting at declarations produces units whose
 * embedding is about one subject, which is the entire premise of the retrieval.
 *
 * The anchor set is deliberately shallow — column-0 declarations only. A real
 * parser would place boundaries better, but it would also mean shipping a
 * tree-sitter grammar per language and keeping ten of them current, for a gain
 * that only shows up on languages with heavy nesting. Since a chunk is a
 * *search result* and not an edit target, a boundary that is roughly right costs
 * a few wasted lines and nothing else.
 *
 * Fixed-width windows remain as the fallback, because a file whose structure the
 * anchor set cannot see — a minified bundle, a template, an unusual dialect —
 * must still be searchable rather than silently absent from the index.
 */

/** Lines a chunk aims for when merging neighbouring declarations. */
export const DEFAULT_CHUNK_LINES = 48

/** Lines a chunk may never exceed; longer ones are split with overlap. */
export const DEFAULT_MAX_CHUNK_LINES = 80

/** Lines shared between the halves of a split chunk, so a boundary is never lost. */
export const DEFAULT_OVERLAP_LINES = 8

/** Chunks taken from one file, so one pathological file cannot dominate the index. */
export const DEFAULT_MAX_CHUNKS_PER_FILE = 250

/** Files whose mean line length exceeds this are treated as generated/minified. */
const MINIFIED_MEAN_LINE = 400

/** Anchor patterns by language family, applied to the raw line. */
const FAMILIES = {
  javascript: /^(?:export\s+)?(?:default\s+)?(?:async\s+)?(?:function|class|const|let|var|interface|type|enum|namespace|declare)\b/,
  python: /^(?:async\s+)?(?:def|class)\s+[A-Za-z_]/,
  go: /^(?:func|type|var|const)\b/,
  rust: /^(?:pub(?:\([^)]*\))?\s+)?(?:async\s+)?(?:unsafe\s+)?(?:extern\s+)?(?:fn|struct|enum|impl|trait|mod|const|static|type)\b/,
  jvm: /^(?:@\w+|public|private|protected|static|final|abstract|sealed|open|class|interface|enum|record|object|fun|val|var|def)\b/,
  dotnet: /^(?:\[[^\]]*\]|public|private|protected|internal|static|final|abstract|sealed|partial|class|interface|enum|record|struct|namespace)\b/,
  c: /^(?:class|struct|enum|union|namespace|template|typedef|using|extern)\b|^[A-Za-z_][\w \t*&:<>]*\s+[A-Za-z_]\w*\s*\([^;]*\)\s*(?:const\s*)?$/,
  ruby: /^(?:def|class|module)\s/,
  php: /^(?:public|private|protected|static|final|abstract|function|class|interface|trait|namespace)\b/,
  shell: /^(?:function\s+[A-Za-z_]|[A-Za-z_]\w*\s*\(\)\s*\{)/,
  sql: /^(?:CREATE|ALTER|DROP|WITH)\b/i,
  markup: /^<(?:template|script|style|svelte:)/,
}

/** Extension to family. Anything unlisted falls back to the window splitter. */
const EXTENSION_FAMILY = {
  ts: 'javascript', tsx: 'javascript', mts: 'javascript', cts: 'javascript',
  js: 'javascript', jsx: 'javascript', mjs: 'javascript', cjs: 'javascript',
  py: 'python', rb: 'ruby', go: 'go', rs: 'rust',
  java: 'jvm', kt: 'jvm', kts: 'jvm', scala: 'jvm', swift: 'jvm',
  c: 'c', h: 'c', cc: 'c', cpp: 'c', cxx: 'c', hpp: 'c', hh: 'c',
  cs: 'dotnet',
  php: 'php', pl: 'ruby', lua: 'ruby',
  sh: 'shell', bash: 'shell', zsh: 'shell',
  sql: 'sql', graphql: 'sql', proto: 'c',
  vue: 'markup', svelte: 'markup',
}

/** Name of the declaration on an anchor line, when one can be read off it. */
const SYMBOL_PATTERNS = [
  /\b(?:function|class|interface|type|enum|namespace|def|func|fn|struct|trait|impl|mod|const|let|var|record|object|fun|val|module|macro_rules!)\s+([A-Za-z_$][\w$]*)/,
  /\b([A-Za-z_$][\w$]*)\s*\([^)]*\)\s*(?:const\s*)?(?:\{|:)?\s*$/,
]

/** The language family for a path, or `undefined` when it has none. */
export function languageOf(path) {
  const name = path.slice(path.lastIndexOf('/') + 1).toLowerCase()
  const dot = name.lastIndexOf('.')
  if (dot <= 0) return undefined
  return EXTENSION_FAMILY[name.slice(dot + 1)]
}

/** The declaration regex for a family, or `undefined`. */
export function declarationPattern(family) {
  return family === undefined ? undefined : FAMILIES[family]
}

/** The name declared on a line, when one reads off it. */
export function symbolOf(line) {
  for (const pattern of SYMBOL_PATTERNS) {
    const match = pattern.exec(line)
    if (match !== null) return match[1]
  }
  return undefined
}

/** Declaration names for a freshly built segment. */
function symbolList(symbol) {
  return symbol === undefined ? [] : [symbol]
}

/**
 * The label for a chunk that may have absorbed several declarations.
 *
 * A small file merges into a single chunk, and that chunk's label is then the
 * file's only description in the index — so dropping every name but the first
 * would describe a whole module as whatever its first declaration happened to
 * be. The list is capped because beyond a handful of names the label stops
 * being a description and starts being noise that dilutes the embedding.
 */
export function symbolLabel(symbols) {
  const unique = [...new Set(symbols)]
  if (unique.length === 0) return undefined
  if (unique.length <= 4) return unique.join(', ')
  return `${unique.slice(0, 4).join(', ')} +${unique.length - 4}`
}

/**
 * Whether a file looks machine-generated.
 *
 * A minified bundle is one enormous line; embedding windows taken from it
 * produces chunks of arbitrary identifiers that match every query weakly and
 * none strongly. The mean line length separates the two cheaply and without
 * reading the whole file.
 */
export function isLikelyMinified(content) {
  const lines = content.split('\n')
  if (lines.length === 0) return true
  // Only consider files large enough for the statistic to mean anything.
  if (lines.length < 4) {
    return content.length > 20_000
  }
  const total = content.length
  return total / lines.length > MINIFIED_MEAN_LINE
}

/** Trim trailing blank lines from a segment, leaving at least one line. */
function trimEnd(lines, start, end) {
  let last = end
  while (last > start && lines[last].trim() === '') last -= 1
  return last
}

/** Split one oversized segment into overlapping windows. */
function splitLong(lines, start, end, maxLines, overlap) {
  const out = []
  const step = Math.max(1, maxLines - overlap)
  for (let at = start; at <= end; at += step) {
    const stop = Math.min(end, at + maxLines - 1)
    out.push({ startLine: at, endLine: stop, symbol: undefined })
    if (stop === end) break
  }
  return out
}

/**
 * Split a file into chunks.
 *
 * @param {object} input
 * @param {string} input.path workspace-relative path, used only for language detection
 * @param {string} input.content file text
 * @param {number} [input.chunkLines] target chunk size
 * @param {number} [input.maxChunkLines] hard cap before splitting
 * @param {number} [input.overlapLines] lines shared when splitting
 * @param {number} [input.maxChunks] cap on chunks returned for one file
 * @returns {{chunks: {startLine: number, endLine: number, symbol?: string, text: string}[], skipped?: string}}
 */
export function chunkFile(input) {
  const {
    path,
    content,
    chunkLines = DEFAULT_CHUNK_LINES,
    maxChunkLines = DEFAULT_MAX_CHUNK_LINES,
    overlapLines = DEFAULT_OVERLAP_LINES,
    maxChunks = DEFAULT_MAX_CHUNKS_PER_FILE,
  } = input

  if (isLikelyMinified(content)) return { chunks: [], skipped: 'generated-or-minified' }
  if (content.trim() === '') return { chunks: [], skipped: 'empty' }

  const lines = content.split('\n')
  const family = languageOf(path)
  const pattern = declarationPattern(family)

  /** Anchors: line indices that begin a declaration. */
  const anchors = []
  if (pattern !== undefined) {
    for (let at = 0; at < lines.length; at += 1) {
      const line = lines[at]
      if (line === '' || /^\s/.test(line)) continue
      if (pattern.test(line)) anchors.push(at)
    }
  }

  /** Segments: [anchor, next anchor) with trailing blanks trimmed. */
  const segments = []
  // One anchor is still structure: a module holding a single top-level function
  // should carry that function's name rather than lose it to the fallback.
  if (anchors.length >= 1) {
    if (anchors[0] > 0) anchors.unshift(0)
    for (let at = 0; at < anchors.length; at += 1) {
      const start = anchors[at]
      const rawEnd = at + 1 < anchors.length ? anchors[at + 1] - 1 : lines.length - 1
      const end = trimEnd(lines, start, rawEnd)
      if (end < start) continue
      segments.push({ startLine: start, endLine: end, symbols: symbolList(symbolOf(lines[start])) })
    }
  } else {
    // No usable structure: one segment, and the splitter below does the work.
    segments.push({ startLine: 0, endLine: trimEnd(lines, 0, lines.length - 1), symbols: [] })
  }

  /**
   * Merge neighbours that fit together.
   *
   * A file of one-line getters would otherwise produce one chunk per getter,
   * each too small to carry meaning. Merging accumulates every declaration name
   * it swallowed, because a chunk tagged only with its first name would be
   * embedded under a label describing a fraction of its contents.
   */
  const merged = []
  let current
  for (const segment of segments) {
    if (current !== undefined) {
      const combined = segment.endLine - current.startLine + 1
      if (combined <= chunkLines) {
        current.endLine = segment.endLine
        current.symbols.push(...segment.symbols)
        continue
      }
    }
    current = { startLine: segment.startLine, endLine: segment.endLine, symbols: [...segment.symbols] }
    merged.push(current)
  }

  const out = []
  for (const segment of merged) {
    const label = symbolLabel(segment.symbols)
    const length = segment.endLine - segment.startLine + 1
    if (length <= maxChunkLines) {
      out.push({ startLine: segment.startLine, endLine: segment.endLine, symbol: label })
      continue
    }
    for (const piece of splitLong(lines, segment.startLine, segment.endLine, maxChunkLines, overlapLines)) {
      out.push({ ...piece, symbol: piece.startLine === segment.startLine ? label : undefined })
    }
  }

  const capped = out.length > maxChunks
  const kept = capped ? out.slice(0, maxChunks) : out
  return {
    chunks: kept.map((chunk) => ({
      // 1-based, so a result can be handed straight to `read` as offset/limit.
      startLine: chunk.startLine + 1,
      endLine: chunk.endLine + 1,
      ...(chunk.symbol === undefined ? {} : { symbol: chunk.symbol }),
      text: lines.slice(chunk.startLine, chunk.endLine + 1).join('\n'),
    })),
    ...(capped ? { skipped: 'chunk-cap-reached' } : {}),
  }
}

/**
 * The exact string that gets embedded for a chunk.
 *
 * The path is included on purpose. A developer asking "where is auth handled"
 * means something closer to "under a path called auth" than to any particular
 * line of code, and the path is the only signal in the chunk that carries that.
 * The symbol follows for the same reason: it is what the question would name.
 */
export function embedText(path, chunk) {
  const header = chunk.symbol === undefined ? path : `${path} :: ${chunk.symbol}`
  return `${header}\n${chunk.text}`
}