git-stk 0.9.0

//! Stack metadata: the `branch.<name>.stkParent`/`stkBase` annotations and
//! the structural queries built on them. Navigation lives in [`nav`], the
//! rebase engine in [`restack`].

use std::collections::{BTreeMap, BTreeSet};

use anyhow::{Context, Result, bail};

use crate::git;
use crate::settings;
use crate::style;

mod nav;
mod restack;
mod snapshot;

pub use nav::{
    behind_parent_hint, checkout_bottom, checkout_child, checkout_parent, checkout_top,
    print_all_stacks, print_children, print_parent, print_stack,
};
pub use restack::{abort_restack, continue_restack, restack};
pub use snapshot::{take as snapshot, undo};

const PARENT_KEY: &str = "stkParent";
const BASE_KEY: &str = "stkBase";
/// Marks a branch as the rename of another that still has an open review, so
/// the next submit can replace and close that review.
const RENAMED_FROM_KEY: &str = "stkRenamedFrom";

pub fn create_branch(branch: &str) -> Result<()> {
    let parent = git::current_branch()?;
    // `new` creates the branch; an existing one is an adopt, not a create.
    if git::local_branches()?
        .iter()
        .any(|existing| existing == branch)
    {
        bail!(
            "branch {branch} already exists - adopt it onto {parent} \
             with `git stk adopt {branch} --parent {parent}`"
        );
    }
    git::create_branch(branch)?;
    set_parent(branch, &parent)?;
    record_base(branch, &parent);
    anstream::println!(
        "created {} with parent {}",
        style::branch(branch),
        style::branch(&parent)
    );
    Ok(())
}

/// Insert a new empty branch directly above the current one, moving the
/// current branch's children onto it. The new branch shares the current tip,
/// so descendants stay correctly based; commit to it, then `restack` to
/// replay them. Any uncommitted changes ride onto the new branch, like `new`.
pub fn insert_branch(branch: &str) -> Result<()> {
    ensure_absent(branch)?;
    let current = git::current_branch()?;
    let children = children_of(&current)?;

    snapshot::take("new --insert");
    git::create_branch(branch)?; // off current; leaves us on the new branch
    set_parent(branch, &current)?;
    record_base(branch, &current);
    for child in &children {
        set_parent(child, branch)?;
        record_base(child, branch);
    }

    anstream::println!(
        "inserted {} above {}",
        style::branch(branch),
        style::branch(&current)
    );
    for child in &children {
        anstream::println!(
            "retargeted {} -> {}",
            style::branch(child),
            style::branch(branch)
        );
    }
    Ok(())
}

/// Insert a new empty branch directly below the current one, moving the
/// current branch onto it. Branches from the current branch's parent, so it
/// requires a clean worktree. Commit to it, then `restack`.
pub fn prepend_branch(branch: &str) -> Result<()> {
    ensure_absent(branch)?;
    let current = git::current_branch()?;
    let parent =
        parent_of(&current)?.context("current branch has no stack parent to prepend below")?;
    if !git::worktree_is_clean()? {
        bail!(
            "working tree has uncommitted changes; commit or stash before `git stk new --prepend`"
        );
    }

    snapshot::take("new --prepend");
    git::checkout(&parent)?;
    git::create_branch(branch)?; // off the parent; leaves us on the new branch
    set_parent(branch, &parent)?;
    record_base(branch, &parent);
    set_parent(&current, branch)?;
    record_base(&current, branch);

    anstream::println!(
        "inserted {} between {} and {}",
        style::branch(branch),
        style::branch(&parent),
        style::branch(&current)
    );
    anstream::println!(
        "retargeted {} -> {}",
        style::branch(&current),
        style::branch(branch)
    );
    Ok(())
}

fn ensure_absent(branch: &str) -> Result<()> {
    if git::local_branches()?
        .iter()
        .any(|existing| existing == branch)
    {
        bail!("branch {branch} already exists");
    }
    Ok(())
}

/// The trunk branch: the remote's default branch when known locally,
/// otherwise a conventional name that exists.
pub fn trunk_branch(branches: &[String]) -> Option<String> {
    let remote = settings::remote().unwrap_or_else(|_| settings::DEFAULT_REMOTE.to_owned());
    if let Some(default) = git::remote_default_branch(&remote) {
        return Some(default);
    }

    ["main", "master"]
        .iter()
        .find(|name| branches.iter().any(|branch| branch == *name))
        .map(|name| (*name).to_owned())
}

