workon/

stack.rs

//! Stacked diff workflow support.
//!
//! This module provides the infrastructure for detecting and interacting with
//! stacked-diff tools alongside git-workon's worktree management.
//!
//! ## Model
//!
//! Stack awareness is two-dimensional:
//!
//! - [`StackModel`] — which tool manages stacks (v1: Graphite and `gh stack`, plus
//!   metadata-less git-inference via [`StackModel::Git`]; future: branchless, sapling)
//! - [`Granularity`] — how worktrees map to stacks (v1: [`Granularity::Stack`], one per stack)
//!
//! ## Default-on behavior
//!
//! When `workon.stackModel` resolves to anything other than [`StackModel::None`] (by explicit
//! config or auto-detection), every command that has a meaningful stack-aware variant uses it
//! by default. A `--no-stack` CLI flag (global across subcommands) downgrades any single
//! invocation back to branch-flat behavior.
//!
//! ## Graphite (v1)
//!
//! Stack metadata is read directly from `refs/branch-metadata/*` git refs — blobs containing
//! JSON written by the `gt` CLI. No `gt` process is needed for detection or visualization.
//! `gt track` is invoked only when registering a new branch (in `workon new` when creating a
//! fork off a stack-worktree's branch).
//!
//! ## gh-stack (v1)
//!
//! Stack metadata is read from the `gh stack` extension's own JSON file — canonical at
//! `<common-dir>/gh-stack`, with a degraded union fallback for unlinked per-worktree copies.
//! See `stack/gh_stack.rs`'s module docs for the read order, the dedupe rule, and why a
//! truncated read is tolerated rather than fatal (the opposite of Graphite's rule below).
//!
//! ## Git-inference (v1)
//!
//! [`StackModel::Git`] covers repositories with no stack tool at all: [`enumerate_stacks`] and
//! [`current_stack`] treat it exactly like [`StackModel::None`] (no branch-level topology to
//! report), but [`crate::assemble_changesets`] gives it meaning by walking `upstream..HEAD`,
//! emitting one changeset per commit. It exists purely as an assembly-time model — see
//! [`StackModel::detect`] for why it is never auto-detected.
//!
//! ## Cross-worktree grouping
//!
//! [`group_by_stack`] partitions a parallel slice of [`Option<Stack>`] values (one per worktree)
//! into [`StackGrouping`]: a list of [`StackGroup`]s (each covering one connected stack) plus an
//! `ungrouped` list for trunk/untracked worktrees. The grouping key is `(trunk, sorted diff
//! set)`, so two worktrees in the same connected stack always collapse into one group regardless
//! of the [`Stack::diffs`] vector order returned by [`current_stack`].

pub(crate) mod gh_stack;
pub(crate) mod graphite;
pub(crate) mod metadata;

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

use git2::Repository;

use crate::error::Result;

/// Which stacked-diff tool is managing stacks in this repository.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StackModel {
    /// No stack tool; today's branch-flat behavior.
    None,
    /// Graphite (`gt`) manages stacks via `refs/branch-metadata/*`.
    Graphite,
    /// `gh stack` (the `github/gh-stack` extension) manages stacks via a JSON file. See
    /// `stack/gh_stack.rs`'s module docs for the canonical-file-plus-symlinks model.
    GhStack,
    /// No stack-metadata tool; changesets are inferred purely from git, one per commit in
    /// `upstream..HEAD`. Unlike [`StackModel::Graphite`]/[`StackModel::GhStack`], this carries
    /// no branch-level stack topology: [`enumerate_stacks`] and [`current_stack`] treat it as
    /// flat (same as [`StackModel::None`]), since there is no metadata to enumerate stacks
    /// from or to group branches into. Only [`crate::assemble_changesets`] (M1 changeset
    /// assembly) gives this variant meaning, walking `upstream..HEAD` per-commit.
    ///
    /// Not reachable via [`StackModel::detect`] — see the module docs' Default-on-behavior
    /// note on why auto-detection never resolves to `Git`.
    Git,
}

