Improve checkouts and handle partial history case

Author: Caleb-T-Owens

crates/but-rebase/src/graph_rebase/commit.rs

@@ -0,0 +1,25 @@
+//! Provides some slightly higher level tools to help with manipulating commits, in preperation for use in the editor.
+
+use anyhow::Result;
+use gix::prelude::ObjectIdExt;
+
+use crate::{
+    commit::{DateMode, create},
+    graph_rebase::Editor,
+};
+
+impl Editor {
+    /// Finds a commit from inside the editor's in memory repository.
+    pub fn find_commit(&self, id: gix::ObjectId) -> Result<but_core::Commit<'_>> {
+        but_core::Commit::from_id(id.attach(&self.repo))
+    }
+
+    /// Writes a commit with correct signing to the in memory repostiory.
+    pub fn write_commit(
+        &self,
+        commit: but_core::Commit<'_>,
+        date_mode: DateMode,
+    ) -> Result<gix::ObjectId> {
+        create(&self.repo, commit.inner, date_mode)
+    }
+}

crates/but-rebase/src/graph_rebase/creation.rs

@@ -1,20 +1,20 @@
 use std::collections::BTreeMap;
 
-use anyhow::Result;
+use anyhow::{Result, bail};
 use but_graph::{Commit, CommitFlags, Graph, Segment};
 use petgraph::Direction;
 