pub fn adopt_branch(branch: &str, parent: &str) -> Result<()> {
    if branch == parent {
        bail!("a branch cannot be its own stack parent");
    }

    let branches: BTreeSet<_> = git::local_branches()?.into_iter().collect();
    if !branches.contains(branch) {
        bail!("branch {branch} does not exist");
    }
    if !branches.contains(parent) {
        bail!("parent branch {parent} does not exist");
    }

    set_parent(branch, parent)?;
    record_base(branch, parent);
    anstream::println!(
        "attached {} to {}",
        style::branch(branch),
        style::branch(parent)
    );
    Ok(())
}

pub fn detach_branch(branch: Option<&str>) -> Result<()> {
    let branch = branch
        .map(str::to_owned)
        .map_or_else(git::current_branch, Ok)?;
    unset_parent(&branch)?;
    unset_base(&branch)?;
    anstream::println!("detached {}", style::branch(&branch));
    Ok(())
}

/// Rename a branch and keep the stack intact. Git moves the branch's own
/// metadata with the rename; children pointing at the old name are
/// retargeted here.
pub fn rename_branch(old: &str, new: &str, dry_run: bool) -> Result<()> {
    let children = children_for_branch(old)?;

    if !dry_run {
        snapshot::take("rename");
        git::rename_branch(old, new)?;
    }
    anstream::println!(
        "{} {} -> {}",
        if dry_run { "would rename" } else { "renamed" },
        style::branch(old),
        style::branch(new)
    );

    for child in &children {
        if !dry_run {
            set_parent_for_branch(child, new)?;
        }
        anstream::println!(
            "{} {} -> {}",
            if dry_run {
                "would retarget"
            } else {
                "retargeted"
            },
            style::branch(child),
            style::branch(new)
        );
    }
    Ok(())
}

pub fn parent_for_branch(branch: &str) -> Result<Option<String>> {
    parent_of(branch)
}

pub fn children_for_branch(branch: &str) -> Result<Vec<String>> {
    children_of(branch)
}

pub fn set_parent_for_branch(branch: &str, parent: &str) -> Result<()> {
    set_parent(branch, parent)
}

pub fn unset_parent_for_branch(branch: &str) -> Result<()> {
    unset_parent(branch)
}

pub fn base_for_branch(branch: &str) -> Result<Option<String>> {
    base_of(branch)
}

pub fn set_base_for_branch(branch: &str, base: &str) -> Result<()> {
    git::config_set(&base_key(branch), base)
}

pub fn unset_base_for_branch(branch: &str) -> Result<()> {
    unset_base(branch)
}

/// Record that `branch` is the rename of `old`, whose open review the next
/// submit should replace and close.
pub fn set_renamed_from(branch: &str, old: &str) -> Result<()> {
    git::config_set(&renamed_from_key(branch), old)
}

/// The branch `branch` was renamed from, if a replaced review is still pending.
pub fn renamed_from(branch: &str) -> Result<Option<String>> {
    git::config_get(&renamed_from_key(branch))
}

/// Drop the rename marker once its review has been handled.
pub fn clear_renamed_from(branch: &str) -> Result<()> {
    git::config_unset(&renamed_from_key(branch))
}

/// Record the fork point between a branch and its parent (best effort; e.g.
/// unrelated histories have no merge base, which is not an error here).
pub fn record_base(branch: &str, parent: &str) {
    if let Ok(base) = git::merge_base(parent, branch) {
        let _ = git::config_set(&base_key(branch), &base);
    }
}

/// The root of the stack containing `branch` (the base everything sits on).
pub fn stack_root(branch: &str) -> Result<String> {
    let parents = parent_map()?;
    Ok(root_for(branch, &parents))
}

pub fn branch_and_descendants(branch: &str) -> Result<Vec<String>> {
    let parents = parent_map()?;
    let children = children_map(&parents);
    let mut branches = vec![branch.to_owned()];
    collect_descendants(branch, &children, &mut branches);
    Ok(branches)
}