impl StackModel {
    /// Auto-detect the active stack model from the repository environment.
    ///
    /// Returns [`StackModel::Graphite`] when the repo has been initialized with `gt init`
    /// (`.graphite_repo_config` or `.graphite_metadata.db` exists). Otherwise returns
    /// [`StackModel::None`].
    ///
    /// Deliberately does NOT consult [`graphite::detect_gt`]: the repo's own metadata is the
    /// ground truth for "is this a Graphite stack," and reading it is pure libgit2 (see the
    /// module docs — no `gt` process is needed for detection or visualization). Gating on the
    /// binary's presence would report `None` for a genuine Graphite stack whenever `gt` happens
    /// to be missing from PATH, silently emptying the review TUI's stack. Whether `gt` can be
    /// *executed* is a separate question, and belongs to the call sites that execute it.
    ///
    /// Never returns [`StackModel::Git`]: auto-detection only distinguishes "a stack tool is
    /// active" from "no stack tool," since CLI routing treats any non-`None` model as
    /// stack-active. Auto-resolving to `Git` for every repository with an upstream-tracking
    /// branch would silently flip that routing for nearly every user. `Git` is reachable via
    /// explicit `workon.stackModel = git` config, or a caller mapping `None` to `Git` before
    /// calling [`crate::assemble_changesets`] (the review crate does this from M3 onward).
    ///
    /// **Graphite wins** when both tools' artifacts are present: `.graphite_repo_config`
    /// comes from an explicit, repo-wide `gt init`, while a `gh-stack` file can appear as a
    /// side effect of one `gh stack add` run in one worktree. The more deliberate,
    /// repo-scoped signal wins, so no repo that resolves to `Graphite` today can silently flip
    /// to `GhStack` just because someone tried the other tool once. The escape hatch is an
    /// explicit `workon.stackModel = gh-stack`.
    pub fn detect(repo: &Repository) -> Self {
        if graphite::is_graphite_repo(repo) {
            Self::Graphite
        } else if gh_stack::is_gh_stack_repo(repo) {
            Self::GhStack
        } else {
            Self::None
        }
    }
}

/// How worktrees map to stacks.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Granularity {
    /// One worktree hosts an entire stack. The user navigates between branches inside it
    /// using the stack tool's own commands (e.g. `gt up` / `gt down`).
    Stack,
}

/// A stack of branches rooted at a trunk branch.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Stack {
    /// The trunk branch this stack is rooted on (e.g., `"main"`).
    pub trunk: String,
    /// All non-trunk diffs (branches) in the stack, in BFS order from bottom to top.
    pub diffs: Vec<String>,
    /// The branch that is currently HEAD in the worktree.
    pub current: String,
    /// Parent map for the diffs in this stack: `diff → parent_branch`.
    ///
    /// The parent may be another diff or the trunk itself. Only covers diffs in
    /// [`Stack::diffs`]; the trunk's own parent is not recorded. Used to reconstruct
    /// the branching tree structure (a flat [`diffs`] list cannot represent forks).
    pub parents: HashMap<String, String>,
    /// Provider-assigned stack number (e.g. gh-stack's `stacks[].number`), for display only.
    ///
    /// This is **display metadata, never an identity key** — [`group_by_stack`] keys on
    /// `(trunk, sorted diff set)` and does not consult this field. Always `None` for
    /// [`StackModel::Graphite`] and [`StackModel::Git`], which have no numbering concept.
    pub number: Option<u64>,
}

/// Return all stacks present in metadata, one per connected component.
///
/// Each returned [`Stack`] corresponds to one potential [`StackGroup`] — the same `(trunk,
/// sorted diffs)` key used by [`group_by_stack`]. Used by the `list` command to surface
/// stacks that have no checked-out worktrees.
pub fn enumerate_stacks(repo: &Repository, model: StackModel) -> Result<Vec<Stack>> {
    match model {
        StackModel::None => Ok(vec![]),
        StackModel::Graphite => graphite::enumerate_stacks(repo).map_err(Into::into),
        StackModel::GhStack => gh_stack::enumerate_stacks(repo).map_err(Into::into),
        // Git-inference has no branch-level stack topology to enumerate — flat, like None.
        StackModel::Git => Ok(vec![]),
    }
}

/// Return the stack for the worktree whose HEAD is `head_branch`, or `None` if the branch
/// is not part of a tracked stack under `model`.
///
/// The returned [`Stack`] includes all branches reachable from the same stack root, not just
/// the ancestors of `head_branch`, so branching stacks are fully represented.
pub fn current_stack(
    repo: &Repository,
    head_branch: &str,
    model: StackModel,
) -> Result<Option<Stack>> {
    match model {
        StackModel::None => Ok(None),
        StackModel::Graphite => graphite::current_stack(repo, head_branch).map_err(Into::into),
        StackModel::GhStack => gh_stack::current_stack(repo, head_branch).map_err(Into::into),
        // Git-inference has no branch-level stack topology — flat, like None.
        StackModel::Git => Ok(None),
    }
}

/// Returns `true` if `gt` is on PATH and this repository has been Graphite-initialized.
pub fn is_graphite_active(repo: &Repository) -> bool {
    graphite::detect_gt() && graphite::is_graphite_repo(repo)
}

