op_walk: add function that reparents (and abandons) operation range

This will be used in "jj op abandon ..op_id" command. The "op_id..@" range will
be reparented onto the root operation.

The current implementation is good enough for local repos, but it won't scale.
We might want to extract it as a trait method or introduce OpIndex for
efficient DAG operation.

lib/src/op_walk.rs

@@ -15,6 +15,8 @@
 //! Utility for operation id resolution and traversal.
 
 use std::cmp::Ordering;
+use std::collections::{HashMap, HashSet};
+use std::slice;
 use std::sync::Arc;
 
 use itertools::Itertools as _;
@@ -230,3 +232,70 @@ pub fn walk_ancestors(head_ops: &[Operation]) -> impl Iterator<Item = OpStoreRes
     )
     .map_ok(|OperationByEndTime(op)| op)
 }
+
+/// Stats about `reparent_range()`.
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct ReparentStats {
+    /// New head operation ids.
+    pub new_head_ids: Vec<OperationId>,
+    /// The number of rewritten operations.
+    pub rewritten_count: usize,
+    /// The number of ancestor operations that become unreachable from the
+    /// rewritten heads.
+    pub unreachable_count: usize,
+}
+
+/// Reparents the operation range `root_ops..head_ops` onto the `dest_op`.
+///
+/// Returns the new head operation ids as well as some stats. If the old
+/// operation heads are remapped to the new heads, the operations within the
+/// range `dest_op..root_ops` become unreachable.
+///
+/// If the source operation range `root_ops..head_ops` was empty, the
+/// `new_head_ids` will be `[dest_op.id()]`, meaning the `dest_op` is the head.
+// TODO: Find better place to host this function. It might be an OpStore method.
+pub fn reparent_range(
+    op_store: &dyn OpStore,
+    root_ops: &[Operation],
+    head_ops: &[Operation],
+    dest_op: &Operation,
+) -> OpStoreResult<ReparentStats> {
+    // Calculate ::root_ops to exclude them from the source range and count the
+    // number of operations that become unreachable.
+    let mut unwanted_ids: HashSet<_> = walk_ancestors(root_ops)
+        .map_ok(|op| op.id().clone())
+        .try_collect()?;
+    let ops_to_reparent: Vec<_> = walk_ancestors(head_ops)
+        .filter_ok(|op| !unwanted_ids.contains(op.id()))
+        .try_collect()?;
+    for op in walk_ancestors(slice::from_ref(dest_op)) {
+        unwanted_ids.remove(op?.id());
+    }
+    let unreachable_ids = unwanted_ids;
+
+    let mut rewritten_ids = HashMap::new();
+    for old_op in ops_to_reparent.into_iter().rev() {
+        let mut data = old_op.store_operation().clone();
+        let mut dest_once = Some(dest_op.id());
+        data.parents = data
+            .parents
+            .iter()
+            .filter_map(|id| rewritten_ids.get(id).or_else(|| dest_once.take()))
+            .cloned()
+            .collect();
+        let new_id = op_store.write_operation(&data)?;
+        rewritten_ids.insert(old_op.id().clone(), new_id);
+    }
+
+    let mut dest_once = Some(dest_op.id());
+    let new_head_ids = head_ops
+        .iter()
+        .filter_map(|op| rewritten_ids.get(op.id()).or_else(|| dest_once.take()))
+        .cloned()
+        .collect();
+    Ok(ReparentStats {
+        new_head_ids,
+        rewritten_count: rewritten_ids.len(),
+        unreachable_count: unreachable_ids.len(),
+    })
+}

lib/tests/test_operations.rs

@@ -13,12 +13,15 @@
 // limitations under the License.
 
 use std::path::Path;
+use std::slice;
+use std::sync::Arc;
 
 use assert_matches::assert_matches;
 use itertools::Itertools as _;
 use jj_lib::backend::{CommitId, ObjectId};
 use jj_lib::op_walk::{self, OpsetEvaluationError, OpsetResolutionError};
-use jj_lib::repo::Repo;
+use jj_lib::operation::Operation;
+use jj_lib::repo::{ReadonlyRepo, Repo};
 use jj_lib::settings::UserSettings;
 use testutils::{create_random_commit, write_random_commit, TestRepo};
 
@@ -182,6 +185,253 @@ fn test_isolation() {
     assert_heads(repo.as_ref(), vec![rewrite1.id(), rewrite2.id()]);
 }
 
