squash: move the core functionality to jj_lib

This allows 'squash' to be executed more easily from a programmatic context

cli/src/commands/squash.rs

@@ -16,9 +16,9 @@ use itertools::Itertools as _;
 use jj_lib::commit::Commit;
 use jj_lib::commit::CommitIteratorExt;
 use jj_lib::matchers::Matcher;
-use jj_lib::merged_tree::MergedTree;
 use jj_lib::object_id::ObjectId;
 use jj_lib::repo::Repo;
+use jj_lib::rewrite;
 use jj_lib::settings::UserSettings;
 use tracing::instrument;
 
@@ -196,12 +196,6 @@ pub fn move_diff(
     tx.base_workspace_helper()
         .check_rewritable(sources.iter().chain(std::iter::once(destination)).ids())?;
 
-    struct SourceCommit<'a> {
-        commit: &'a Commit,
-        parent_tree: MergedTree,
-        selected_tree: MergedTree,
-        abandon: bool,
-    }
     let mut source_commits = vec![];
     for source in sources {
         let parent_tree = source.parent_tree(tx.repo())?;
@@ -227,105 +221,52 @@ from the source will be moved into the destination.
         let selected_tree_id =
             diff_selector.select(&parent_tree, &source_tree, matcher, format_instructions)?;
         let selected_tree = tx.repo().store().get_root_tree(&selected_tree_id)?;
-        let abandon = !keep_emptied && selected_tree.id() == source_tree.id();
+
-        if !abandon && selected_tree_id == parent_tree.id() {
+        source_commits.push(rewrite::CommitToSquash {
-            // Nothing selected from this commit. If it's abandoned (i.e. already empty), we
+            commit: source.clone(),
-            // still include it so `jj squash` can be used for abandoning an empty commit in
-            // the middle of a stack.
-            continue;
-        }
-        // TODO: Do we want to optimize the case of moving to the parent commit (`jj
-        // squash -r`)? The source tree will be unchanged in that case.
-        source_commits.push(SourceCommit {
-            commit: source,
-            parent_tree,
             selected_tree,
-            abandon,
+            parent_tree,
         });
     }
-    if source_commits.is_empty() {
-        if diff_selector.is_interactive() {
-            return Err(user_error("No changes selected"));
-        }
 
-        if let [only_path] = path_arg {
+    let repo_path = tx.base_workspace_helper().repo_path().to_owned();
-            if no_rev_arg
+    match rewrite::squash_commits(
-                && tx
+        settings,
-                    .base_workspace_helper()
+        tx.repo_mut(),
-                    .parse_revset(ui, &RevisionArg::from(only_path.to_owned()))
+        &source_commits,
-                    .is_ok()
+        destination,
-            {
+        keep_emptied,
-                writeln!(
+        |abandoned_commits| match description {
-                    ui.warning_default(),
+            SquashedDescription::Exact(description) => Ok(description),
-                    "The argument {only_path:?} is being interpreted as a path. To specify a \
+            SquashedDescription::UseDestination => Ok(destination.description().to_owned()),
-                     revset, pass -r {only_path:?} instead."
+            SquashedDescription::Combine => {
-                )?;
+                let abandoned_commits = abandoned_commits.iter().map(|c| &c.commit).collect_vec();
+                combine_messages(&repo_path, &abandoned_commits, destination, settings)
+            }
+        },
+    )? {
+        rewrite::SquashResult::NoChanges => {
+            if diff_selector.is_interactive() {
+                return Err(user_error("No changes selected"));
             }
-        }
 
-        return Ok(());
+            if let [only_path] = path_arg {
-    }
+                if no_rev_arg
+                    && tx
+                        .base_workspace_helper()
+                        .parse_revset(ui, &RevisionArg::from(only_path.to_owned()))
+                        .is_ok()
+                {
+                    writeln!(
+                        ui.warning_default(),
+                        "The argument {only_path:?} is being interpreted as a path. To specify a \
+                         revset, pass -r {only_path:?} instead."
+                    )?;
+                }
+            }
 
-    for source in &source_commits {
+            Ok(())
-        if source.abandon {
-            tx.repo_mut()
-                .record_abandoned_commit(source.commit.id().clone());
-        } else {
-            let source_tree = source.commit.tree()?;
-            // Apply the reverse of the selected changes onto the source
-            let new_source_tree = source_tree.merge(&source.selected_tree, &source.parent_tree)?;
-            tx.repo_mut()
-                .rewrite_commit(settings, source.commit)
-                .set_tree_id(new_source_tree.id().clone())
-                .write()?;
         }
+        rewrite::SquashResult::NewCommit(_) => Ok(()),
     }
-
-    let mut rewritten_destination = destination.clone();
-    if sources
-        .iter()
-        .any(|source| tx.repo().index().is_ancestor(source.id(), destination.id()))
-    {
-        // If we're moving changes to a descendant, first rebase descendants onto the
-        // rewritten sources. Otherwise it will likely already have the content
-        // changes we're moving, so applying them will have no effect and the
-        // changes will disappear.
-        let rebase_map = tx.repo_mut().rebase_descendants_return_map(settings)?;
-        let rebased_destination_id = rebase_map.get(destination.id()).unwrap().clone();
-        rewritten_destination = tx.repo().store().get_commit(&rebased_destination_id)?;
-    }
-    // Apply the selected changes onto the destination
-    let mut destination_tree = rewritten_destination.tree()?;
-    for source in &source_commits {
-        destination_tree = destination_tree.merge(&source.parent_tree, &source.selected_tree)?;
-    }
-    let description = match description {
-        SquashedDescription::Exact(description) => description,
-        SquashedDescription::UseDestination => destination.description().to_owned(),
-        SquashedDescription::Combine => {
-            let abandoned_commits = source_commits
-                .iter()
-                .filter_map(|source| source.abandon.then_some(source.commit))
-                .collect_vec();
-            combine_messages(
-                tx.base_workspace_helper().repo_path(),
-                &abandoned_commits,
-                destination,
-                settings,
-            )?
-        }
-    };
-    let mut predecessors = vec![destination.id().clone()];
-    predecessors.extend(
-        source_commits
-            .iter()
-            .map(|source| source.commit.id().clone()),
-    );
-    tx.repo_mut()
-        .rewrite_commit(settings, &rewritten_destination)
-        .set_tree_id(destination_tree.id().clone())
-        .set_predecessors(predecessors)
-        .set_description(description)
-        .write()?;
-    Ok(())
 }

lib/src/rewrite.rs

@@ -830,3 +830,125 @@ pub fn move_commits(
         num_skipped_rebases,
     })
 }
+
+pub struct CommitToSquash {
+    pub commit: Commit,
+    pub selected_tree: MergedTree,
+    pub parent_tree: MergedTree,
+}
+
+impl CommitToSquash {
+    /// Returns true if the selection contains all changes in the commit.
+    fn is_full_selection(&self) -> bool {
+        &self.selected_tree.id() == self.commit.tree_id()
+    }
+
+    /// Returns true if the selection matches the parent tree (contains no
+    /// changes from the commit).
+    ///
+    /// Both `is_full_selection()` and `is_empty_selection()`
+    /// can be true if the commit is itself empty.
+    fn is_empty_selection(&self) -> bool {
+        self.selected_tree.id() == self.parent_tree.id()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum SquashResult {
+    /// No inputs contained actual changes.
+    NoChanges,
+    /// Destination was rewritten.
+    NewCommit(Commit),
+}
+
+/// Squash `sources` into `destination` and return a CommitBuilder for the
+/// resulting commit. Caller is responsible for setting the description and
+/// finishing the commit.
+pub fn squash_commits<E>(
+    settings: &UserSettings,
+    repo: &mut MutableRepo,
+    sources: &[CommitToSquash],
+    destination: &Commit,
+    keep_emptied: bool,
+    description_fn: impl FnOnce(&[&CommitToSquash]) -> Result<String, E>,
+) -> Result<SquashResult, E>
+where
+    E: From<BackendError>,
+{
+    struct SourceCommit<'a> {
+        commit: &'a CommitToSquash,
+        abandon: bool,
+    }
+    let mut source_commits = vec![];
+    for source in sources {
+        let abandon = !keep_emptied && source.is_full_selection();
+        if !abandon && source.is_empty_selection() {
+            // Nothing selected from this commit. If it's abandoned (i.e. already empty), we
+            // still include it so `jj squash` can be used for abandoning an empty commit in
+            // the middle of a stack.
+            continue;
+        }
+
+        // TODO: Do we want to optimize the case of moving to the parent commit (`jj
+        // squash -r`)? The source tree will be unchanged in that case.
+        source_commits.push(SourceCommit {
+            commit: source,
+            abandon,
+        });
+    }
+
+    if source_commits.is_empty() {
+        return Ok(SquashResult::NoChanges);
+    }
+
+    let mut abandoned_commits = vec![];
+    for source in &source_commits {
+        if source.abandon {
+            repo.record_abandoned_commit(source.commit.commit.id().clone());
+            abandoned_commits.push(source.commit);
+        } else {
+            let source_tree = source.commit.commit.tree()?;
+            // Apply the reverse of the selected changes onto the source
+            let new_source_tree =
+                source_tree.merge(&source.commit.selected_tree, &source.commit.parent_tree)?;
+            repo.rewrite_commit(settings, &source.commit.commit)
+                .set_tree_id(new_source_tree.id().clone())
+                .write()?;
+        }
+    }
+
+    let mut rewritten_destination = destination.clone();
+    if sources.iter().any(|source| {
+        repo.index()
+            .is_ancestor(source.commit.id(), destination.id())
+    }) {
+        // If we're moving changes to a descendant, first rebase descendants onto the
+        // rewritten sources. Otherwise it will likely already have the content
+        // changes we're moving, so applying them will have no effect and the
+        // changes will disappear.
+        let rebase_map = repo.rebase_descendants_return_map(settings)?;
+        let rebased_destination_id = rebase_map.get(destination.id()).unwrap().clone();
+        rewritten_destination = repo.store().get_commit(&rebased_destination_id)?;
+    }
+    // Apply the selected changes onto the destination
+    let mut destination_tree = rewritten_destination.tree()?;
+    for source in &source_commits {
+        destination_tree =
+            destination_tree.merge(&source.commit.parent_tree, &source.commit.selected_tree)?;
+    }
+    let mut predecessors = vec![destination.id().clone()];
+    predecessors.extend(
+        source_commits
+            .iter()
+            .map(|source| source.commit.commit.id().clone()),
+    );
+
+    let destination = repo
+        .rewrite_commit(settings, &rewritten_destination)
+        .set_tree_id(destination_tree.id().clone())
+        .set_predecessors(predecessors)
+        .set_description(description_fn(&abandoned_commits)?)
+        .write()?;
+
+    Ok(SquashResult::NewCommit(destination))
+}