/// Return the first trunk branch name from `.graphite_repo_config`, or `None` if
/// the file is missing, unparseable, or contains no trunk entries.
///
/// Use this when you need to pass `--parent <trunk>` to `gt track` and want to
/// avoid hardcoding `"main"`. Returns `None` rather than a hardcoded fallback so
/// callers can omit `--parent` entirely and let `gt` infer when the trunk is unknown.
pub fn graphite_trunk(repo: &Repository) -> Option<String> {
    graphite::graphite_trunk(repo)
}

/// Plant `gh-stack`/`gh-stack.lock` symlinks for `worktree_name`, pointing at the canonical
/// `<common-dir>/gh-stack` store. Idempotent; never replaces a real file — see
/// [`migrate_worktree`] for that. Callers gate this on `StackModel::GhStack` themselves (see
/// `workon new`'s hook); planting is harmless under any other model, but pointless.
pub fn link_worktree(repo: &Repository, worktree_name: &str) -> Result<()> {
    gh_stack::link_worktree(repo, worktree_name).map_err(Into::into)
}

/// Append `branch` to the canonical gh-stack file's stack that currently ends at
/// `base_branch`. Callers gate this on `StackModel::GhStack` themselves (see `workon new`'s
/// hook), matching [`link_worktree`]. See `stack/gh_stack.rs`'s `register_branch` for the
/// locking, CAS, and target-selection rules.
pub fn register_branch(repo: &Repository, branch: &str, base_branch: &str) -> Result<()> {
    gh_stack::register_branch(repo, branch, base_branch).map_err(Into::into)
}

/// Merge `worktree_name`'s real `gh-stack` file into canonical, then replace it with a
/// symlink, leaving `gh-stack.bak` behind. Reachable only from `doctor --fix` — never
/// automatically. See `stack/gh_stack.rs`'s module docs for the shared-canonical-file model.
pub fn migrate_worktree(repo: &Repository, worktree_name: &str) -> Result<()> {
    gh_stack::migrate_worktree(repo, worktree_name).map_err(Into::into)
}

/// `true` if the repository has been Graphite-initialized (`.graphite_repo_config` or
/// `.graphite_metadata.db` exists), independent of whether `gt` is on PATH. Unlike
/// [`is_graphite_active`], this doesn't gate on `gt`'s presence — `doctor`'s
/// `BothStackToolsDetected` check needs to know whether the repo's *metadata* says Graphite,
/// not whether `gt` is currently executable.
pub fn is_graphite_repo(repo: &Repository) -> bool {
    graphite::is_graphite_repo(repo)
}

/// `true` if this repository has a gh-stack file anywhere workon knows to look — canonical
/// or an unlinked worktree file. Used by `doctor`'s `GhStackNotInitialized` and
/// `BothStackToolsDetected` checks.
pub fn is_gh_stack_repo(repo: &Repository) -> bool {
    gh_stack::is_gh_stack_repo(repo)
}

/// Per-worktree gh-stack link status, for `doctor`'s `GhStackWorktreeNotLinked` check.
pub use gh_stack::LinkStatus as GhStackLinkStatus;

/// Compute [`GhStackLinkStatus`] for `worktree_name`'s `gh-stack` admin-dir path.
pub fn gh_stack_worktree_link_status(repo: &Repository, worktree_name: &str) -> GhStackLinkStatus {
    gh_stack::worktree_link_status(repo, worktree_name)
}

/// gh-stack files (canonical + unlinked worktree copies) that exist but fail to parse or use
/// an unsupported schema version, as `(path, reason)` pairs. For `doctor`'s
/// `GhStackFileUnreadable` check.
pub fn gh_stack_readability_errors(repo: &Repository) -> Vec<(std::path::PathBuf, String)> {
    gh_stack::readability_errors(repo)
        .into_iter()
        .map(|(path, err)| (path, err.to_string()))
        .collect()
}

/// Stack numbers that appear in more than one gh-stack source, only reachable in degraded
/// (union-read) mode. For `doctor`'s `GhStackDivergentStacks` check.
pub fn gh_stack_divergent_stack_numbers(repo: &Repository) -> Vec<u64> {
    gh_stack::divergent_stack_numbers(repo)
}

/// One group of worktrees that share a connected stack, identified by trunk + diff set.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StackGroup {
    /// Representative stack (trunk + diffs in BFS bottom→top order) for the group.
    pub stack: Stack,
    /// Indices into the caller's worktree slice for members of this stack, ordered by
    /// the member branch's position in `stack.diffs` (bottom→top). Members whose
    /// branch isn't found in `stack.diffs` sort last, in original input order.
    pub members: Vec<usize>,
}

