git_stk/stack/
snapshot.rs

1//! Undo support: capture the current stack's branch tips and metadata
2//! before a mutating command rewrites them, and restore that capture on
3//! `git stk undo`. Local only - pushes and platform merges are not
4//! reverted.
5
6use std::sync::atomic::{AtomicBool, Ordering};
7
8use anyhow::{Context, Result};
9use serde_json::{Value, json};
10
11use super::{base_of, branch_and_descendants, parent_of, stack_root};
12use crate::git;
13use crate::style;
14
15const SNAPSHOT_FILE: &str = "stk-undo";
16
17// One snapshot per process: the outermost mutating command captures state;
18// inner calls (sync's restack, merge's sync) must not overwrite it.
19static TAKEN: AtomicBool = AtomicBool::new(false);
20
21/// Record the current stack so `undo` can restore it. The `label` names the
22/// operation being undone. No-ops after the first call in a process, and is
23/// best effort: a snapshot failure never blocks the command itself.
24pub fn take(label: &str) {
25    if TAKEN.swap(true, Ordering::Relaxed) {
26        return;
27    }
28    if let Err(error) = capture(label) {
29        // The command should still run; we just lose undo for it.
30        let _ = error;
31    }
32}
33
34fn capture(label: &str) -> Result<()> {
35    let head = git::current_branch()?;
36    let root = stack_root(&head)?;
37
38    let branches: Vec<Value> = branch_and_descendants(&root)?
39        .into_iter()
40        .map(|branch| {
41            json!({
42                "name": branch,
43                "sha": git::branch_sha(&branch),
44                "parent": parent_of(&branch).ok().flatten(),
45                "base": base_of(&branch).ok().flatten(),
46            })
47        })
48        .collect();
49
50    let snapshot = json!({
51        "label": label,
52        "head": head,
53        "branches": branches,
54    });
55    let path = git::git_path(SNAPSHOT_FILE)?;
56    std::fs::write(&path, snapshot.to_string())
57        .with_context(|| format!("failed to write {path}"))?;
58    Ok(())
59}
60
61/// Restore the most recent snapshot: reset branch tips and metadata to their
62/// pre-mutation state. Refuses on a dirty worktree (it resets the current
63/// branch) and consumes the snapshot so it is one-shot.
64pub fn undo() -> Result<()> {
65    let path = git::git_path(SNAPSHOT_FILE)?;
66    let Ok(contents) = std::fs::read_to_string(&path) else {
67        anyhow::bail!("nothing to undo");
68    };
69    let snapshot: Value = serde_json::from_str(&contents).context("failed to parse undo state")?;
70
71    if super::restack::in_progress() {
72        anyhow::bail!(
73            "a restack is in progress; finish with `git stk continue` or `git stk abort` first"
74        );
75    }
76    if !git::worktree_is_clean()? {
77        anyhow::bail!(
78            "worktree has uncommitted changes; commit or stash them before `git stk undo`"
79        );
80    }
81
82    let label = snapshot["label"].as_str().unwrap_or("the last operation");
83    let head = snapshot["head"].as_str().unwrap_or_default().to_owned();
84    let branches = snapshot["branches"].as_array().cloned().unwrap_or_default();
85
86    let mut restored = 0;
87    for entry in &branches {
88        let name = entry["name"].as_str().unwrap_or_default();
89        if name.is_empty() {
90            continue;
91        }
92
93        // Refs first: recreate deleted branches, rewind moved ones.
94        if let Some(sha) = entry["sha"].as_str() {
95            git::update_ref(name, sha)?;
96        }
97
98        // Then metadata, set or cleared to match the snapshot.
99        restore_config(name, "stkParent", entry["parent"].as_str())?;
100        restore_config(name, "stkBase", entry["base"].as_str())?;
101        restored += 1;
102    }
103
104    // Put HEAD back where it was and sync the worktree to the restored tip
105    // (clean-tree precondition makes this lossless).
106    if !head.is_empty() && git::branch_sha(&head).is_some() {
107        if git::current_branch().ok().as_deref() != Some(&head) {
108            git::checkout(&head)?;
109        }
110        git::reset_hard()?;
111    }
112
113    std::fs::remove_file(&path).ok();
114
115    anstream::println!(
116        "{}",
117        style::success(&format!("undid {label}: restored {restored} branches"))
118    );
119    anstream::println!(
120        "{}",
121        style::dim("local refs and metadata only; pushes and merged reviews are not reverted")
122    );
123    Ok(())
124}
125
126fn restore_config(branch: &str, key: &str, value: Option<&str>) -> Result<()> {
127    let full = format!("branch.{branch}.{key}");
128    match value {
129        Some(value) => git::config_set(&full, value),
130        None => git::config_unset(&full),
131    }
132}
git_stk/stack/snapshot.rs

git_stk/stack/
snapshot.rs
