patch.py

#!/usr/bin/env python3

"""
Simple binary and text patch utilities.

Binary patches are based on difflib from the standard library.
For those, we use a simple JSON+zlib "copy/add" patch format that reconstructs a target
byte-string from a base byte-string by:
  - copy: (offset, length) slices from the base
  - add: literal bytes (base64-encoded in JSON)

The main use case of this is in agenda.py: when an object is updated, we
compute and store a compressed patch that undoes the change (i.e., it transforms
the new content back to the old content). Then, in materialize.py, when we're
unpacking agenda checkpoints, we reconstruct all versions of each file.

For text files, we have a simple line-based diff format (see apply_text_diff).
These diffs are generated by various LLM workers (e.g., EditorWorker) to
modify programs.
"""

import base64
import json
import zlib
from difflib import SequenceMatcher


def compute_reverse_patch(new_bytes: bytes, old_bytes: bytes) -> bytes:
    """
    Compute a patch that, when applied to `new_bytes`, reconstructs `old_bytes`.

    Returns a zlib-compressed JSON patch (bytes).
    """
    sm = SequenceMatcher(None, new_bytes, old_bytes, autojunk=False)
    ops: list[list] = []
    for tag, i1, i2, j1, j2 in sm.get_opcodes():
        if tag == "equal":
            if i2 > i1:
                ops.append(["copy", i1, i2 - i1])
        elif tag in ("replace", "insert"):
            lit = old_bytes[j1:j2]
            if lit:
                ops.append(["add", base64.b64encode(lit).decode("ascii")])
        elif tag == "delete":
            # We omit regions from new that are not in old by simply not copying them.
            pass

    patch_obj = {
        "algo": "copy-add-v1",
        "ops": ops,
    }
    raw = json.dumps(patch_obj, separators=(",", ":")).encode("utf-8")
    return zlib.compress(raw)


def apply_patch(base: bytes, patch_bytes: bytes) -> bytes:
    """
    Apply a zlib-compressed JSON patch to `base` and return reconstructed bytes.
    """
    patch_obj = json.loads(zlib.decompress(patch_bytes).decode("utf-8"))

    out = bytearray()
    for op in patch_obj.get("ops", []):
        kind = op[0]
        if kind == "copy":
            _, offset, length = op
            out += base[offset:offset + length]
        elif kind == "add":
            _, b64 = op
            out += base64.b64decode(b64)
        else:
            raise ValueError(f"Unknown op {kind}")

    return bytes(out)


def test_bytes_patches() -> None:
    strs = ["a",
            "a\nadd line 2\nnew line 3",
            "a\nchange line 2\nold line 3",
            "a\nold line 3"]

    strs = [s.encode("utf-8") for s in strs]

    p1 = compute_reverse_patch(strs[1], strs[0])
    p2 = compute_reverse_patch(strs[2], strs[1])
    p3 = compute_reverse_patch(strs[3], strs[2])

    patch_history = [p1, p2, p3]
    cur = strs[3]
    prevs = []
    for rev in reversed(patch_history):
        cur = apply_patch(cur, rev)
        prevs.append(cur)

    assert prevs == list(reversed(strs[:-1]))


def apply_text_diff(text: str, diff: str) -> str:
    """
    Apply a simple, line-based diff to `text`.

    Directives (one per line in `diff`):
      - Lines starting and ending with "@@" are anchors. Consecutive anchors form a
        *sequence*: we search forward from the current cursor for consecutive lines
        matching all anchors, then set the cursor to just after the last matched line.
        An empty anchor ("@@@@") is a no-op synchronization point.
      - Lines starting with '=' are "keep" directives: search forward for that line
        and move cursor just after it.
      - Lines starting with '-' are "delete" directives: search forward for that line
        and delete it. Cursor stays at the deletion point.
      - Lines starting with '+' are "add" directives: insert that line at the current
        cursor position and advance cursor past the inserted line.

    Matching uses stripped text of the line (without trailing whitespace/newlines).
    Inserted lines end with a newline.
    """
    lines = text.splitlines(keepends=True)
    cursor = 0

    def line_content(i: int) -> str:
        return lines[i].rstrip("\n")

    def find_forward(target: str, start: int) -> int | None:
        for idx in range(start, len(lines)):
            if line_content(idx).strip() == target.strip():
                return idx
        return None

    def find_sequence(anchors: list[str], start: int) -> int | None:
        """Find first position >= start where all anchors match consecutive lines."""
        for idx in range(start, len(lines) - len(anchors) + 1):
            if all(line_content(idx + j).strip() == anchors[j].strip()
                   for j in range(len(anchors))):
                return idx
        return None

    # Parse diff lines, grouping consecutive anchors into sequences.
    diff_lines = diff.splitlines()
    pending_anchors: list[str] = []

    def flush_anchors():
        nonlocal cursor, pending_anchors
        if not pending_anchors:
            return
        j = find_sequence(pending_anchors, cursor)
        if j is not None:
            cursor = j + len(pending_anchors)
        pending_anchors = []

    for raw in diff_lines:
        if not raw:
            continue
        if raw.startswith("@@") and raw.endswith("@@"):
            anchor = raw[2:-2]
            if anchor:
                pending_anchors.append(anchor)
            # else: empty anchor = no-op
            continue

        # Non-anchor line: flush any pending anchor sequence first.
        flush_anchors()

        op = raw[0]
        payload = raw[1:]
        if payload.startswith(" "):
            payload = payload[1:]

        if op == '=':
            j = find_forward(payload, cursor)
            if j is not None:
                cursor = j + 1
        elif op == '-':
            j = find_forward(payload, cursor)
            if j is not None:
                del lines[j]
                cursor = j
        elif op == '+':
            lines.insert(cursor, payload + "\n")
            cursor += 1

    return "".join(lines)