/// Result of partitioning a worktree slice by stacks via [`group_by_stack`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StackGrouping {
    /// Groups of worktrees sharing a connected stack, ordered by `(trunk, sorted branch set)`.
    pub groups: Vec<StackGroup>,
    /// Indices of worktrees with no stack (trunk branches, untracked branches), in input order.
    pub ungrouped: Vec<usize>,
}

/// Partition a slice of optional stack descriptors into groups sharing the same connected stack.
///
/// The grouping key is `(trunk, sorted diff set)`, so worktrees in the same connected stack
/// always collapse into one group regardless of the [`Stack::diffs`] vector order returned
/// by [`current_stack`]. Members within each group are ordered by their branch's position in
/// the representative stack's `diffs` list (bottom→top).
///
/// Groups are returned in deterministic order: sorted by `(trunk, sorted branch set)`.
///
/// When all entries are `None` (stack model inactive or `--no-stack`), `groups` is empty and
/// every index appears in `ungrouped` — callers get identical flat-list behavior.
pub fn group_by_stack(stacks: &[Option<Stack>]) -> StackGrouping {
    // BTreeMap provides sorted iteration: (trunk, branch_set) order automatically.
    let mut map: BTreeMap<(String, Vec<String>), (Stack, Vec<usize>)> = BTreeMap::new();
    let mut ungrouped: Vec<usize> = Vec::new();

    for (i, opt) in stacks.iter().enumerate() {
        match opt {
            None => ungrouped.push(i),
            Some(stack) => {
                let mut key_branches = stack.diffs.clone();
                key_branches.sort();
                let key = (stack.trunk.clone(), key_branches);
                let entry = map
                    .entry(key)
                    .or_insert_with(|| (stack.clone(), Vec::new()));
                entry.1.push(i);
            }
        }
    }

    let groups = map
        .into_values()
        .map(|(rep_stack, mut members)| {
            // Sort members by their branch's position in rep_stack.diffs (bottom→top).
            members.sort_by_key(|&idx| {
                let branch = stacks[idx]
                    .as_ref()
                    .map(|s| s.current.as_str())
                    .unwrap_or("");
                rep_stack
                    .diffs
                    .iter()
                    .position(|b| b == branch)
                    .unwrap_or(usize::MAX)
            });
            StackGroup {
                stack: rep_stack,
                members,
            }
        })
        .collect();

    StackGrouping { groups, ungrouped }
}

#[cfg(test)]
mod tests {
    use super::*;
    use git_workon_fixture::prelude::*;

    /// `detect` must resolve `Graphite` from repo metadata alone, with no dependency on whether
    /// `gt` is installed on the host running the tests.
    ///
    /// The "gt absent" half cannot be forced from inside this process: it would mean mutating
    /// `PATH`, which is process-global (`unsafe` under Rust 2024) and would race every other test
    /// in this binary. Only the CLI suite can scrub `PATH` hermetically, by passing it to a child
    /// process (`git-workon/tests/suite/new.rs`'s `path_without_gt_new`). So the guard here is
    /// environmental rather than hermetic — CI has no `gt`, so a reintroduced `detect_gt()` gate
    /// turns this red there. Its value is making that failure say "detection must not require gt"
    /// instead of surfacing as a handful of unrelated review-TUI changeset-count assertions.
    #[test]
    fn detect_resolves_graphite_from_repo_metadata_without_requiring_the_gt_binary() {
        let fixture = FixtureBuilder::new()
            .graphite_config(&["main"])
            .branch_metadata("a", "main")
            .build()
            .unwrap();
        let repo = fixture.repo().unwrap();

        assert_eq!(
            StackModel::detect(repo),
            StackModel::Graphite,
            "a gt-initialized repo is a Graphite stack regardless of PATH"
        );
    }

    #[test]
    fn detect_resolves_gh_stack_from_canonical_file() {
        let fixture = FixtureBuilder::new()
            .gh_stack(None, 1, "main", &["feat-a"])
            .build()
            .unwrap();
        let repo = fixture.repo().unwrap();

        assert_eq!(StackModel::detect(repo), StackModel::GhStack);
    }

    #[test]
    fn detect_prefers_graphite_when_both_tools_artifacts_are_present() {
        let fixture = FixtureBuilder::new()
            .graphite_config(&["main"])
            .branch_metadata("a", "main")
            .gh_stack(None, 1, "main", &["feat-a"])
            .build()
            .unwrap();
        let repo = fixture.repo().unwrap();

        assert_eq!(
            StackModel::detect(repo),
            StackModel::Graphite,
            "Graphite's repo-wide gt init must win over a gh-stack file appearing alongside it"
        );
    }