/// Every branch in the stack containing `branch`, parent-first: the line from
/// the stack bottom up through `branch`, plus everything above it. Sibling
/// stacks that share only the trunk are left out - they branch off the trunk
/// separately, not through `branch`. The trunk itself is excluded; an
/// unanchored root stays in (`path_from_root` keeps it).
pub fn stack_line(branch: &str) -> Result<Vec<String>> {
    let mut line = path_from_root(branch)?; // [bottom ..= branch]
    let above = branch_and_descendants(branch)?; // [branch, ..descendants]
    line.extend(above.into_iter().skip(1)); // append above-branch, dropping the duplicate

    // `path_from_root` keeps its starting branch even when that is the trunk
    // (you are standing on it); a trunk is never part of a stack.
    let trunk = trunk_branch(&git::local_branches()?);
    line.retain(|candidate| Some(candidate) != trunk.as_ref());
    Ok(line)
}

/// The stack path from the bottom up to (and including) `branch`,
/// parent-first; descendants above it are left out.
pub fn path_from_root(branch: &str) -> Result<Vec<String>> {
    let trunk = trunk_branch(&git::local_branches()?);
    let mut path = vec![branch.to_owned()];
    let mut seen = BTreeSet::from([branch.to_owned()]);

    let mut cursor = branch.to_owned();
    while let Some(parent) = parent_of(&cursor)? {
        if Some(&parent) == trunk.as_ref() || !seen.insert(parent.clone()) {
            break;
        }
        path.push(parent.clone());
        cursor = parent;
    }

    path.reverse();
    Ok(path)
}

/// (branch, parent) pairs for the branches that have a stack parent;
/// branches without one are skipped.
pub fn branch_parents(branches: &[String]) -> Result<Vec<(String, String)>> {
    let mut pairs = Vec::new();
    for branch in branches {
        if let Some(parent) = parent_of(branch)? {
            pairs.push((branch.clone(), parent));
        }
    }
    Ok(pairs)
}

fn parent_map() -> Result<BTreeMap<String, String>> {
    let mut parents = BTreeMap::new();
    for branch in git::local_branches()? {
        if let Some(parent) = parent_of(&branch)? {
            parents.insert(branch, parent);
        }
    }
    Ok(parents)
}

fn collect_descendants(
    branch: &str,
    children: &BTreeMap<String, Vec<String>>,
    branches: &mut Vec<String>,
) {
    if let Some(branch_children) = children.get(branch) {
        for child in branch_children {
            branches.push(child.to_owned());
            collect_descendants(child, children, branches);
        }
    }
}

fn children_of(parent: &str) -> Result<Vec<String>> {
    Ok(parent_map()?
        .into_iter()
        .filter_map(|(branch, branch_parent)| (branch_parent == parent).then_some(branch))
        .collect())
}

fn children_map(parents: &BTreeMap<String, String>) -> BTreeMap<String, Vec<String>> {
    let mut children: BTreeMap<String, Vec<String>> = BTreeMap::new();
    for (branch, parent) in parents {
        children
            .entry(parent.to_owned())
            .or_default()
            .push(branch.to_owned());
    }
    children
}

fn root_for(branch: &str, parents: &BTreeMap<String, String>) -> String {
    let mut root = branch.to_owned();
    let mut seen = BTreeSet::new();

    while let Some(parent) = parents.get(&root) {
        if !seen.insert(root.clone()) {
            break;
        }
        root = parent.to_owned();
    }

    root
}

fn parent_of(branch: &str) -> Result<Option<String>> {
    git::config_get(&parent_key(branch))
}

fn base_of(branch: &str) -> Result<Option<String>> {
    git::config_get(&base_key(branch))
}

fn set_parent(branch: &str, parent: &str) -> Result<()> {
    git::config_set(&parent_key(branch), parent)
}

fn unset_parent(branch: &str) -> Result<()> {
    git::config_unset(&parent_key(branch))
}

fn unset_base(branch: &str) -> Result<()> {
    git::config_unset(&base_key(branch))
}

fn parent_key(branch: &str) -> String {
    format!("branch.{branch}.{PARENT_KEY}")
}

fn base_key(branch: &str) -> String {
    format!("branch.{branch}.{BASE_KEY}")
}

fn renamed_from_key(branch: &str) -> String {
    format!("branch.{branch}.{RENAMED_FROM_KEY}")
}