def compute_text_diff(before: str, after: str) -> str:
    """Compute a text diff in the format expected by apply_text_diff.

    Returns a diff string that, when applied to ``before``, produces ``after``.

    For each change block, emits anchor lines (``@@...@@``) with enough
    preceding context to disambiguate the position, followed by ``-``/``+``
    directives for deleted/inserted lines.

    Context lines are drawn only from the immediately preceding equal region
    so that earlier insertions/deletions don't invalidate the anchor sequence.
    """
    before_lines = before.splitlines(keepends=False)
    after_lines = after.splitlines(keepends=False)

    # autojunk=False prevents SequenceMatcher from treating frequent lines
    # like '}' as junk, which would cause replace ops instead of insert+equal
    # and break round-tripping when the last line has no trailing newline.
    matcher = SequenceMatcher(None, before_lines, after_lines, autojunk=False)
    opcodes = matcher.get_opcodes()

    # Build index: stripped content -> list of line indices in before_lines.
    line_index: dict[str, list[int]] = {}
    for i, line in enumerate(before_lines):
        line_index.setdefault(line.strip(), []).append(i)

    # Track the start of the immediately preceding equal region for each opcode.
    prev_equal_start = 0  # start of the equal region just before the current change

    diff_parts: list[str] = []

    for op_idx, (tag, i1, i2, j1, j2) in enumerate(opcodes):
        if tag == 'equal':
            prev_equal_start = i1
            continue

        if i1 == 0:
            diff_parts.append("@@@@")
        else:
            # Context lines come only from the immediately preceding equal
            # region [prev_equal_start, i1).  This ensures they are still
            # consecutive in the file even after earlier changes are applied.
            max_context = i1 - prev_equal_start

            n_context = 1
            while n_context <= max_context:
                ctx_start = i1 - n_context
                ctx = [before_lines[ctx_start + k].strip()
                       for k in range(n_context)]
                first_key = ctx[0]
                candidates = line_index.get(first_key, [])
                n_matches = sum(
                    1 for c in candidates
                    if c + n_context <= len(before_lines)
                    and all(before_lines[c + k].strip() == ctx[k]
                            for k in range(n_context))
                )
                if n_matches <= 1:
                    break
                n_context += 1

            # Clamp to available context (may still be ambiguous, but
            # apply_text_diff searches forward from cursor which resolves it).
            n_context = min(n_context, max_context)
            ctx_start = i1 - n_context
            for k in range(n_context):
                line = before_lines[ctx_start + k]
                # A blank line would produce @@@@ which is the empty (no-op)
                # anchor.  Use a single space instead, so .strip() matching
                # in apply_text_diff will still match blank lines.
                if not line.strip():
                    line = ' '
                diff_parts.append(f"@@{line}@@")

        if tag == 'replace':
            for line in before_lines[i1:i2]:
                diff_parts.append(f"- {line}")
            for line in after_lines[j1:j2]:
                diff_parts.append(f"+ {line}")
        elif tag == 'delete':
            for line in before_lines[i1:i2]:
                diff_parts.append(f"- {line}")
        elif tag == 'insert':
            for line in after_lines[j1:j2]:
                diff_parts.append(f"+ {line}")

    return "\n".join(diff_parts)


def test_apply_text_diff_basic() -> None:
    # Use the same example constants we surface in prompts
    text = TEXT_BEFORE_EXAMPLE
    diff = TEXT_DIFF_EXAMPLE
    out = apply_text_diff(text, diff)
    assert out == TEXT_AFTER_EXAMPLE

# Example diff that illustrates the simple format (used in prompts)
TEXT_DIFF_EXAMPLE = """
@@ a @@
+ inserted-after-a
@@ b cd @@
- line2
= line3
+ after3
""".strip("\n")

# Example source and result texts corresponding to TEXT_DIFF_EXAMPLE
TEXT_BEFORE_EXAMPLE = "hello\nworld\na\nb cd\nline2\nline3\n"
TEXT_AFTER_EXAMPLE = "hello\nworld\na\ninserted-after-a\nb cd\nline3\nafter3\n"