+#[test]
+fn test_reparent_range_linear() {
+    let settings = testutils::user_settings();
+    let test_repo = TestRepo::init();
+    let repo_0 = test_repo.repo;
+    let op_store = repo_0.op_store();
+
+    let read_op = |id| {
+        let data = op_store.read_operation(id).unwrap();
+        Operation::new(op_store.clone(), id.clone(), data)
+    };
+
+    fn op_parents<const N: usize>(op: &Operation) -> [Operation; N] {
+        let parents: Vec<_> = op.parents().try_collect().unwrap();
+        parents.try_into().unwrap()
+    }
+
+    // Set up linear operation graph:
+    // D
+    // C
+    // B
+    // A
+    // 0 (initial)
+    let random_tx = |repo: &Arc<ReadonlyRepo>| {
+        let mut tx = repo.start_transaction(&settings);
+        write_random_commit(tx.mut_repo(), &settings);
+        tx
+    };
+    let repo_a = random_tx(&repo_0).commit("op A");
+    let repo_b = random_tx(&repo_a).commit("op B");
+    let repo_c = random_tx(&repo_b).commit("op C");
+    let repo_d = random_tx(&repo_c).commit("op D");
+
+    // Reparent B..D (=C|D) onto A:
+    // D'
+    // C'
+    // A
+    // 0 (initial)
+    let stats = op_walk::reparent_range(
+        op_store.as_ref(),
+        slice::from_ref(repo_b.operation()),
+        slice::from_ref(repo_d.operation()),
+        repo_a.operation(),
+    )
+    .unwrap();
+    assert_eq!(stats.new_head_ids.len(), 1);
+    assert_eq!(stats.rewritten_count, 2);
+    assert_eq!(stats.unreachable_count, 1);
+    let new_op_d = read_op(&stats.new_head_ids[0]);
+    assert_eq!(
+        new_op_d.store_operation().metadata,
+        repo_d.operation().store_operation().metadata
+    );
+    assert_eq!(
+        new_op_d.store_operation().view_id,
+        repo_d.operation().store_operation().view_id
+    );
+    let [new_op_c] = op_parents(&new_op_d);
+    assert_eq!(
+        new_op_c.store_operation().metadata,
+        repo_c.operation().store_operation().metadata
+    );
+    assert_eq!(
+        new_op_c.store_operation().view_id,
+        repo_c.operation().store_operation().view_id
+    );
+    assert_eq!(new_op_c.parent_ids(), slice::from_ref(repo_a.op_id()));
+
+    // Reparent empty range onto A
+    let stats = op_walk::reparent_range(
+        op_store.as_ref(),
+        slice::from_ref(repo_d.operation()),
+        slice::from_ref(repo_d.operation()),
+        repo_a.operation(),
+    )
+    .unwrap();
+    assert_eq!(stats.new_head_ids, vec![repo_a.op_id().clone()]);
+    assert_eq!(stats.rewritten_count, 0);
+    assert_eq!(stats.unreachable_count, 3);
+}
+
+#[test]
+fn test_reparent_range_branchy() {
+    let settings = testutils::user_settings();
+    let test_repo = TestRepo::init();
+    let repo_0 = test_repo.repo;
+    let op_store = repo_0.op_store();
+
+    let read_op = |id| {
+        let data = op_store.read_operation(id).unwrap();
+        Operation::new(op_store.clone(), id.clone(), data)
+    };
+
+    fn op_parents<const N: usize>(op: &Operation) -> [Operation; N] {
+        let parents: Vec<_> = op.parents().try_collect().unwrap();
+        parents.try_into().unwrap()
+    }
+
+    // Set up branchy operation graph:
+    // G
+    // |\
+    // | F
+    // E |
+    // D |
+    // |/
+    // C
+    // B
+    // A
+    // 0 (initial)
+    let random_tx = |repo: &Arc<ReadonlyRepo>| {
+        let mut tx = repo.start_transaction(&settings);
+        write_random_commit(tx.mut_repo(), &settings);
+        tx
+    };
+    let repo_a = random_tx(&repo_0).commit("op A");
+    let repo_b = random_tx(&repo_a).commit("op B");
+    let repo_c = random_tx(&repo_b).commit("op C");
+    let repo_d = random_tx(&repo_c).commit("op D");
+    let tx_e = random_tx(&repo_d);
+    let tx_f = random_tx(&repo_c);
+    let repo_g = testutils::commit_transactions(&settings, vec![tx_e, tx_f]);
+    let [op_e, op_f] = op_parents(repo_g.operation());
+
+    // Reparent D..G (= E|F|G) onto B:
+    // G'
+    // |\
+    // | F'
+    // E'|
+    // |/
+    // B
+    // A
+    // 0 (initial)
+    let stats = op_walk::reparent_range(
+        op_store.as_ref(),
+        slice::from_ref(repo_d.operation()),
+        slice::from_ref(repo_g.operation()),
+        repo_b.operation(),
+    )
+    .unwrap();
+    assert_eq!(stats.new_head_ids.len(), 1);
+    assert_eq!(stats.rewritten_count, 3);
+    assert_eq!(stats.unreachable_count, 2);
+    let new_op_g = read_op(&stats.new_head_ids[0]);
+    assert_eq!(
+        new_op_g.store_operation().metadata,
+        repo_g.operation().store_operation().metadata
+    );
+    assert_eq!(
+        new_op_g.store_operation().view_id,
+        repo_g.operation().store_operation().view_id
+    );
+    let [new_op_e, new_op_f] = op_parents(&new_op_g);
+    assert_eq!(new_op_e.parent_ids(), slice::from_ref(repo_b.op_id()));
+    assert_eq!(new_op_f.parent_ids(), slice::from_ref(repo_b.op_id()));
+
+    // Reparent B..G (=C|D|E|F|G) onto A:
+    // G'
+    // |\
+    // | F'
+    // E'|
+    // D'|
+    // |/
+    // C'
+    // A
+    // 0 (initial)
+    let stats = op_walk::reparent_range(
+        op_store.as_ref(),
+        slice::from_ref(repo_b.operation()),
+        slice::from_ref(repo_g.operation()),
+        repo_a.operation(),
+    )
+    .unwrap();
+    assert_eq!(stats.new_head_ids.len(), 1);
+    assert_eq!(stats.rewritten_count, 5);
+    assert_eq!(stats.unreachable_count, 1);
+    let new_op_g = read_op(&stats.new_head_ids[0]);
+    assert_eq!(
+        new_op_g.store_operation().metadata,
+        repo_g.operation().store_operation().metadata
+    );
+    assert_eq!(
+        new_op_g.store_operation().view_id,
+        repo_g.operation().store_operation().view_id
+    );
+    let [new_op_e, new_op_f] = op_parents(&new_op_g);
+    let [new_op_d] = op_parents(&new_op_e);
+    assert_eq!(new_op_d.parent_ids(), new_op_f.parent_ids());
+    let [new_op_c] = op_parents(&new_op_d);
+    assert_eq!(new_op_c.parent_ids(), slice::from_ref(repo_a.op_id()));
+
+    // Reparent (E|F)..G (=G) onto D:
+    // G'
+    // D
+    // C
+    // B
+    // A
+    // 0 (initial)
+    let stats = op_walk::reparent_range(
+        op_store.as_ref(),
+        &[op_e.clone(), op_f.clone()],
+        slice::from_ref(repo_g.operation()),
+        repo_d.operation(),
+    )
+    .unwrap();
+    assert_eq!(stats.new_head_ids.len(), 1);
+    assert_eq!(stats.rewritten_count, 1);
+    assert_eq!(stats.unreachable_count, 2);
+    let new_op_g = read_op(&stats.new_head_ids[0]);
+    assert_eq!(
+        new_op_g.store_operation().metadata,
+        repo_g.operation().store_operation().metadata
+    );
+    assert_eq!(
+        new_op_g.store_operation().view_id,
+        repo_g.operation().store_operation().view_id
+    );
+    assert_eq!(new_op_g.parent_ids(), slice::from_ref(repo_d.op_id()));
+
+    // Reparent C..F (=F) onto D (ignoring G):
+    // F'
+    // D
+    // C
+    // B
+    // A
+    // 0 (initial)
+    let stats = op_walk::reparent_range(
+        op_store.as_ref(),
+        slice::from_ref(repo_c.operation()),
+        slice::from_ref(&op_f),
+        repo_d.operation(),
+    )
+    .unwrap();
+    assert_eq!(stats.new_head_ids.len(), 1);
+    assert_eq!(stats.rewritten_count, 1);
+    assert_eq!(stats.unreachable_count, 0);
+    let new_op_f = read_op(&stats.new_head_ids[0]);
+    assert_eq!(
+        new_op_f.store_operation().metadata,
+        op_f.store_operation().metadata
+    );
+    assert_eq!(
+        new_op_f.store_operation().view_id,
+        op_f.store_operation().view_id
+    );
+    assert_eq!(new_op_f.parent_ids(), slice::from_ref(repo_d.op_id()));
+}
+
 fn stable_op_id_settings() -> UserSettings {
     UserSettings::from_config(
         testutils::base_config()