Add graph overlay function on successful rebase

This function gives us access to what the but-graph would look like after materialization.

Author: Caleb-T-Owens

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

@@ -10,7 +10,10 @@ use std::collections::{BTreeMap, HashMap};
 
 use anyhow::{Context, Result, bail};
 use but_core::RefMetadata;
+use but_graph::init::Overlay;
 use gix::refs::transaction::RefEdit;
+
+use crate::graph_rebase::util::collect_ordered_parents;
 pub mod cherry_pick;
 pub mod commit;
 pub mod materialize;
@@ -211,7 +214,7 @@ pub struct Editor<'ws, 'meta, M: RefMetadata> {
     /// Provides data about how the editor instance was transformed.
     history: RevisionHistory,
     /// A reference to the workspace that the editor was created for.
-    pub workspace: &'ws mut but_graph::projection::Workspace,
+    workspace: &'ws mut but_graph::projection::Workspace,
     /// A reference to the metadata that the editor was created for.
     meta: &'meta mut M,
 }
@@ -235,6 +238,68 @@ pub struct SuccessfulRebase<'ws, 'meta, M: RefMetadata> {
     meta: &'meta mut M,
 }
 
+impl<'ws, 'meta, M: RefMetadata> SuccessfulRebase<'ws, 'meta, M> {
+    /// Returns a preview of what the but-graph will look like after
+    /// materialization.
+    ///
+    /// Any objects referenced in the resulting graph must be accessed via the
+    /// in memory repo owned by the [`Editor`] since they might be in-memory
+    /// only.
+    pub fn overlayed_graph(&self) -> Result<but_graph::Graph> {
+        let dropped_refs = self.ref_edits.iter().filter_map(|edit| match &edit.change {
+            gix::refs::transaction::Change::Delete { .. } => Some(edit.name.clone()),
+            _ => None,
+        });
+        let updated_refs = self.ref_edits.iter().filter_map(|edit| match &edit.change {
+            gix::refs::transaction::Change::Update { new, .. } => Some(gix::refs::Reference {
+                name: edit.name.clone(),
+                target: new.clone(),
+                // TODO(CTO): Peeled is only relevant for symbolic refs?
+                peeled: None,
+            }),
+            _ => None,
+        });
+
+        let Some((entrypoint_id, entrypoint_refname)) = self
+            .checkouts
+            .iter()
+            .filter_map(|checkout| match checkout {
+                Checkout::Head(selector) => {
+                    let selector = self.history.normalize_selector(*selector).ok()?;
+                    let step = &self.graph[selector.id];
+
+                    match step {
+                        Step::None => None,
+                        Step::Pick(Pick { id, .. }) => Some((*id, None)),
+                        Step::Reference { refname } => {
+                            let parents = collect_ordered_parents(&self.graph, selector.id);
+
+                            if let Some(to_reference) = parents.first()
+                                && let Step::Pick(Pick { id, .. }) = self.graph[*to_reference]
+                            {
+                                Some((id, Some(refname.clone())))
+                            } else {
+                                None
+                            }
+                        }
+                    }
+                }
+            })
+            .next()
+        else {
+            bail!("BUG: Tried to construct rebase engine graph overlay with no entrypoints");
+        };
+
+        let overlay = Overlay::default()
+            .with_references(updated_refs)
+            .with_dropped_references(dropped_refs)
+            .with_entrypoint(entrypoint_id, entrypoint_refname);
+        self.workspace
+            .graph
+            .redo_traversal_with_overlay(&self.repo, self.meta, overlay)
+    }
+}
+
 /// The outcome of a materialize
 #[derive(Debug)]
 pub struct MaterializeOutcome<'ws, 'meta, M: RefMetadata> {

crates/but-rebase/tests/rebase/graph_rebase/rebase_identities.rs

@@ -3,7 +3,7 @@
 use anyhow::Result;
 use but_graph::Graph;
 use but_rebase::graph_rebase::Editor;
-use but_testsupport::{graph_workspace, visualize_commit_graph_all};
+use but_testsupport::{graph_tree, graph_workspace, visualize_commit_graph_all};
 
 use crate::utils::{fixture_writable, standard_options};
 
@@ -21,9 +21,15 @@ fn four_commits() -> Result<()> {
 
     let graph = Graph::from_head(&repo, &*meta, standard_options())?.validated()?;
 
-    let mut ws = graph.into_workspace()?;
+    let mut ws = graph.clone().into_workspace()?;
     let editor = Editor::create(&mut ws, &mut *meta, &repo)?;
     let outcome = editor.rebase()?;
+
+    assert_eq!(
+        graph_tree(&graph).to_string(),
+        graph_tree(&outcome.overlayed_graph()?).to_string()
+    );
+
     let outcome = outcome.materialize()?;
 
     assert_eq!(visualize_commit_graph_all(&repo)?, before);
@@ -46,7 +52,7 @@ fn four_commits_with_short_traversal() -> Result<()> {
 
     let options = standard_options().with_hard_limit(4);
     let graph = Graph::from_head(&repo, &*meta, options)?.validated()?;
-    let mut ws = graph.into_workspace()?;
+    let mut ws = graph.clone().into_workspace()?;
 
     insta::assert_snapshot!(graph_workspace(&ws), @"
     ⌂:0:main[🌳] <> ✓!
@@ -58,6 +64,12 @@ fn four_commits_with_short_traversal() -> Result<()> {
 
     let editor = Editor::create(&mut ws, &mut *meta, &repo)?;
     let outcome = editor.rebase()?;
+
+    assert_eq!(
+        graph_tree(&graph).to_string(),
+        graph_tree(&outcome.overlayed_graph()?).to_string()
+    );
+
     let outcome = outcome.materialize()?;
 
     assert_eq!(visualize_commit_graph_all(&repo)?, before);
@@ -83,9 +95,15 @@ fn merge_in_the_middle() -> Result<()> {
 
     let graph = Graph::from_head(&repo, &*meta, standard_options())?.validated()?;
 
-    let mut ws = graph.into_workspace()?;
+    let mut ws = graph.clone().into_workspace()?;
     let editor = Editor::create(&mut ws, &mut *meta, &repo)?;
     let outcome = editor.rebase()?;
+
+    assert_eq!(
+        graph_tree(&graph).to_string(),
+        graph_tree(&outcome.overlayed_graph()?).to_string()
+    );
+
     let outcome = outcome.materialize()?;
 
     assert_eq!(visualize_commit_graph_all(&repo)?, before);
@@ -115,9 +133,15 @@ fn three_branches_merged() -> Result<()> {
 
     let graph = Graph::from_head(&repo, &*meta, standard_options())?.validated()?;
 
-    let mut ws = graph.into_workspace()?;
+    let mut ws = graph.clone().into_workspace()?;
     let editor = Editor::create(&mut ws, &mut *meta, &repo)?;
     let outcome = editor.rebase()?;
+
+    assert_eq!(
+        graph_tree(&graph).to_string(),
+        graph_tree(&outcome.overlayed_graph()?).to_string()
+    );
+
     let outcome = outcome.materialize()?;
 
     assert_eq!(visualize_commit_graph_all(&repo)?, before);

crates/but-workspace/src/branch/move_branch.rs

@@ -48,14 +48,14 @@ pub(super) mod function {
     ///
     /// Returns the in memory update [outcome](Outcome) that can then used for materialisation.
     pub fn tear_off_branch<'ws, 'meta, M: RefMetadata>(
-        mut editor: Editor<'ws, 'meta, M>,
+        editor: Editor<'ws, 'meta, M>,
         subject_branch_name: &FullNameRef,
         stack_id_override: Option<StackId>,
     ) -> anyhow::Result<Outcome<'ws, 'meta, M>> {
-        let Some(source) = editor
-            .workspace
-            .find_segment_and_stack_by_refname(subject_branch_name)
-        else {
+        let sucessful_rebase = editor.rebase()?;
+        let workspace = sucessful_rebase.overlayed_graph()?.into_workspace()?;
+        let mut editor = sucessful_rebase.into_editor();
+        let Some(source) = workspace.find_segment_and_stack_by_refname(subject_branch_name) else {
             bail!(
                 "Couldn't find branch to move in workspace with reference name: {subject_branch_name}"
             );
@@ -64,7 +64,7 @@ pub(super) mod function {
         // We're currently stopping the move branch operations imperatively at this stage, in order to
         // reduce the scope of this first iteration of moving the branches.
         // TODO: Enable and test that we can move branches in any kind of workspace.
-        match &editor.workspace.kind {
+        match &workspace.kind {
             WorkspaceKind::Managed { .. } => {}
             WorkspaceKind::ManagedMissingWorkspaceCommit { .. } => {
                 bail!("Moving branches currently need a workspace commit")
@@ -74,7 +74,7 @@ pub(super) mod function {
             }
         };
 
-        let mut ws_meta = editor.workspace.metadata.clone();
+        let mut ws_meta = workspace.metadata.clone();
 
         let (source_stack, subject_segment) = source;
 
@@ -86,17 +86,16 @@ pub(super) mod function {
             });
         }
 
-        let Some(workspace_head) = editor.workspace.tip_commit().map(|commit| commit.id) else {
+        let Some(workspace_head) = workspace.tip_commit().map(|commit| commit.id) else {
             bail!("Couldn't find workspace head.")
         };
         let head_selector = editor
             .select_commit(workspace_head)
             .context("Failed to find the workspace head in the graph.")?;
 
-        let Some(lower_bound_ref) = editor
-            .workspace
+        let Some(lower_bound_ref) = workspace
             .lower_bound_segment_id
-            .map(|segment_id| &editor.workspace.graph[segment_id])
+            .map(|segment_id| &workspace.graph[segment_id])
             .and_then(|segment| segment.ref_name())
         else {
             bail!("Tearing off a branch requires a workspace common base");
@@ -107,7 +106,13 @@ pub(super) mod function {
             .context("Failed to find target reference in graph.")?;
 
         let (subject_delimiter, children_to_disconnect, parents_to_disconnect) =
-            get_disconnect_parameters(&editor, source_stack, subject_segment, workspace_head)?;
+            get_disconnect_parameters(
+                &editor,
+                &workspace,
+                source_stack,
+                subject_segment,
+                workspace_head,
+            )?;
 
         editor.disconnect_segment_from(
             subject_delimiter.clone(),
@@ -158,24 +163,25 @@ pub(super) mod function {
     ///
     /// Returns an [outcome](Outcome) for potential materialisation.
     pub fn move_branch<'ws, 'meta, M: RefMetadata>(
-        mut editor: Editor<'ws, 'meta, M>,
+        editor: Editor<'ws, 'meta, M>,
         subject_branch_name: &FullNameRef,
         target_branch_name: &FullNameRef,
     ) -> anyhow::Result<Outcome<'ws, 'meta, M>> {
-        let (source, destination) = retrieve_branches_and_containers(
-            editor.workspace,
-            subject_branch_name,
-            target_branch_name,
-        )?;
+        let sucessful_rebase = editor.rebase()?;
+        let workspace = sucessful_rebase.overlayed_graph()?.into_workspace()?;
+        let mut editor = sucessful_rebase.into_editor();
+
+        let (source, destination) =
+            retrieve_branches_and_containers(&workspace, subject_branch_name, target_branch_name)?;
 
-        let Some(workspace_head) = editor.workspace.tip_commit().map(|commit| commit.id) else {
+        let Some(workspace_head) = workspace.tip_commit().map(|commit| commit.id) else {
             bail!("Couldn't find workspace head.")
         };
 
         // We're currently stopping the move branch operations imperatively at this stage, in order to
         // reduce the scope of this first iteration of moving the branches.
         // TODO: Enable and test that we can move branches in any kind of workspace.
-        match &editor.workspace.kind {
+        match &workspace.kind {
             WorkspaceKind::Managed { .. } => {}
             WorkspaceKind::ManagedMissingWorkspaceCommit { .. } => {
                 bail!("Moving branches currently need a workspace commit")
@@ -185,7 +191,7 @@ pub(super) mod function {
             }
         };
 
-        let mut ws_meta = editor.workspace.metadata.clone();
+        let mut ws_meta = workspace.metadata.clone();
 
         let (source_stack, subject_segment) = source;
         let (_, target_segment) = destination;
@@ -197,7 +203,13 @@ pub(super) mod function {
             .context("Failed to find target reference in graph.")?;
 
         let (subject_delimiter, children_to_disconnect, parents_to_disconnect) =
-            get_disconnect_parameters(&editor, &source_stack, &subject_segment, workspace_head)?;
+            get_disconnect_parameters(
+                &editor,
+                &workspace,
+                &source_stack,
+                &subject_segment,
+                workspace_head,
+            )?;
 
         let skip_reconnect_step = source_stack.segments.len() == 1;
         editor.disconnect_segment_from(
@@ -306,6 +318,7 @@ pub(super) mod function {
 /// as well as the right segment delimiter to move.
 fn get_disconnect_parameters<'ws, 'meta, M: RefMetadata>(
     editor: &Editor<'ws, 'meta, M>,
+    workspace: &but_graph::projection::Workspace,
     source_stack: &Stack,
     subject_segment: &StackSegment,
     workspace_head: gix::ObjectId,
@@ -358,7 +371,7 @@ fn get_disconnect_parameters<'ws, 'meta, M: RefMetadata>(
     // graph data, and it's probably the target branch, which is not included in the workspace.
     let graph_base_segment = subject_segment
         .base_segment_id
-        .map(|segment_idx| &editor.workspace.graph[segment_idx]);
+        .map(|segment_idx| &workspace.graph[segment_idx]);
 
     let parents_to_disconnect = if let Some(stack_base_segment) = stack_base_segment {
         // Base segment is part of the source stack.

crates/but-workspace/src/commit/move_commit.rs

@@ -25,15 +25,18 @@ pub(crate) mod function {
     /// The subject commit will be detached from the source segment, and inserted relative
     /// to a given anchor (branch or commit).
     pub fn move_commit<'ws, 'meta, M: RefMetadata>(
-        mut editor: Editor<'ws, 'meta, M>,
+        editor: Editor<'ws, 'meta, M>,
         subject_commit: impl ToCommitSelector,
         anchor: impl ToSelector,
         side: InsertSide,
     ) -> anyhow::Result<SuccessfulRebase<'ws, 'meta, M>> {
+        let sucessful_rebase = editor.rebase()?;
+        let workspace = sucessful_rebase.overlayed_graph()?.into_workspace()?;
+        let mut editor = sucessful_rebase.into_editor();
         let (subject_commit_selector, subject_commit) =
             editor.find_selectable_commit(subject_commit)?;
 
-        let subject = retrieve_commit_and_containers(editor.workspace, &subject_commit)?;
+        let subject = retrieve_commit_and_containers(&workspace, &subject_commit)?;
 
         let (source_stack, source_segment, _) = subject;
 
@@ -54,6 +57,7 @@ pub(crate) mod function {
 
         let parent_to_disconnect = determine_parent_selector(
             &editor,
+            &workspace,
             source_stack,
             source_segment,
             index_of_subject_commit,
@@ -139,6 +143,7 @@ pub(crate) mod function {
     /// and the position of the source segment in the source stack.
     fn determine_parent_selector<'ws, 'meta, M: RefMetadata>(
         editor: &Editor<'ws, 'meta, M>,
+        workspace: &but_graph::projection::Workspace,
         source_stack: &but_graph::projection::Stack,
         source_segment: &but_graph::projection::StackSegment,
         index_of_subject_commit: usize,
@@ -161,7 +166,7 @@ pub(crate) mod function {
             // Look for the base segment in the graph data, as a fallback if there's no stack segment found.
             let graph_base_segment_ref_name = source_segment
                 .base_segment_id
-                .map(|base_segment_id| &editor.workspace.graph[base_segment_id])
+                .map(|base_segment_id| &workspace.graph[base_segment_id])
                 .and_then(|segment| segment.ref_name());
 
             match stack_base_segment_ref_name.or(graph_base_segment_ref_name) {