feat(undo): per-container UndoScope for UndoManager

crates/loro-internal/Cargo.toml

@@ -128,3 +128,7 @@ harness = false
 [[bench]]
 name = "jsonpath"
 harness = false
+
+[[bench]]
+name = "undo"
+harness = false

crates/loro-internal/benches/undo.rs

@@ -0,0 +1,186 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+#[cfg(feature = "test_utils")]
+mod run {
+    use super::*;
+    use loro_internal::{
+        handler::{HandlerTrait, UpdateOptions},
+        LoroDoc, UndoManager, UndoScope,
+    };
+
+    /// Number of commits to record per benchmark iteration. Small enough to keep
+    /// runs fast, large enough to amortize per-iteration timer noise so we can
+    /// see sub-microsecond per-commit deltas between configurations.
+    const N_COMMITS: usize = 1_000;
+
+    fn one_text_edit(loro: &LoroDoc, value: &str) {
+        let text = loro.get_text("text");
+        text.update(value, UpdateOptions::default()).unwrap();
+        loro.commit_then_renew();
+    }
+
+    /// Record-time cost: build a doc, attach an UndoManager (in three different
+    /// configurations), and measure the time to record `N_COMMITS` local commits.
+    /// This is the hot path the subscription callback sits on.
+    pub fn record_local_commits(c: &mut Criterion) {
+        let mut g = c.benchmark_group("undo/record_local_commits");
+        g.sample_size(50);
+
+        // Baseline: no UndoManager attached at all. Establishes the cost of just
+        // committing N edits, so we can isolate the manager's overhead.
+        g.bench_function("no_manager", |b| {
+            b.iter(|| {
+                let loro = LoroDoc::default();
+                for i in 0..N_COMMITS {
+                    one_text_edit(&loro, &format!("v{}", i));
+                }
+            });
+        });
+
+        // Default: UndoManager with UndoScope::Doc (current behavior, our changes
+        // must not regress this). Compares directly against the same code on main.
+        g.bench_function("undo_manager_default_scope", |b| {
+            b.iter(|| {
+                let loro = LoroDoc::default();
+                let _undo = UndoManager::new(&loro);
+                for i in 0..N_COMMITS {
+                    one_text_edit(&loro, &format!("v{}", i));
+                }
+            });
+        });
+
+        // Scoped: UndoManager with UndoScope::Containers([text_id]). Quantifies
+        // the cost users opt into when they enable scope. Single-container scope
+        // is the smallest possible set; larger sets only affect FxHashSet lookup
+        // (constant-time average).
+        g.bench_function("undo_manager_scoped_one_container", |b| {
+            b.iter(|| {
+                let loro = LoroDoc::default();
+                let text = loro.get_text("text");
+                let _undo = UndoManager::new(&loro)
+                    .with_scope(UndoScope::containers([text.id()]));
+                for i in 0..N_COMMITS {
+                    one_text_edit(&loro, &format!("v{}", i));
+                }
+            });
+        });
+    }
+
+    /// Mixed-scope workload: alternate edits between an in-scope and an
+    /// out-of-scope container. Out-of-scope commits hit the
+    /// `compose_remote_event` branch instead of `record_checkpoint`. This
+    /// stresses the path the scope feature actually exists for.
+    pub fn record_mixed_scope(c: &mut Criterion) {
+        let mut g = c.benchmark_group("undo/record_mixed_scope");
+        g.sample_size(50);
+
+        // Doc-wide for reference: every commit is recorded.
+        g.bench_function("doc_scope_all_recorded", |b| {
+            b.iter(|| {
+                let loro = LoroDoc::default();
+                let text_a = loro.get_text("a");
+                let text_b = loro.get_text("b");
+                let _undo = UndoManager::new(&loro);
+                for i in 0..(N_COMMITS / 2) {
+                    text_a
+                        .update(&format!("a{}", i), UpdateOptions::default())
+                        .unwrap();
+                    loro.commit_then_renew();
+                    text_b
+                        .update(&format!("b{}", i), UpdateOptions::default())
+                        .unwrap();
+                    loro.commit_then_renew();
+                }
+            });
+        });
+
+        // Container scope = {a}: half the commits are filtered to the
+        // compose-as-remote branch.
+        g.bench_function("scoped_a_half_filtered", |b| {
+            b.iter(|| {
+                let loro = LoroDoc::default();
+                let text_a = loro.get_text("a");
+                let text_b = loro.get_text("b");
+                let _undo = UndoManager::new(&loro)
+                    .with_scope(UndoScope::containers([text_a.id()]));
+                for i in 0..(N_COMMITS / 2) {
+                    text_a
+                        .update(&format!("a{}", i), UpdateOptions::default())
+                        .unwrap();
+                    loro.commit_then_renew();
+                    text_b
+                        .update(&format!("b{}", i), UpdateOptions::default())
+                        .unwrap();
+                    loro.commit_then_renew();
+                }
+            });
+        });
+    }
+
+    /// Replay-time cost: build a doc with N recorded commits, then measure the
+    /// time to undo every commit followed by redoing every commit. Exercises
+    /// `undo_internal_with_scope` end-to-end including the optional mask block.
+    pub fn undo_redo_all(c: &mut Criterion) {
+        let mut g = c.benchmark_group("undo/replay_all");
+        g.sample_size(20);
+
+        g.bench_function("doc_scope", |b| {
+            b.iter_batched(
+                || {
+                    let loro = LoroDoc::default();
+                    let undo = UndoManager::new(&loro);
+                    for i in 0..N_COMMITS {
+                        one_text_edit(&loro, &format!("v{}", i));
+                    }
+                    (loro, undo)
+                },
+                |(_loro, undo)| {
+                    while undo.can_undo() {
+                        undo.undo().unwrap();
+                    }
+                    while undo.can_redo() {
+                        undo.redo().unwrap();
+                    }
+                },
+                criterion::BatchSize::SmallInput,
+            );
+        });
+
+        g.bench_function("scoped_one_container", |b| {
+            b.iter_batched(
+                || {
+                    let loro = LoroDoc::default();
+                    let text = loro.get_text("text");
+                    let undo = UndoManager::new(&loro)
+                        .with_scope(UndoScope::containers([text.id()]));
+                    for i in 0..N_COMMITS {
+                        one_text_edit(&loro, &format!("v{}", i));
+                    }
+                    (loro, undo)
+                },
+                |(_loro, undo)| {
+                    while undo.can_undo() {
+                        undo.undo().unwrap();
+                    }
+                    while undo.can_redo() {
+                        undo.redo().unwrap();
+                    }
+                },
+                criterion::BatchSize::SmallInput,
+            );
+        });
+    }
+}
+
+pub fn dumb(_c: &mut Criterion) {}
+
+#[cfg(feature = "test_utils")]
+criterion_group!(
+    benches,
+    run::record_local_commits,
+    run::record_mixed_scope,
+    run::undo_redo_all,
+);
+#[cfg(not(feature = "test_utils"))]
+criterion_group!(benches, dumb);
+criterion_main!(benches);

