From 72c8d6f9152eb4de556bf8dc2d598cb01310edce Mon Sep 17 00:00:00 2001
From: Rebecca Turner <rbt@sent.as>
Date: Sat, 6 Apr 2024 19:39:48 -0700
Subject: [PATCH] Restore previously checked out commit after restack (#57)

Closes #44
---
 src/restack.rs | 40 +++++++++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/src/restack.rs b/src/restack.rs
index 0799621..df633e0 100644
--- a/src/restack.rs
+++ b/src/restack.rs
@@ -28,6 +28,7 @@ const CONTINUE_MESSAGE: &str = "Fix conflicts and then use `git-gr restack conti
 /// TODO: Add versioning?
 #[derive(serde::Deserialize, serde::Serialize, Debug, Clone)]
 pub struct RestackTodo {
+    before: RepositoryState,
     pub graph: DependencyGraph,
     /// Restack steps left to perform.
     steps: VecDeque<Step>,
@@ -112,6 +113,12 @@ impl RestackTodo {
     }
 }
 
+#[derive(serde::Deserialize, serde::Serialize, Debug, Clone)]
+pub struct RepositoryState {
+    change: Option<ChangeNumber>,
+    commit: CommitHash,
+}
+
 #[derive(serde::Deserialize, serde::Serialize, Debug, Clone)]
 struct Step {
     change: ChangeNumber,
@@ -285,6 +292,8 @@ pub fn restack(
 
     fs::remove_file(todo_path(&git)?).into_diagnostic()?;
 
+    let restore = todo.before.clone();
+
     let mut todo = PushTodo::from(todo);
     if todo.is_empty() {
         tracing::info!("Restack completed; no changes");
@@ -304,6 +313,17 @@ pub fn restack(
         tracing::info!("Restack completed but changes have not been pushed; run `git-gr restack push` to sync changes with the remote.");
     }
 
+    let restore_commit = match restore.change {
+        Some(restore_change) => todo
+            .refs
+            .get(&restore_change)
+            .map(|update| &update.new)
+            .unwrap_or(&restore.commit),
+        None => &restore.commit,
+    };
+
+    git.checkout(restore_commit)?;
+
     Ok(())
 }
 
@@ -374,11 +394,25 @@ pub fn create_todo(gerrit: &mut GerritGitRemote, branch: &str) -> miette::Result
         return Err(miette!("Restack todo already exists at `{todo_path}`"));
     }
 
-    let change_id = git
-        .change_id(branch)
-        .wrap_err("Failed to get Change-Id for HEAD")?;
+    let head = git.rev_parse("HEAD")?;
+    let head_change = match git
+        .change_id(&head)
+        .and_then(|change_id| gerrit.get_change(change_id))
+    {
+        Ok(change) => Some(change.number),
+        Err(error) => {
+            tracing::debug!("Failed to get HEAD change ID: {error}");
+            None
+        }
+    };
+
+    let change_id = git.change_id(branch)?;
     let change = gerrit.get_change(change_id)?;
     let mut todo = RestackTodo {
+        before: RepositoryState {
+            change: head_change,
+            commit: head,
+        },
         graph: gerrit.dependency_graph(change.number)?,
         steps: Default::default(),
         refs: Default::default(),