Make safe_checkout() fully functional

Author: Byron

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

@@ -45,9 +45,9 @@ pub(crate) mod function {
     use gix::diff::rewrites::tracker::{Change as _, ChangeKind};
     use gix::diff::tree::visit;
     use gix::index::entry::Stage;
-    use gix::object::tree::EntryKind;
+    use gix::merge::tree::TreatAsUnresolved;
     use gix::objs::TreeRefIter;
-    use gix::prelude::ObjectIdExt as _;
+    use gix::prelude::{ObjectIdExt as _, Write};
     use gix::refs::Target;
     use gix::refs::transaction::{Change, LogChange, PreviousValue, RefEdit, RefLog};
     use std::collections::{BTreeSet, VecDeque};
@@ -113,7 +113,7 @@ pub(crate) mod function {
     ) -> anyhow::Result<Outcome> {
         let source_tree = current_head_id.attach(repo).object()?.peel_to_tree()?;
         let new_object = new_head_id.attach(repo).object()?;
-        let destination_tree = new_object.clone().peel_to_tree()?;
+        let mut destination_tree = new_object.clone().peel_to_tree()?;
 
         let mut delegate = Delegate::default();
         gix::diff::tree(
@@ -125,7 +125,8 @@ pub(crate) mod function {
         )?;
 
         let mut opts = git2::build::CheckoutBuilder::new();
-        let snapshot_tree = if !delegate.changed_files.is_empty() {
+        let mut snapshot_tree = None;
+        if !delegate.changed_files.is_empty() {
             let changes = but_core::diff::worktree_changes_no_renames(repo)?;
             if !changes.changes.is_empty() {
                 let actual_head_tree_id = repo.head_tree_id_or_empty()?;
@@ -138,7 +139,7 @@ pub(crate) mod function {
                 // Figure out which added or modified files are actually touched. Deletions we ignore, and allow
                 // these files to be recreated during checkout even if they were part in a rename
                 // (we don't do rename tracking here)
-                let mut change_lut = repo.empty_tree().edit()?;
+                let mut change_lut = tree::Lut::default();
                 for change in &changes.changes {
                     match change.status {
                         TreeStatus::Deletion { .. } => {
@@ -148,27 +149,25 @@ pub(crate) mod function {
                         }
                         TreeStatus::Addition { .. } | TreeStatus::Modification { .. } => {
                             // It's not about the actual values, just to have a lookup for overlapping paths.
-                            change_lut.upsert(
-                                &change.path,
-                                EntryKind::Blob,
-                                repo.object_hash().empty_blob(),
-                            )?;
+                            change_lut.track_file(change.path.as_ref());
                         }
                         TreeStatus::Rename { .. } => {
                             unreachable!("rename tracking was disabled")
                         }
                     }
                 }
 
-                let selection_of_changes_checkout_would_affect = BTreeSet::new();
-                // TODO: find uncommitted that would be overwritten.
-                for _file_to_be_modified in &delegate.changed_files {
-                    // selection.extend(change_lut.extend_leafs(&file_to_be_modified))
+                let mut selection_of_changes_checkout_would_affect = BTreeSet::new();
+                for (_kind, file_to_be_modified) in &delegate.changed_files {
+                    change_lut.get_intersecting(
+                        file_to_be_modified.as_ref(),
+                        &mut selection_of_changes_checkout_would_affect,
+                    );
                 }
 
                 if !selection_of_changes_checkout_would_affect.is_empty() {
-                    let repo_in_memory = repo.clone().with_object_memory();
-                    let _out = crate::snapshot::create_tree(
+                    let mut repo_in_memory = repo.clone().with_object_memory();
+                    let out = crate::snapshot::create_tree(
                         source_tree.id.attach(&repo_in_memory),
                         snapshot::create_tree::State {
                             changes,
@@ -178,20 +177,45 @@ pub(crate) mod function {
                         no_workspace_and_meta(),
                     )?;
 
-                    match uncommitted_changes {
-                        UncommitedWorktreeChanges::KeepAndAbortOnConflict => {}
-                        UncommitedWorktreeChanges::KeepConflictingInSnapshotAndOverwrite => {}
+                    if !out.is_empty() {
+                        let resolve = crate::snapshot::resolve_tree(
+                            out.snapshot_tree.attach(&repo_in_memory),
+                            destination_tree.id,
+                            snapshot::resolve_tree::Options {
+                                worktree_cherry_pick: None,
+                            },
+                        )?;
+                        if let Some(mut worktree_cherry_pick) = resolve.worktree_cherry_pick {
+                            // re-apply snapshot of just what we need and see if they apply cleanly.
+                            match uncommitted_changes {
+                                UncommitedWorktreeChanges::KeepAndAbortOnConflict => {
+                                    let unresolved = TreatAsUnresolved::git();
+                                    if worktree_cherry_pick.has_unresolved_conflicts(unresolved) {
+                                       let paths = worktree_cherry_pick
+                                           .conflicts
+                                           .iter()
+                                           .filter(|c| c.is_unresolved(unresolved)).map(|c| format!("{:?}", c.ours.location()))
+                                           .collect::<Vec<_>>() ;
+                                        bail!("Worktree changes would be overwritten by checkout: {}", paths.join(", "));
+                                    }
+                                }
+                                UncommitedWorktreeChanges::KeepConflictingInSnapshotAndOverwrite => {}
+                            }
+                            destination_tree.id = worktree_cherry_pick.tree.write()?.detach();
+                            if let Some(memory) = repo_in_memory.objects.take_object_memory() {
+                                for (kind, data) in memory.values() {
+                                    repo_in_memory
+                                        .write_buf(*kind, data)
+                                        .map_err(anyhow::Error::from_boxed)?;
+                                }
+                            }
+                        }
+                        snapshot_tree = Some(out.snapshot_tree);
+                        // TODO: deal with index, but to do that it needs to be merged with destination tree!
                     }
-                    todo!("deal with snapshot")
-                } else {
-                    None
                 }
-            } else {
-                None
             }
-        } else {
-            None
-        };
+        }
 
         // Finally, perform the actual checkout
         // TODO(gix): use unconditional `gix` checkout implementation as pre-cursor to the real deal (not needed here).
@@ -277,6 +301,206 @@ pub(crate) mod function {
         })
     }
 
+    mod tree {
+        use bstr::{BStr, BString, ByteVec};
+        use std::collections::{BTreeSet, HashMap};
+
+        /// A lookup table as a tree, with trees represented by their path component.
+        pub struct Lut {
+            nodes: Vec<TreeNode>,
+        }
+
+        impl Default for Lut {
+            fn default() -> Self {
+                Lut {
+                    nodes: vec![TreeNode {
+                        children: Default::default(),
+                    }],
+                }
+            }
+        }
+
+        #[allow(clippy::indexing_slicing)]
+        impl Lut {
+            /// Insert a node for each component in slash-separated `rela_path`.
+            pub fn track_file(&mut self, rela_path: &BStr) {
+                let mut next_index = self.nodes.len();
+                let mut cursor = &mut self.nodes[0];
+                for component in to_components(rela_path) {
+                    match cursor.children.get(component).copied() {
+                        None => {
+                            cursor.children.insert(component.to_owned(), next_index);
+                            self.nodes.push(TreeNode::default());
+                            cursor = &mut self.nodes[next_index];
+                            next_index += 1;
+                        }
+                        Some(existing_idx) => {
+                            cursor = &mut self.nodes[existing_idx];
+                        }
+                    }
+                }
+            }
+
+            /// Given `rela_path`, place all leaf paths into `out` which is partially or fully
+            /// intersecting with `rela_path`.
+            /// Note that all paths in `out` will be `/` separated relative paths.
+            ///
+            /// For example, in a tree with paths `a`, `b/c` and `b/d`, `rela_path` of `a/b` will match `a`,
+            /// while `b` would match `b/c` and `b/d`.
+            pub fn get_intersecting(&self, rela_path: &BStr, out: &mut BTreeSet<BString>) {
+                let mut cur_path = BString::default();
+                let mut cursor = &self.nodes[0];
+
+                for component in to_components(rela_path) {
+                    match cursor.children.get(component).copied() {
+                        None => {
+                            if !cur_path.is_empty() {
+                                out.insert(cur_path.clone());
+                            }
+                            return;
+                        }
+                        Some(existing_idx) => {
+                            cursor = &self.nodes[existing_idx];
+                            push_component(&mut cur_path, component)
+                        }
+                    }
+                }
+
+                if cursor.children.is_empty() {
+                    if !cur_path.is_empty() {
+                        out.insert(cur_path.clone());
+                    }
+                    return;
+                }
+
+                let mut queue: Vec<_> = cursor
+                    .children
+                    .iter()
+                    .map(|(component, idx)| (cur_path.len(), (component, *idx)))
+                    .collect();
+
+                while let Some((cur_path_len, (component, idx))) = queue.pop() {
+                    cur_path.truncate(cur_path_len);
+                    push_component(&mut cur_path, component.as_ref());
+
+                    let node = &self.nodes[idx];
+                    if node.children.is_empty() {
+                        if !cur_path.is_empty() {
+                            out.insert(cur_path.clone());
+                        }
+                    } else {
+                        queue.extend(
+                            node.children
+                                .iter()
+                                .map(|(component, idx)| (cur_path.len(), (component, *idx))),
+                        );
+                    }
+                }
+            }
+        }
+
+        #[derive(Debug, Default, Clone)]
+        struct TreeNode {
+            /// A mapping of path component names (always without slash) to their entry in the [`Lut::nodes`].
+            /// Note that the node is only a leaf if it has no `children` itself, which is when it can be assumed to be a file.
+            children: HashMap<BString, usize>,
+        }
+
+        pub fn to_components(rela_path: &BStr) -> impl Iterator<Item = &BStr> {
+            rela_path.split(|b| *b == b'/').map(Into::into)
+        }
+
+        pub fn push_component(path: &mut BString, component: &BStr) {
+            if !path.is_empty() {
+                path.push_byte(b'/');
+            }
+            path.extend_from_slice(component);
+        }
+
+        #[cfg(test)]
+        mod tests {
+            use super::*;
+
+            #[test]
+            fn journey() {
+                let mut lut = Lut::default();
+                for path in ["a", "b/c", "b/d"] {
+                    lut.track_file(path.into());
+                }
+                let mut out = BTreeSet::new();
+                lut.get_intersecting("".into(), &mut out);
+                assert_eq!(out.len(), 0, "empty means nothing, instead of everything");
+
+                lut.get_intersecting("d".into(), &mut out);
+
+                // This one could not be found.
+                insta::assert_debug_snapshot!(out, @r"{}");
+
+                lut.get_intersecting("a".into(), &mut out);
+                // Perfect match
+                insta::assert_compact_debug_snapshot!(out, @r#"{"a"}"#);
+
+                out.clear();
+                lut.get_intersecting("a/b".into(), &mut out);
+                // indirect match, prefix
+                insta::assert_compact_debug_snapshot!(out, @r#"{"a"}"#);
+
+                out.clear();
+                lut.get_intersecting("b".into(), &mut out);
+                // indirect match, suffix/leafs
+                insta::assert_debug_snapshot!(out, @r#"
+                {
+                    "b/c",
+                    "b/d",
+                }
+                "#);
+            }
+
+            #[test]
+            fn complex_journey() {
+                let mut lut = Lut::default();
+                for path in [
+                    "a/1/2/4/5",
+                    "a/1/2/3",
+                    "a/2/3",
+                    "a/3",
+                    "b/3/4/5",
+                    "b/3/2/1",
+                    "b/2/3",
+                    "b/1",
+                ] {
+                    lut.track_file(path.into());
+                }
+                let mut out = BTreeSet::new();
+
+                lut.get_intersecting("a".into(), &mut out);
+                insta::assert_debug_snapshot!(out, @r#"
+                {
+                    "a/1/2/3",
+                    "a/1/2/4/5",
+                    "a/2/3",
+                    "a/3",
+                }
+                "#);
+
+                // It's additive
+                lut.get_intersecting("b".into(), &mut out);
+                insta::assert_debug_snapshot!(out, @r#"
+                {
+                    "a/1/2/3",
+                    "a/1/2/4/5",
+                    "a/2/3",
+                    "a/3",
+                    "b/1",
+                    "b/2/3",
+                    "b/3/2/1",
+                    "b/3/4/5",
+                }
+                "#);
+            }
+        }
+    }
+
     #[derive(Default)]
     struct Delegate {
         path_deque: VecDeque<BString>,

crates/but-workspace/src/snapshot/resolve_tree.rs

@@ -1,4 +1,5 @@
 /// The information extracted from [`resolve_tree`](function::resolve_tree()).
+#[must_use]
 pub struct Outcome<'repo> {
     /// The cherry-pick result as merge between the target worktree and the snapshot, **possibly with conflicts**.
     ///
@@ -95,7 +96,7 @@ pub(super) mod function {
         let index = match (&head_tree, index, index_conflicts) {
             (_, Some(index_tree), Some(index_conflicts)) => {
                 let (mut index, _path) = repo.index_from_tree(&index_tree.id())?.into_parts();
-                resolve_conflicts(&mut index, index_conflicts.id())?;
+                replace_entries_with_their_restored_conflicts(&mut index, index_conflicts.id())?;
                 Some(index)
             }
             (_, Some(index_tree), None) => {
@@ -104,7 +105,7 @@ pub(super) mod function {
             }
             (Some(worktree_base), None, Some(index_conflicts)) => {
                 let (mut index, _path) = repo.index_from_tree(&worktree_base.id())?.into_parts();
-                resolve_conflicts(&mut index, index_conflicts.id())?;
+                replace_entries_with_their_restored_conflicts(&mut index, index_conflicts.id())?;
                 Some(index)
             }
             (None, None, Some(_index_conflicts)) => bail!(
@@ -123,7 +124,7 @@ pub(super) mod function {
     }
 
     #[expect(clippy::indexing_slicing)]
-    fn resolve_conflicts(
+    fn replace_entries_with_their_restored_conflicts(
         index: &mut gix::index::State,
         conflict_tree: gix::Id,
     ) -> anyhow::Result<()> {