From 2ea871e750a0999351c3a9d3c43bf1bf62cf47f9 Mon Sep 17 00:00:00 2001
From: Mike Frysinger <vapier@gentoo.org>
Date: Sun, 14 Feb 2021 23:30:02 -0500
Subject: [PATCH] bin: improve rebase helpers

Rewrite rb-all in python to make it faster.

Add a new "catchup" mode which tries to rebase onto the most recent
commit as possible.
---
 .bin/git-rb-all     | 242 ++++++++++++++++++++++++++++++++------------
 .bin/git-rb-catchup | 137 +++++++++++++++++++++++++
 2 files changed, 316 insertions(+), 63 deletions(-)
 create mode 100755 .bin/git-rb-catchup

diff --git a/.bin/git-rb-all b/.bin/git-rb-all
index be05ba6..1d9ded3 100755
--- a/.bin/git-rb-all
+++ b/.bin/git-rb-all
@@ -1,63 +1,179 @@
-#!/bin/bash
-# Helper to rewrite all local branches.
-
-rb_one() {
-	local b=$1
-	shift
-
-	printf "${BRACKET}### ${GREEN}${b}${NORMAL}"
-	if ! git config --local "branch.${b}.merge" >/dev/null; then
-		echo " -> skipping due to missing merge branch"
-	else
-		echo
-		git checkout -q "${b}" || return
-		git rebase "${opts[@]}" | sed -e '/^Fast-forwarded/d' -e "s:^:${BAD}:" -e "s:$:${NORMAL}:"
-		if [[ ${PIPESTATUS[0]} -ne 0 ]] ; then
-			git rebase --abort
-		fi
-	fi
-}
-
-usage() {
-	cat <<EOF
-Usage: rb-all [-q]
-EOF
-	exit 1
-}
-
-main() {
-	local opts=(-q)
-	while [[ $# -ne 0 ]] ; do
-		case $1 in
-		-q|-v|-n|--no-stat|--stat) opts+=( "$1" );;
-		*) usage ;;
-		esac
-		shift
-	done
-
-	[[ -z ${GOOD} ]] && eval "$(bash-colors --env)"
-
-	# Switch to the top dir in case the working dir doesn't exist in every branch.
-	cd "$(git rev-parse --show-toplevel)" || return
-
-	# Skip if rebase is in progress.
-	if [[ -e $(git rev-parse --git-path rebase-merge) || \
-          -e $(git rev-parse --git-path rebase-apply) ]] ; then
-		printf "${BAD}skipping due to active rebase${NORMAL}\n"
-		exit 1
-	fi
-
-	local orig b branches
-	orig=$(git rev-parse --abbrev-ref HEAD) || return
-
-	branches=( $(git for-each-ref --format='%(refname:short)' 'refs/heads/*') )
-	for b in "${branches[@]}" ; do
-		# If it's a branch in another worktree, ignore it.
-		if [[ $(git branch --list "${b}") != "+"* ]] ; then
-			rb_one "${b}" "${opts[@]}"
-		fi
-	done
-
-	git checkout -q "${orig}"
-}
-main "$@"
+#!/usr/bin/env python3
+
+"""Helper to rebase all local branches."""
+
+import argparse
+import os
+from pathlib import Path
+import re
+import subprocess
+import sys
+
+
+PROG = os.path.basename(__file__)
+
+
+class Terminal:
+    """Terminal escape sequences."""
+
+    CSI_PREFIX = '\033['
+    SGR_SUFFIX = 'm'
+    NORMAL = ''
+    BOLD = '1'
+    _BLACK, _RED, _GREEN, _YELLOW, _BLUE, _MAGENTA, _CYAN, _WHITE = range(0, 8)
+    _FG = 30
+    _BG = 40
+    FG_BLACK = str(_FG + _BLACK)
+    FG_RED = str(_FG + _RED)
+    FG_GREEN = str(_FG + _GREEN)
+    FG_YELLOW = str(_FG + _YELLOW)
+    FG_BLUE = str(_FG + _BLUE)
+    FG_MAGENTA = str(_FG + _MAGENTA)
+    FG_CYAN = str(_FG + _CYAN)
+    FG_WHITE = str(_FG + _WHITE)
+
+
+class Color:
+    """Helper colors."""
+
+    _combine = lambda *args: Terminal.CSI_PREFIX + ';'.join(args) + Terminal.SGR_SUFFIX
+    NORMAL = _combine(Terminal.NORMAL)
+    GOOD = _combine(Terminal.FG_GREEN)
+    WARN = _combine(Terminal.FG_YELLOW)
+    BAD = _combine(Terminal.FG_RED)
+    HILITE = _combine(Terminal.BOLD, Terminal.FG_CYAN)
+    BRACKET = _combine(Terminal.BOLD, Terminal.FG_BLUE)
+
+
+def git(args, **kwargs):
+    """Run git."""
+    kwargs.setdefault('check', True)
+    kwargs.setdefault('capture_output', True)
+    kwargs.setdefault('encoding', 'utf-8')
+    #print('+', 'git', *args)
+    return subprocess.run(['git'] + args, **kwargs)
+
+
+def rebase_inprogress():
+    """Determine whether a rebase is already in progress."""
+    output = git(['rev-parse', '--git-path', 'rebase-merge']).stdout.strip()
+    if Path(output).exists():
+        return True
+
+    output = git(['rev-parse', '--git-path', 'rebase-apply']).stdout.strip()
+    if Path(output).exists():
+        return True
+
+    return False
+
+
+class AppendOption(argparse.Action):
+    """Append the command line option (with no arguments) to dest.
+
+    parser.add_argument('-b', '--barg', dest='out', action='append_option')
+    options = parser.parse_args(['-b', '--barg'])
+    options.out == ['-b', '--barg']
+    """
+
+    def __init__(self, option_strings, dest, **kwargs):
+        if 'nargs' in kwargs:
+            raise ValueError('nargs is not supported for append_option action')
+        super().__init__(option_strings, dest, nargs=0, **kwargs)
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        if getattr(namespace, self.dest, None) is None:
+            setattr(namespace, self.dest, [])
+        getattr(namespace, self.dest).append(option_string)
+
+
+def get_parser():
+    """Get CLI parser."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        '--catchup', action='store_true',
+        help='run git-rb-catchup when rebasing')
+    parser.add_argument(
+        '-q', '--quiet', dest='git_options', action=AppendOption, default=['-q'],
+        help='passthru to git rebase')
+    return parser
+
+
+def main(argv):
+    """The main entry point for scripts."""
+    parser = get_parser()
+    opts = parser.parse_args(argv)
+
+    # Skip if rebase is in progress.
+    if rebase_inprogress():
+        print(f'{Color.BAD}{PROG}: skipping due to active rebase{Color.NORMAL}')
+        return 1
+
+    # Switch to the top dir in case the working dir doesn't exist in every branch.
+    topdir = git(['rev-parse', '--show-toplevel']).stdout.strip()
+    os.chdir(topdir)
+
+    # Example output:
+    #  ||m|refs/remotes/origin/master|ahead 2, behind 203
+    # *||master|refs/remotes/origin/master|ahead 1
+    #  |/usr/local/src/gnu/gdb/build/release|release||behind 10
+    #  ||s-stash||
+    state = git(
+        ['for-each-ref', '--format=%(HEAD)|%(worktreepath)|%(refname:short)|%(upstream)|%(upstream:track,nobracket)',
+         'refs/heads/*']).stdout.splitlines()
+
+    curr_state = None
+    branch_width = 0
+    for line in state:
+        head, worktreepath, branch, tracking, ahead_behind = line.split('|')
+        branch_width = max(branch_width, len(branch))
+        if head == '*':
+            curr_state = branch
+    if not curr_state:
+        print('Unable to resolve current branch', file=sys.stderr)
+        return 1
+
+    branches = {}
+    for line in state:
+        head, worktreepath, branch, tracking, ahead_behind = line.split('|')
+
+        # If it's a branch in another worktree, ignore it.
+        if worktreepath and worktreepath != topdir:
+            continue
+
+        print(f'{Color.BRACKET}### {Color.GOOD}{branch:{branch_width}}{Color.NORMAL} ',
+              end='', flush=True)
+        if not tracking:
+            print(f'{Color.WARN}skipping{Color.NORMAL} due to missing merge branch')
+            continue
+
+        m = re.match(r'ahead ([0-9]+)', ahead_behind)
+        ahead = int(m.group(1)) if m else 0
+        m = re.search(r'behind ([0-9]+)', ahead_behind)
+        behind = int(m.group(1)) if m else 0
+        if not behind:
+            print('Up-to-date!')
+            continue
+        elif not ahead:
+            git(['update-ref', f'refs/heads/{branch}', tracking])
+            print('fast forwarded')
+            continue
+
+        print(f'rebasing [{ahead_behind}] ', end='', flush=True)
+        git(['checkout', '-q', branch])
+        if opts.catchup:
+            print()
+            result = git(['rb-catchup'], capture_output=False, check=False)
+        else:
+            result = git(['rebase'] + opts.git_options, check=False)
+            if result.returncode:
+                git(['rebase', '--abort'])
+                print(f'{Color.BAD}failed{Color.NORMAL}\n' + result.stdout.strip())
+            else:
+                print('OK!')
+
+    git(['checkout', '-q', curr_state])
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv[1:]))
diff --git a/.bin/git-rb-catchup b/.bin/git-rb-catchup
new file mode 100755
index 0000000..bb6b910
--- /dev/null
+++ b/.bin/git-rb-catchup
@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+
+"""Helper to automatically rebase onto latest commit possible."""
+
+import argparse
+import subprocess
+import sys
+
+
+def git(args, **kwargs):
+    """Run git."""
+    kwargs.setdefault('check', True)
+    kwargs.setdefault('capture_output', True)
+    kwargs.setdefault('encoding', 'utf-8')
+    return subprocess.run(['git'] + args, **kwargs)
+
+
+def rebase(target):
+    """Try to rebase onto |target|."""
+    try:
+        git(['rebase', target])
+        return True
+    except KeyboardInterrupt:
+        git(['rebase', '--abort'])
+        print('aborted')
+        sys.exit(1)
+    except:
+        git(['rebase', '--abort'])
+        return False
+
+
+def rebase_bisect(lbranch, rbranch, behind):
+    """Try to rebase branch as close to |rbranch| as possible."""
+    def attempt(pos):
+        target = f'{rbranch}~{pos}'
+        print(f'Rebasing onto {target} ', end='', flush=True)
+        print('.', end='', flush=True)
+#        git(['checkout', '-f', target])
+        print('.', end='', flush=True)
+#        git(['checkout', '-f', lbranch])
+        print('. ', end='', flush=True)
+        ret = rebase(target)
+        print('OK' if ret else 'failed')
+        return ret
+
+    #"min" is the latest branch commit while "max" is where we're now.
+    min = 0
+    max = behind
+    old_mid = None
+    while True:
+        mid = min + (max - min) // 2
+        if mid == old_mid or mid < min or mid >= max:
+            break
+        if attempt(mid):
+            max = mid
+        else:
+            min = mid
+        old_mid = mid
+    print('Done')
+
+
+def get_ahead_behind(lbranch, rbranch):
+    """Return number of commits |lbranch| is ahead & behind relative to |rbranch|."""
+    output = git(['rev-list', '--left-right', '--count', f'{lbranch}...{rbranch}']).stdout
+    return [int(x) for x in output.split()]
+
+
+def get_tracking_branch(branch):
+    """Return remote branch that |branch| is tracking."""
+    merge = git(['config', '--local', f'branch.{branch}.merge']).stdout.strip()
+    if not merge:
+        return None
+
+    remote = git(['config', '--local', f'branch.{branch}.remote']).stdout.strip()
+    if remote:
+        if merge.startswith('refs/heads/'):
+            merge = merge[11:]
+        return f'{remote}/{merge}'
+    else:
+        return merge
+
+
+def get_local_branch():
+    """Return the name of the local checked out branch."""
+    return git(['branch', '--show-current']).stdout.strip()
+
+
+def get_parser():
+    """Get CLI parser."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        '--skip-initial-rebase-latest', dest='initial_rebase',
+        action='store_false', default=True,
+        help='skip initial rebase attempt onto the latest branch')
+    parser.add_argument(
+        'branch', nargs='?',
+        help='branch to rebase onto')
+    return parser
+
+
+def main(argv):
+    """The main entry point for scripts."""
+    parser = get_parser()
+    opts = parser.parse_args(argv)
+
+    lbranch = get_local_branch()
+    print(f'Local branch resolved to "{lbranch}"')
+    if not lbranch:
+        print('Unable to resolve local branch', file=sys.stderr)
+        return 1
+
+    if opts.branch:
+        rbranch = opts.branch
+    else:
+        rbranch = get_tracking_branch(lbranch)
+    print(f'Remote branch resolved to "{rbranch}"')
+
+    ahead, behind = get_ahead_behind(lbranch, rbranch)
+    print(f'Branch is {ahead} commits ahead and {behind} commits behind')
+
+    if not behind:
+        print('Up-to-date!')
+    elif not ahead:
+        print('Fast forwarding ...')
+        git(['merge'])
+    else:
+        if opts.initial_rebase:
+            print(f'Trying to rebase onto latest {rbranch} ... ', end='', flush=True)
+            if rebase(rbranch):
+                print('OK!')
+                return 0
+            print('failed; falling back to bisect')
+        rebase_bisect(lbranch, rbranch, behind)
+
+
+if __name__ == '__main__':
+    sys.exit(main(sys.argv[1:]))
-- 
2.39.2