ask/chunk.ts

import type { Law } from "./corpus.js";

// A section of a state's vehicle code, cut into the pieces the box retrieves.
// The unit is a subsection: the model reads a handful of laws, and a handful
// of whole sections of a parking statute would be a chapter. Every piece is
// lines of the committed source joined with a space, so a test can find each
// one verbatim in its file.

/** Roughly a paragraph: long enough for a whole rule, short enough that eight are a page. */
export const CHUNK_CHARS = 900;

export interface CodeSection {
  /** The state: "MD". */
  code: string;
  /** "16-112". */
  section: string;
  /** "Md. Code, Transp. § 16-112". */
  cite: string;
  /** The section's catchline, where the state publishes one. */
  heading: string;
  /** The day the section was fetched. */
  asOf: string;
  text: string;
}

type Kind = "lower" | "digit";

interface Label {
  kind: Kind;
  /** The ordinal: a=1, b=2; 1=1. */
  value: number;
  text: string;
}

// "(a)", "(b)" or "(1)", "(2)": whichever the state writes its top-level
// subsections in. The production chunker reads a dozen more shapes -- "A.",
// "1-A.", "I.", "(d–1)" -- each added when a state's code filed a rule under
// the wrong label.
const LABEL = /^\(([a-z]{1,2}|\d{1,3})\)/;

function labelOf(line: string): Label | null {
  const match = LABEL.exec(line);
  if (match === null) return null;
  const body = match[1];
  if (/^\d+$/.test(body)) return { kind: "digit", value: Number(body), text: body };
  return { kind: "lower", value: body.charCodeAt(0) - 96, text: body };
}

/** A page's own banner above the section, and a heading line that is only the section's own number. */
function isHeader(line: string, section: string): boolean {
  return /^Article\s+[-–—]\s+\S/.test(line) || line.replace(/[–—‑]/g, "-") === `§${section}.`;
}

export function bodyLines(text: string, section: string): string[] {
  const lines = text.split("\n").map((line) => line.trim()).filter((line) => line.length > 0);
  let start = 0;
  while (start < Math.min(lines.length - 1, 2) && isHeader(lines[start], section)) start += 1;
  return lines.slice(start);
}

/** Lines into pieces of at most CHUNK_CHARS, cut at line boundaries. */
function pieces(lines: string[]): string[][] {
  const out: string[][] = [];
  let current: string[] = [];
  for (const line of lines) {
    if (current.length > 0 && [...current, line].join(" ").length > CHUNK_CHARS) {
      out.push(current);
      current = [];
    }
    current.push(line);
  }
  if (current.length > 0) out.push(current);
  return out;
}

export function chunkSection(section: CodeSection): Law[] {
  const lines = bodyLines(section.text, section.section);
  const first = lines.findIndex((line) => labelOf(line) !== null);
  const top = first < 0 ? null : labelOf(lines[first]);

  // Group the lines by top-level label: a new group opens on the next label
  // of the top kind in sequence, so a nested "(1)" under "(c)" stays inside
  // (c), and "(i)" after "(h)" is the ninth subsection.
  const groups: { label: string | null; lines: string[] }[] = [];
  if (top === null) {
    groups.push({ label: null, lines });
  } else {
    if (first > 0) groups.push({ label: null, lines: lines.slice(0, first) });
    let current = top.value;
    let group = { label: top.text, lines: [lines[first]] };
    for (const line of lines.slice(first + 1)) {
      const label = labelOf(line);
      if (label !== null && label.kind === top.kind && label.value === current + 1) {
        groups.push(group);
        current = label.value;
        group = { label: label.text, lines: [] };
      }
      group.lines.push(line);
    }
    groups.push(group);
  }

  const base = `${section.code}-C-${section.section.replace(/[^A-Za-z0-9.]+/g, "-")}`;
  const chunks: Law[] = [];
  let paragraph = 0;
  for (const group of groups) {
    for (const [n, piece] of pieces(group.lines).entries()) {
      const text = piece.join(" ");
      const suffix = group.label === null ? `p${(paragraph += 1)}` : n === 0 ? group.label : `${group.label}-p${n + 1}`;
      const reference = group.label === null ? "" : `(${group.label})`;
      chunks.push({
        id: `${base}-${suffix}`,
        law: text,
        question: "",
        concept: section.heading,
        citation: `${section.cite}${reference} — '${text}'`,
        kind: "statute",
        asOf: section.asOf,
      });
    }
  }
  return chunks;
}