-use crate::graph_rebase::{Edge, Editor, Step, StepGraph, StepGraphIndex};
+use crate::graph_rebase::{Checkouts, Edge, Editor, Step, StepGraph, StepGraphIndex};
 
 /// Provides an extension for creating an Editor out of the segment graph
 pub trait GraphExt {
     /// Creates an editor.
-    fn to_editor(&self) -> Result<Editor>;
+    fn to_editor(&self, repo: &gix::Repository) -> Result<Editor>;
 }
 
 impl GraphExt for Graph {
     /// Creates an editor out of the segment graph.
-    fn to_editor(&self) -> Result<Editor> {
+    fn to_editor(&self, repo: &gix::Repository) -> Result<Editor> {
         // TODO(CTO): Look into traversing "in workspace" segments that are not reachable from the entrypoint
         // TODO(CTO): Look into stopping at the common base
         let entrypoint = self.lookup_entrypoint()?;
@@ -67,14 +67,43 @@ impl GraphExt for Graph {
         let mut steps_for_commits: BTreeMap<gix::ObjectId, StepGraphIndex> = BTreeMap::new();
         let mut graph = StepGraph::new();
 
-        let mut last_inserted = None;
         while let Some(c) = commits.pop() {
-            let mut ni = graph.add_node(Step::Pick { id: c.id });
-
-            for (order, p) in c.parent_ids.iter().enumerate() {
-                if let Some(parent_ni) = steps_for_commits.get(p) {
-                    graph.add_edge(ni, *parent_ni, Edge { order });
-                }
+            let parent_steps = c
+                .parent_ids
+                .iter()
+                .enumerate()
+                .filter_map(|(o, step)| Some((o, steps_for_commits.get(step)?)))
+                .collect::<Vec<_>>();
+
+            let has_no_parents = c.parent_ids.is_empty();
+            let has_missing_parent_steps = parent_steps.len() != c.parent_ids.len();
+            let has_no_parent_steps = parent_steps.is_empty();
+
+            // If we are missing _some_ but not all parents, something has gone wrong.
+            if has_missing_parent_steps && !has_no_parent_steps {
+                bail!(
+                    "Rebase creation has failed. Expected {} parent steps, found {}",
+                    c.parent_ids.len(),
+                    parent_steps.len()
+                );
+            };
+
+            // If the commit has parents in the commit graph, but none of
+            // them are in the graph, this means but-graph did a partial
+            // traversal and we want to preserve the commit as it is.
+            let preserved_parents = if !has_no_parents && has_no_parent_steps {
+                Some(c.parent_ids)
+            } else {
+                None
+            };
+
+            let mut ni = graph.add_node(Step::Pick {
+                id: c.id,
+                preserved_parents,
+            });
+
+            for (order, parent_ni) in parent_steps {
+                graph.add_edge(ni, *parent_ni, Edge { order });
             }
 
             if let Some(refs) = references.get_mut(&c.id) {
@@ -87,14 +116,16 @@ impl GraphExt for Graph {
                 }
             }
 
-            last_inserted = Some(ni);
             steps_for_commits.insert(c.id, ni);
         }
 
         Ok(Editor {
             graph,
             initial_references: references.values().flatten().cloned().collect(),
-            heads: last_inserted.into_iter().collect(),
+            // TODO(CTO): We need to eventually list all worktrees that we own
+            // here so we can `safe_checkout` them too.
+            checkouts: vec![Checkouts::Head],
+            repo: repo.clone().with_object_memory(),
         })
     }
 }

crates/but-rebase/src/graph_rebase/materialize.rs

@@ -1,25 +1,41 @@
 //! Functions for materializing a rebase
-use crate::graph_rebase::rebase::SuccessfulRebase;
+use crate::graph_rebase::{Checkouts, rebase::SuccessfulRebase};
 use anyhow::Result;
-use but_core::worktree::{
-    checkout::{Options, UncommitedWorktreeChanges},
-    safe_checkout,
+use but_core::{
+    ObjectStorageExt as _,
+    worktree::{
+        checkout::{Options, UncommitedWorktreeChanges},
+        safe_checkout_from_head,
+    },
 };
 
 impl SuccessfulRebase {
     /// Materializes a history rewrite
-    pub fn materialize(self, repo: &gix::Repository) -> Result<()> {
+    pub fn materialize(mut self) -> Result<()> {
+        let repo = self.repo.clone();
+        if let Some(memory) = self.repo.objects.take_object_memory() {
+            memory.persist(self.repo)?;
+        }
+
         for checkout in self.checkouts {
-            // TODO(CTO): improve safe_checkout to allow for speculation
-            safe_checkout(
-                checkout.old_head_id,
-                checkout.head_id,
-                repo,
-                Options {
-                    uncommitted_changes: UncommitedWorktreeChanges::KeepAndAbortOnConflict,
-                    skip_head_update: true,
-                },
-            )?;
+            match checkout {
+                Checkouts::Head => {
+                    let head_oid = repo.head_commit()?.id;
+                    if let Some(new_head) = self.commit_mapping.get(&head_oid) {
+                        // If the head has changed (which means it's in the
+                        // commit mapping), perform a safe checkout.
+                        safe_checkout_from_head(
+                            *new_head,
+                            &repo,
+                            Options {
+                                uncommitted_changes:
+                                    UncommitedWorktreeChanges::KeepAndAbortOnConflict,
+                                skip_head_update: true,
+                            },
+                        )?;
+                    }
+                }
+            }
         }
 
         repo.edit_references(self.ref_edits.clone())?;

crates/but-rebase/src/graph_rebase/mod.rs

@@ -8,6 +8,7 @@ mod creation;
 pub mod rebase;
 pub use creation::GraphExt;
 pub mod cherry_pick;
+pub mod commit;
 pub mod materialize;
 pub mod mutate;
 
@@ -21,6 +22,15 @@ pub enum Step {
     Pick {
         /// The ID of the commit getting picked
         id: gix::ObjectId,
+        /// If we are dealing with a sub-graph with an incomplete history, we
+        /// need to represent the bottom most commits in a way that we preserve
+        /// their parents.
+        ///
+        /// If this is Some, the commit WILL NOT be picked onto the parents the
+        /// graph implies but instead on to the parents listed here.
+        ///
+        /// This is intened to be a private API
+        preserved_parents: Option<Vec<gix::ObjectId>>,
     },
     /// Represents applying a reference to the commit found at it's first parent
     Reference {
@@ -31,6 +41,16 @@ pub enum Step {
     None,
 }
 
+impl Step {
+    /// Creates a pick with the expected defaults
+    pub fn new_pick(id: gix::ObjectId) -> Self {
+        Self::Pick {
+            id,
+            preserved_parents: None,
+        }
+    }
+}
+
 /// Used to represent a connection between a given commit.
 #[derive(Debug, Clone)]
 pub(crate) struct Edge {
@@ -52,6 +72,13 @@ pub struct Selector {
     id: StepGraphIndex,
 }
 
+/// Represents places where `safe_checkout` should be called from
+#[derive(Debug, Clone)]
+pub(crate) enum Checkouts {
+    /// The HEAD of the `repo` the editor was created for.
+    Head,
+}
+
 /// Used to manipulate a set of picks.
 #[derive(Debug, Clone)]
 pub struct Editor {
@@ -60,5 +87,8 @@ pub struct Editor {
     /// Initial references. This is used to track any references that might need
     /// deleted.
     initial_references: Vec<gix::refs::FullName>,
-    heads: Vec<StepGraphIndex>,
+    /// Worktrees that we might need to perform `safe_checkout` on.
+    checkouts: Vec<Checkouts>,
+    /// The in-memory repository that the rebase engine works with.
+    repo: gix::Repository,
 }

crates/but-rebase/src/graph_rebase/mutate.rs

@@ -20,7 +20,7 @@ impl Editor {
     /// Get a selector to a particular commit in the graph
     pub fn select_commit(&self, target: gix::ObjectId) -> Option<Selector> {
         for node_idx in self.graph.node_indices() {
-            if let Step::Pick { id } = self.graph[node_idx]
+            if let Step::Pick { id, .. } = self.graph[node_idx]
                 && id == target
             {
                 return Some(Selector { id: node_idx });