corpus.py

"""Ingest a code/text corpus into per-file chunk lists.

Chunks are fixed line-windows (robust, no language parsing). Samples for the JEPA are
sliding windows of T consecutive chunks WITHIN a file -- so the predictive task is
"reconstruct a masked chunk from its neighbours in the same file."
"""

from __future__ import annotations

import functools
import os
from pathlib import Path


@functools.lru_cache(maxsize=None)
def _contains(path_str: str, needle: str) -> bool:
    try:
        return needle in Path(path_str).read_text(encoding="utf-8", errors="ignore")
    except OSError:
        return False


_SKIP_DIRS = {"Intermediate", "Binaries", ".git", "__pycache__", "DerivedDataCache", "Saved"}
_SKIP_SUFFIXES = (".generated.h", ".gen.cpp")  # UnrealHeaderTool auto-generated -- not domain code


def _keep(path: Path) -> bool:
    if any(part in _SKIP_DIRS for part in path.parts):
        return False
    name = path.name
    return not any(name.endswith(s) for s in _SKIP_SUFFIXES)


def iter_files(roots: list[Path], exts: tuple[str, ...]) -> list[Path]:
    files: list[Path] = []
    for root in roots:
        root = Path(root)
        if not root.exists():
            continue
        for ext in exts:
            files.extend(f for f in root.rglob(f"*{ext}") if _keep(f))
    filt = os.environ.get("CARTRIDGE_FILTER")   # keep only API-rich headers (e.g. UFUNCTION) in huge domains like UE
    if filt:
        files = [f for f in files if f.suffix != ".h" or _contains(str(f), filt)]
    cap = _file_cap()                    # optional global cap (env CARTRIDGE_MAXFILES)
    return files[:cap] if cap else files


def _file_cap():
    v = os.environ.get("CARTRIDGE_MAXFILES")
    return int(v) if v else None


def domain_roots() -> list[Path]:
    """Active source roots. Default RiftSuite; env CARTRIDGE_DOMAIN=ue switches to the UE engine sample."""
    return UE_ROOTS if os.environ.get("CARTRIDGE_DOMAIN") == "ue" else RIFTSUITE_ROOTS


def chunk_text(text: str, lines_per_chunk: int = 12, min_chars: int = 40) -> list[str]:
    lines = text.splitlines()
    chunks: list[str] = []
    for i in range(0, len(lines), lines_per_chunk):
        c = "\n".join(lines[i : i + lines_per_chunk]).strip()
        if len(c) >= min_chars:
            chunks.append(c)
    return chunks


def file_chunk_lists(
    roots: list[Path],
    exts: tuple[str, ...],
    lines_per_chunk: int = 12,
    min_file_chunks: int = 1,
    max_files: int | None = None,
) -> list[list[str]]:
    """Return a list (per file) of that file's chunk strings."""
    out: list[list[str]] = []
    files = iter_files(roots, exts)
    for path in files:
        try:
            text = path.read_text(encoding="utf-8", errors="ignore")
        except OSError:
            continue
        chunks = chunk_text(text, lines_per_chunk)
        if len(chunks) >= min_file_chunks:
            out.append(chunks)
        if max_files and len(out) >= max_files:
            break
    return out


# Corpus definitions used by the plugin experiment.
UPROJ = Path(r"C:\Users\layth\Documents\Unreal Projects")
# repos-as-source-of-truth (Aaron 2026-07-18): the Laythros/<bird>-AIE clones are canonical.
# The old MagicReborn\Plugins tree went stale in early July and the cart silently verified
# against it (the fact-553 osprey incident). Globbed so new bird clones join automatically.
_AIE = UPROJ / "RiftSuite58AIE" / "Plugins"
RIFTSUITE_ROOTS = (sorted(p for p in _AIE.iterdir() if p.name.endswith("-AIE"))
                   if _AIE.exists() else [UPROJ / "MagicReborn" / "Plugins" / p
                                          for p in ("RiftVault", "RiftReach", "RiftGAS", "RiftPulse")])
LYRA_ROOTS = [UPROJ / "LyraStarterGame" / "Source", UPROJ / "LyraStarterGame" / "Plugins"]
PYTHON_ROOTS = [Path(__file__).resolve().parent / ".venv" / "Lib" / "site-packages"]
_UE = Path(r"C:\Program Files\Epic Games\UE_5.8\Engine\Source\Runtime")
# the WHOLE Runtime tree (Aaron 2026-07-01: full cart, not a sample) — the UFUNCTION content
# filter (CARTRIDGE_FILTER) is what keeps this to API-rich headers rather than everything
UE_ROOTS = [_UE]

CPP_EXTS = (".h", ".cpp")
PY_EXTS = (".py",)


def _summary(name: str, fcl: list[list[str]]) -> None:
    nchunks = sum(len(f) for f in fcl)
    print(f"  {name:10s} files={len(fcl):4d}  chunks={nchunks:5d}  "
          f"avg_chunks/file={nchunks / max(len(fcl),1):.1f}")


if __name__ == "__main__":
    print("corpus chunk counts (lines_per_chunk=12):")
    _summary("riftsuite", file_chunk_lists(RIFTSUITE_ROOTS, CPP_EXTS))
    _summary("lyra", file_chunk_lists(LYRA_ROOTS, CPP_EXTS, max_files=400))
    _summary("python", file_chunk_lists(PYTHON_ROOTS, PY_EXTS, max_files=400))