cleanup pruning

Author: wbbradley

src/main.rs

@@ -134,7 +134,7 @@ enum Command {
     },
     /// Open the git-stack state file in an editor for manual editing.
     Edit,
-    /// Restack your active branch and all branches in its related stack.
+    /// Restack your active branch onto its parent branch.
     Restack {
         /// The name of the branch to restack.
         #[arg(long, short)]
@@ -145,6 +145,9 @@ enum Command {
         /// Push any changes up to the remote after restacking.
         #[arg(long, short)]
         push: bool,
+        /// Restack all parent branches recursively up to trunk.
+        #[arg(long, short = 'a', default_value_t = false)]
+        all_parents: bool,
     },
     /// Shows the log between the given branch and its parent (git-stack tree) branch.
     Log {
@@ -368,6 +371,7 @@ fn inner_main() -> Result<()> {
             branch,
             fetch,
             push,
+            all_parents,
         }) => {
             let restack_branch = branch.clone().unwrap_or_else(|| current_branch.clone());
             state.try_auto_mount(&git_repo, &repo, &restack_branch)?;
@@ -380,6 +384,7 @@ fn inner_main() -> Result<()> {
                 current_branch,
                 fetch,
                 push,
+                all_parents,
             )
         }
         Some(Command::Mount { parent_branch }) => {
@@ -969,15 +974,33 @@ fn restack(
     orig_branch: String,
     fetch: bool,
     push: bool,
+    all_parents: bool,
 ) -> Result<(), anyhow::Error> {
     let restack_branch = restack_branch.unwrap_or(orig_branch.clone());
 
     if fetch {
         git_fetch()?;
     }
 
+    // Ensure target branch exists locally (check it out from remote if needed)
+    if !git_repo.branch_exists(&restack_branch) {
+        let remote_ref = format!("{DEFAULT_REMOTE}/{restack_branch}");
+        if git_repo.ref_exists(&remote_ref) {
+            run_git(&["checkout", "-b", &restack_branch, &remote_ref])?;
+            println!(
+                "Created local branch {} from remote.",
+                restack_branch.yellow()
+            );
+        } else {
+            bail!(
+                "Branch {} does not exist locally or on remote.",
+                restack_branch
+            );
+        }
+    }
+
     // Find starting_branch in the stacks of branches to determine which stack to use.
-    let plan = state.plan_restack(git_repo, repo, &restack_branch)?;
+    let plan = state.plan_restack(git_repo, repo, &restack_branch, all_parents)?;
 
     tracing::debug!(?plan, "Restacking branches with plan. Checking out main...");
     git_checkout_main(git_repo, None)?;
@@ -1002,6 +1025,7 @@ fn restack(
             if push
                 && !git_repo.shas_match(&format!("{DEFAULT_REMOTE}/{}", branch.name), &branch.name)
             {
+                println!("Pushing branch '{}' to remote...", branch.name);
                 run_git(&[
                     "push",
                     match branch.stack_method {
@@ -1106,7 +1130,7 @@ fn restack(
         "git checkout {} failed",
         restack_branch
     );
-    tracing::info!("Done.");
+    println!("Done.");
     state.refresh_lkgs(git_repo, repo)?;
 
     // Note: PR sync is now handled separately via `git stack sync`

src/state.rs

@@ -288,11 +288,34 @@ impl State {
         git_repo: &GitRepo,
         repo: &str,
         starting_branch: &str,
+        all_parents: bool,
     ) -> Result<Vec<RestackStep<'_>>> {
-        tracing::debug!("Planning restack for {starting_branch}");
-        let trunk = git_trunk(git_repo)?;
-        // Find all the descendents of the starting branch.
-        // Traverse the tree from the starting branch to the root,
+        tracing::debug!("Planning restack for {starting_branch} (all_parents={all_parents})");
+
+        // Single-step mode: only restack the target branch onto its immediate parent
+        if !all_parents {
+            let parent = self
+                .get_parent_branch_of(repo, starting_branch)
+                .ok_or_else(|| {
+                    anyhow!("Branch {starting_branch} not found in the git-stack tree.")
+                })?;
+
+            let branch = self.get_tree_branch(repo, starting_branch).ok_or_else(|| {
+                anyhow!("Branch {starting_branch} not found in the git-stack tree.")
+            })?;
+
+            // Resolve parent ref to local or origin/branch
+            let parent_ref = git_repo
+                .resolve_branch_ref(&parent.name)
+                .unwrap_or_else(|| parent.name.clone());
+
+            return Ok(vec![RestackStep {
+                parent: parent_ref,
+                branch,
+            }]);
+        }
+
+        // All-parents mode: restack the entire ancestry chain
         let mut path: Vec<&Branch> = vec![];
         let repo_state = self
             .repos

src/sync.rs

@@ -183,6 +183,10 @@ pub struct SyncPlan {
     pub local_changes: Vec<LocalChange>,
     pub remote_changes: Vec<RemoteChange>,
     pub warnings: Vec<String>,
+    /// Branches that will be unmounted (for checkout logic if user is on one)
+    pub branches_to_unmount: Vec<String>,
+    /// Branches safe to delete locally (work preserved on remote)
+    pub branches_to_delete: Vec<String>,
 }
 
 impl SyncPlan {
@@ -712,9 +716,10 @@ fn compute_sync_plan(
         }
     }
 
-    // Detect merged PRs and handle unmounting (pull direction)
+    // Detect merged/closed PRs and handle unmounting (pull direction)
     // This runs regardless of push_only/pull_only since it's about reconciling state
     let mut branches_to_unmount: Vec<(String, String)> = Vec::new(); // (branch, repoint_to)
+    let mut branches_to_delete: Vec<String> = Vec::new();
 
     tracing::debug!(
         "Checking for merged PRs. Local branches: {:?}, Closed PRs: {:?}",
@@ -738,14 +743,31 @@ fn compute_sync_plan(
                 closed_pr.number,
                 closed_pr.state
             );
-            if closed_pr.state == RemotePrState::Merged {
-                // This branch's PR was merged - it should be unmounted
+
+            // Any closed PR (merged or just closed) should unmount from git-stack
+            if matches!(closed_pr.state, RemotePrState::Merged | RemotePrState::Closed) {
+                // This branch's PR was merged/closed - it should be unmounted
                 // Children should be repointed to this branch's parent
                 let repoint_to = local_branch
                     .parent
                     .clone()
                     .unwrap_or_else(|| local.trunk.clone());
                 branches_to_unmount.push((branch_name.clone(), repoint_to));
+
+                // Determine if local branch is safe to delete
+                // Safe if: merged, OR (closed AND remote exists AND local is ancestor of remote)
+                let safe_to_delete = if closed_pr.state == RemotePrState::Merged {
+                    true
+                } else {
+                    // Closed but not merged - check if remote has our work
+                    let remote_ref = format!("{}/{}", DEFAULT_REMOTE, branch_name);
+                    git_repo.ref_exists(&remote_ref)
+                        && git_repo.is_ancestor(branch_name, &remote_ref).unwrap_or(false)
+                };
+
+                if safe_to_delete {
+                    branches_to_delete.push(branch_name.clone());
+                }
             }
         }
     }
@@ -953,6 +975,11 @@ fn compute_sync_plan(
         local_changes,
         remote_changes,
         warnings,
+        branches_to_unmount: branches_to_unmount
+            .iter()
+            .map(|(name, _)| name.clone())
+            .collect(),
+        branches_to_delete,
     }
 }
 
@@ -1037,6 +1064,38 @@ fn apply_plan(
     repo_id: &RepoIdentifier,
     plan: &SyncPlan,
 ) -> Result<()> {
+    // If current branch is being unmounted, checkout a safe ancestor first
+    if !plan.branches_to_unmount.is_empty() {
+        let current_branch = git_repo.current_branch().unwrap_or_default();
+        let unmount_set: HashSet<&str> = plan
+            .branches_to_unmount
+            .iter()
+            .map(|s| s.as_str())
+            .collect();
+
+        if unmount_set.contains(current_branch.as_str()) {
+            // Find the first ancestor that's NOT being unmounted
+            let trunk = git_trunk(git_repo)?;
+            let mut safe_branch = trunk.main_branch.clone();
+            let mut current = current_branch.clone();
+
+            while let Some(parent) = state.get_parent_branch_of(repo, &current) {
+                if !unmount_set.contains(parent.name.as_str()) {
+                    safe_branch = parent.name.clone();
+                    break;
+                }
+                current = parent.name.clone();
+            }
+
+            println!(
+                "Branch {} was merged. Switching to {}...",
+                current_branch.yellow(),
+                safe_branch.green()
+            );
+            run_git(&["checkout", &safe_branch])?;
+        }
+    }
+
     // Apply local changes first (checkout, mount, update pr_number)
     for change in &plan.local_changes {
         apply_local_change(git_repo, state, repo, change)?;
@@ -1053,6 +1112,16 @@ fn apply_plan(
     // Save state again if PR numbers were updated
     state.save_state()?;
 
+    // Delete local branches that are safe to delete (work preserved on remote)
+    for branch_name in &plan.branches_to_delete {
+        if git_repo.branch_exists(branch_name) {
+            println!("Deleting local branch {}...", branch_name.yellow());
+            if let Err(e) = run_git(&["branch", "-D", branch_name]) {
+                tracing::warn!("Failed to delete local branch {}: {}", branch_name, e);
+            }
+        }
+    }
+
     Ok(())
 }