    fn stack(trunk: &str, diffs: &[&str], current: &str) -> Stack {
        Stack {
            trunk: trunk.to_string(),
            diffs: diffs.iter().map(|s| s.to_string()).collect(),
            current: current.to_string(),
            parents: HashMap::new(),
            number: None,
        }
    }

    #[test]
    fn empty_input_yields_empty_grouping() {
        let result = group_by_stack(&[]);
        assert!(result.groups.is_empty());
        assert!(result.ungrouped.is_empty());
    }

    #[test]
    fn all_none_yields_all_ungrouped() {
        let stacks: Vec<Option<Stack>> = vec![None, None, None];
        let result = group_by_stack(&stacks);
        assert!(result.groups.is_empty());
        assert_eq!(result.ungrouped, vec![0, 1, 2]);
    }

    #[test]
    fn single_stacked_worktree_forms_one_group() {
        let stacks = vec![Some(stack("main", &["feat-a"], "feat-a"))];
        let result = group_by_stack(&stacks);
        assert!(result.ungrouped.is_empty());
        assert_eq!(result.groups.len(), 1);
        assert_eq!(result.groups[0].members, vec![0]);
        assert_eq!(result.groups[0].stack.trunk, "main");
    }

    #[test]
    fn two_worktrees_same_stack_collapse_into_one_group_ordered_bottom_to_top() {
        // Two worktrees in the same 2-branch stack: feat-b is bottom, feat-a is top.
        // Input idx 0 has current="feat-a" (position 1), input idx 1 has current="feat-b" (position 0).
        let stacks = vec![
            Some(stack("main", &["feat-b", "feat-a"], "feat-a")), // top worktree
            Some(stack("main", &["feat-b", "feat-a"], "feat-b")), // bottom worktree
        ];
        let result = group_by_stack(&stacks);
        assert!(result.ungrouped.is_empty());
        assert_eq!(result.groups.len(), 1);
        // Members ordered by branch position: idx 1 (feat-b, pos 0) first, idx 0 (feat-a, pos 1) second.
        assert_eq!(result.groups[0].members, vec![1, 0]);
    }

    #[test]
    fn same_stack_different_branches_vector_order_still_collapses() {
        // The two Stacks have the same set but different vector orderings — must collapse.
        let stacks = vec![
            Some(stack("main", &["feat-a", "feat-b"], "feat-a")),
            Some(stack("main", &["feat-b", "feat-a"], "feat-b")),
        ];
        let result = group_by_stack(&stacks);
        assert_eq!(
            result.groups.len(),
            1,
            "different branch vector order must still collapse"
        );
        assert!(result.ungrouped.is_empty());
    }

    #[test]
    fn two_distinct_stacks_on_same_trunk_stay_separate() {
        let stacks = vec![
            Some(stack("main", &["stack1-a"], "stack1-a")),
            Some(stack("main", &["stack2-a"], "stack2-a")),
        ];
        let result = group_by_stack(&stacks);
        assert_eq!(result.groups.len(), 2);
        assert!(result.ungrouped.is_empty());
    }

    #[test]
    fn mixed_stacked_and_ungrouped() {
        let stacks = vec![
            None, // trunk worktree
            Some(stack("main", &["feat-a", "feat-b"], "feat-a")),
            None, // another ungrouped
            Some(stack("main", &["feat-a", "feat-b"], "feat-b")),
        ];
        let result = group_by_stack(&stacks);
        assert_eq!(result.ungrouped, vec![0, 2]);
        assert_eq!(result.groups.len(), 1);
        // feat-a is position 0, feat-b is position 1 → member idx 1 (current=feat-a) first, then idx 3
        assert_eq!(result.groups[0].members, vec![1, 3]);
    }

    #[test]
    fn groups_ordered_deterministically_by_trunk_then_branch_set() {
        let stacks = vec![
            Some(stack("main", &["z-feat"], "z-feat")),
            Some(stack("dev", &["d-feat"], "d-feat")),
            Some(stack("main", &["a-feat"], "a-feat")),
        ];
        let result = group_by_stack(&stacks);
        assert_eq!(result.groups.len(), 3);
        // Ordered: (dev, [d-feat]), (main, [a-feat]), (main, [z-feat])
        assert_eq!(result.groups[0].stack.trunk, "dev");
        assert_eq!(result.groups[1].stack.diffs, vec!["a-feat"]);
        assert_eq!(result.groups[2].stack.diffs, vec!["z-feat"]);
    }
}