crates/loro-internal/src/lib.rs

@@ -40,7 +40,7 @@ pub use state::DocState;
 pub use state::{TreeNode, TreeNodeWithChildren, TreeParentId};
 use subscription::{LocalUpdateCallback, Observer, PeerIdUpdateCallback};
 use txn::Transaction;
-pub use undo::UndoManager;
+pub use undo::{UndoManager, UndoScope};
 pub use utils::subscription::SubscriberSetWithQueue;
 pub use utils::subscription::Subscription;
 pub mod allocation;

crates/loro-internal/src/loro.rs

@@ -1154,13 +1154,40 @@ impl LoroDoc {
     /// This implementation is kinda slow, but it's simple and maintainable. We can optimize it
     /// further when it's needed. The time complexity is O(n + m), n is the ops in the id_span, m is the
     /// distance from id_span to the current latest version.
-    #[instrument(level = "info", skip_all)]
+    #[inline]
     pub fn undo_internal(
         &self,
         id_span: IdSpan,
         container_remap: &mut FxHashMap<ContainerID, ContainerID>,
         post_transform_base: Option<&DiffBatch>,
         before_diff: &mut dyn FnMut(&DiffBatch),
+    ) -> LoroResult<CommitWhenDrop<'_>> {
+        // Tracing instrumentation lives on `undo_internal_with_scope`; keeping
+        // it off this thin wrapper avoids duplicate spans for direct callers.
+        self.undo_internal_with_scope(
+            id_span,
+            container_remap,
+            post_transform_base,
+            before_diff,
+            None,
+        )
+    }
+
+    /// Like [`Self::undo_internal`], but additionally accepts a `scope_filter`.
+    ///
+    /// When `scope_filter` is `Some`, the diff computed for the undo is masked
+    /// so only containers in the set are mutated. This lets callers (notably
+    /// [`UndoManager`] with [`crate::UndoScope::Containers`]) revert just the
+    /// in-scope portion of a commit even when the original commit also touched
+    /// out-of-scope containers.
+    #[instrument(level = "info", skip_all)]
+    pub fn undo_internal_with_scope(
+        &self,
+        id_span: IdSpan,
+        container_remap: &mut FxHashMap<ContainerID, ContainerID>,
+        post_transform_base: Option<&DiffBatch>,
+        before_diff: &mut dyn FnMut(&DiffBatch),
+        scope_filter: Option<&FxHashSet<ContainerID>>,
     ) -> LoroResult<CommitWhenDrop<'_>> {
         if !self.can_edit() {
             return Err(LoroError::EditWhenDetached);
@@ -1208,6 +1235,26 @@ impl LoroDoc {
         }
         drop(txn);
         self.start_auto_commit();
+
+        // If a scope filter was supplied (UndoManager with UndoScope::Containers),
+        // mask the diff so only in-scope containers are reverted. This is what
+        // makes mixed commits — single commits touching both in-scope and
+        // out-of-scope containers — undo only their in-scope portion.
+        // The filter walks container_remap so a remapped target counts as in-scope
+        // when the remap destination is in scope, mirroring _apply_diff's own walk.
+        let mut diff = diff;
+        if let Some(scope) = scope_filter {
+            let in_scope = |cid: &ContainerID| -> bool {
+                let mut id = cid.clone();
+                while let Some(rid) = container_remap.get(&id) {
+                    id = rid.clone();
+                }
+                scope.contains(&id)
+            };
+            diff.cid_to_events.retain(|cid, _| in_scope(cid));
+            diff.order.retain(|cid| in_scope(cid));
+        }
+
         // Try applying the diff, but ignore the error if it happens.
         // MovableList's undo behavior is too tricky to handle in a collaborative env
         // so in edge cases this may be an Error