chore(scripts): add branch restack helper

Author: ignatremizov

scripts/restack_branch_refs.py

@@ -0,0 +1,216 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import argparse
+import subprocess
+import sys
+from collections import defaultdict
+from dataclasses import dataclass
+
+
+def git(*args: str, check: bool = True) -> str:
+    result = subprocess.run(
+        ["git", *args],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if check and result.returncode != 0:
+        stderr = result.stderr.strip()
+        raise SystemExit(stderr or f"git {' '.join(args)} failed")
+    return result.stdout.strip()
+
+
+def git_ok(*args: str) -> bool:
+    result = subprocess.run(
+        ["git", *args],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    return result.returncode == 0
+
+
+@dataclass(frozen=True)
+class BranchMove:
+    branch: str
+    old_sha: str
+    new_sha: str
+    distance_from_tip: int
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Restack local branch refs after rebasing the tip branch of a linear stack. "
+            "By default this only prints the inferred mapping."
+        )
+    )
+    parser.add_argument(
+        "--base",
+        default="main",
+        help="Base branch name. Defaults to 'main'.",
+    )
+    parser.add_argument(
+        "--tip",
+        default=None,
+        help="Tip branch name. Defaults to the currently checked out branch.",
+    )
+    parser.add_argument(
+        "--old-tip",
+        default=None,
+        help="Old pre-rebase tip revision. Defaults to '<tip>@{1}'.",
+    )
+    parser.add_argument(
+        "--remote",
+        default="origin",
+        help="Remote used by --push. Defaults to 'origin'.",
+    )
+    parser.add_argument(
+        "--apply",
+        action="store_true",
+        help="Move local branch refs instead of only printing the plan.",
+    )
+    parser.add_argument(
+        "--push",
+        action="store_true",
+        help="Force-push the updated branch refs after applying them.",
+    )
+    return parser.parse_args()
+
+
+def current_branch() -> str:
+    branch = git("symbolic-ref", "--quiet", "--short", "HEAD")
+    if not branch:
+        raise SystemExit("Detached HEAD is not supported for auto-detecting the tip branch.")
+    return branch
+
+
+def first_parent_chain(tip_sha: str) -> list[str]:
+    chain = git("rev-list", "--first-parent", tip_sha).splitlines()
+    if not chain:
+        raise SystemExit(f"Could not read first-parent history for {tip_sha}.")
+    return chain
+
+
+def branch_refs() -> list[tuple[str, str]]:
+    rows = git("for-each-ref", "refs/heads", "--format=%(refname:short)\t%(objectname)").splitlines()
+    refs: list[tuple[str, str]] = []
+    for row in rows:
+        branch, sha = row.split("\t", 1)
+        refs.append((branch, sha))
+    return refs
+
+
+def infer_moves(base: str, tip: str, old_tip_rev: str) -> tuple[list[BranchMove], str, str]:
+    new_tip_sha = git("rev-parse", tip)
+    old_tip_sha = git("rev-parse", old_tip_rev)
+    old_chain = first_parent_chain(old_tip_sha)
+    old_chain_set = set(old_chain)
+
+    candidates: list[tuple[str, str, int]] = []
+    for branch, sha in branch_refs():
+        if branch in {base, tip}:
+            continue
+        if sha not in old_chain_set:
+            continue
+        if not git_ok("merge-base", "--is-ancestor", sha, old_tip_sha):
+            continue
+        distance = int(git("rev-list", "--first-parent", "--count", f"{sha}..{old_tip_sha}"))
+        candidates.append((branch, sha, distance))
+
+    if not candidates:
+        raise SystemExit(
+            f"No stacked branch refs found on the old {tip} first-parent chain ({old_tip_sha[:12]})."
+        )
+
+    groups: dict[str, list[tuple[str, int]]] = defaultdict(list)
+    for branch, sha, distance in candidates:
+        groups[sha].append((branch, distance))
+
+    ordered_groups = sorted(
+        groups.items(),
+        key=lambda item: item[1][0][1],
+        reverse=True,
+    )
+
+    max_distance = ordered_groups[0][1][0][1]
+    new_chain = first_parent_chain(new_tip_sha)
+    if max_distance >= len(new_chain):
+        raise SystemExit(
+            f"New tip {new_tip_sha[:12]} does not have enough first-parent history for distance {max_distance}."
+        )
+
+    moves: list[BranchMove] = []
+    previous_distance: int | None = None
+    for old_sha, branch_entries in ordered_groups:
+        distance = branch_entries[0][1]
+        if previous_distance is not None and distance >= previous_distance:
+            raise SystemExit("Inferred branch stack is not strictly ordered by ancestry.")
+        previous_distance = distance
+        new_sha = new_chain[distance]
+        for branch, _ in sorted(branch_entries):
+            if branch == tip and old_sha == old_tip_sha:
+                continue
+            if git("rev-parse", branch) == new_sha:
+                continue
+            moves.append(
+                BranchMove(
+                    branch=branch,
+                    old_sha=old_sha,
+                    new_sha=new_sha,
+                    distance_from_tip=distance,
+                )
+            )
+
+    if not moves:
+        raise SystemExit("All inferred branch refs already point at the rebased commits.")
+
+    return moves, old_tip_sha, new_tip_sha
+
+
+def print_plan(moves: list[BranchMove], old_tip_sha: str, new_tip_sha: str, tip: str) -> None:
+    print(f"Old {tip} tip: {old_tip_sha[:12]}")
+    print(f"New {tip} tip: {new_tip_sha[:12]}")
+    print("Planned branch updates:")
+    for move in moves:
+        print(
+            f"  {move.branch}: {move.old_sha[:12]} -> {move.new_sha[:12]}"
+            f" (distance {move.distance_from_tip} from {tip})"
+        )
+
+
+def apply_moves(moves: list[BranchMove]) -> None:
+    for move in moves:
+        git("branch", "-f", move.branch, move.new_sha)
+
+
+def push_moves(moves: list[BranchMove], remote: str) -> None:
+    branches = [move.branch for move in moves]
+    git("push", "--force-with-lease", remote, *branches)
+
+
+def main() -> int:
+    args = parse_args()
+    tip = args.tip or current_branch()
+    old_tip_rev = args.old_tip or f"{tip}@{{1}}"
+    moves, old_tip_sha, new_tip_sha = infer_moves(args.base, tip, old_tip_rev)
+
+    print_plan(moves, old_tip_sha, new_tip_sha, tip)
+    if not args.apply and not args.push:
+        print("\nDry run only. Re-run with --apply to move refs, and add --push to update the remote.")
+        return 0
+
+    apply_moves(moves)
+    print("\nUpdated local refs.")
+
+    if args.push:
+        push_moves(moves, args.remote)
+        print(f"Pushed updated refs to {args.remote}.")
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())