shockbot/diff.ts

import type { DiffChunk } from "./types.ts";

// ~4 chars per token is a reasonable rough estimate for code
function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4);
}

const SKIP_EXTENSIONS = new Set([
  "lock",
  "sum",
  "map",
  "png",
  "jpg",
  "jpeg",
  "gif",
  "svg",
  "ico",
  "webp",
  "woff",
  "woff2",
  "ttf",
  "eot",
  "pdf",
  "zip",
  "tar",
  "gz",
  "bin",
  "exe",
  "dll",
  "so",
  "dylib",
]);

const SKIP_PATTERNS = [
  /^node_modules\//,
  /^\.git\//,
  /\/generated\//,
  /\.pb\.\w+$/,
  /\.min\.[jt]s$/,
];

function shouldSkip(filename: string): boolean {
  if (!filename.includes(".")) return true; // no extension = likely binary
  const ext = filename.split(".").pop()!.toLowerCase();
  return (
    SKIP_EXTENSIONS.has(ext) || SKIP_PATTERNS.some((p) => p.test(filename))
  );
}

const EXT_TO_LANG: Record<string, string> = {
  ts: "TypeScript",
  tsx: "TypeScript",
  js: "JavaScript",
  jsx: "JavaScript",
  mjs: "JavaScript",
  cjs: "JavaScript",
  py: "Python",
  go: "Go",
  rs: "Rust",
  rb: "Ruby",
  java: "Java",
  kt: "Kotlin",
  kts: "Kotlin",
  swift: "Swift",
  cs: "C#",
  cpp: "C++",
  cc: "C++",
  cxx: "C++",
  hpp: "C++",
  c: "C",
  php: "PHP",
  sh: "Shell",
  bash: "Shell",
  zsh: "Shell",
  yaml: "YAML",
  yml: "YAML",
  json: "JSON",
  md: "Markdown",
  sql: "SQL",
  html: "HTML",
  htm: "HTML",
  css: "CSS",
  scss: "CSS",
  sass: "CSS",
  less: "CSS",
  vue: "Vue",
  svelte: "Svelte",
  toml: "TOML",
  xml: "XML",
  tf: "Terraform",
  hcl: "HCL",
};

function detectLanguage(filename: string): string {
  const ext = filename.split(".").pop()?.toLowerCase() ?? "";
  return EXT_TO_LANG[ext] ?? "plain";
}

export interface DiffFile {
  filename: string;
  language: string;
  hunks: Array<{ header: string; body: string }>;
}

export function parseDiff(rawDiff: string): DiffFile[] {
  console.log("Parsing diff...");
  const files: DiffFile[] = [];

  const sections = rawDiff.split(/^diff --git /m).filter((s) => s.trim());

  for (const section of sections) {
    const lines = section.split("\n");

    // Use +++ / --- lines — more reliable than first line for renames and paths with spaces
    const plusLine = lines.find((l) => l.startsWith("+++ "));
    const minusLine = lines.find((l) => l.startsWith("--- "));

    let filename: string;
    if (plusLine?.startsWith("+++ b/")) {
      filename = plusLine.slice(6);
    } else if (
      plusLine === "+++ /dev/null" &&
      minusLine?.startsWith("--- a/")
    ) {
      filename = minusLine.slice(6); // deleted file — use old path
    } else {
      continue;
    }

    if (shouldSkip(filename)) continue;

    const language = detectLanguage(filename);
    const hunks: DiffFile["hunks"] = [];
    let header = "";
    let bodyLines: string[] = [];

    for (const line of lines) {
      if (line.startsWith("@@")) {
        if (header) hunks.push({ header, body: bodyLines.join("\n") });
        header = line;
        bodyLines = [];
      } else if (header) {
        bodyLines.push(line);
      }
    }
    if (header) hunks.push({ header, body: bodyLines.join("\n") });

    if (hunks.length > 0) files.push({ filename, language, hunks });
  }

  return files;
}

export function chunkFile(file: DiffFile, maxTokens: number): DiffChunk[] {
  console.log(
    `Chunking file ${file.filename} with ${file.hunks.length} hunks...`,
  );
  return file.hunks.map(({ header, body }) => {
    let hunk = body;

    if (estimateTokens(body) > maxTokens) {
      // Truncate from bottom — top of the hunk has the most useful context
      const lines = body.split("\n");
      const charBudget = maxTokens * 4;
      let chars = 0;
      let i = 0;
      while (
        i < lines.length &&
        chars + (lines[i]?.length ?? 0) + 1 <= charBudget
      ) {
        chars += (lines[i]?.length ?? 0) + 1;
        i++;
      }
      hunk = lines.slice(0, i).join("\n");
    }

    return {
      filename: file.filename,
      language: file.language,
      hunk,
      context: header,
    };
  });
}