Merge pull request #1664 from zed-industries/slow-undo

Use a `SumTree` as the backing storage of `UndoMap`

crates/text/src/text.rs

@@ -14,6 +14,7 @@ mod selection;
 pub mod subscription;
 #[cfg(test)]
 mod tests;
+mod undo_map;
 
 pub use anchor::*;
 use anyhow::Result;
@@ -46,6 +47,7 @@ use std::{
 pub use subscription::*;
 pub use sum_tree::Bias;
 use sum_tree::{FilterCursor, SumTree, TreeMap};
+use undo_map::UndoMap;
 
 lazy_static! {
     static ref CARRIAGE_RETURNS_REGEX: Regex = Regex::new("\r\n|\r").unwrap();
@@ -66,7 +68,7 @@ pub struct Buffer {
     version_barriers: Vec<(clock::Global, barrier::Sender)>,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone)]
 pub struct BufferSnapshot {
     replica_id: ReplicaId,
     remote_id: u64,
@@ -335,44 +337,6 @@ impl History {
     }
 }
 
-#[derive(Clone, Default, Debug)]
-struct UndoMap(HashMap<clock::Local, Vec<(clock::Local, u32)>>);
-
-impl UndoMap {
-    fn insert(&mut self, undo: &UndoOperation) {
-        for (edit_id, count) in &undo.counts {
-            self.0.entry(*edit_id).or_default().push((undo.id, *count));
-        }
-    }
-
-    fn is_undone(&self, edit_id: clock::Local) -> bool {
-        self.undo_count(edit_id) % 2 == 1
-    }
-
-    fn was_undone(&self, edit_id: clock::Local, version: &clock::Global) -> bool {
-        let undo_count = self
-            .0
-            .get(&edit_id)
-            .unwrap_or(&Vec::new())
-            .iter()
-            .filter(|(undo_id, _)| version.observed(*undo_id))
-            .map(|(_, undo_count)| *undo_count)
-            .max()
-            .unwrap_or(0);
-        undo_count % 2 == 1
-    }
-
-    fn undo_count(&self, edit_id: clock::Local) -> u32 {
-        self.0
-            .get(&edit_id)
-            .unwrap_or(&Vec::new())
-            .iter()
-            .map(|(_, undo_count)| *undo_count)
-            .max()
-            .unwrap_or(0)
-    }
-}
-
 struct Edits<'a, D: TextDimension, F: FnMut(&FragmentSummary) -> bool> {
     visible_cursor: rope::Cursor<'a>,
     deleted_cursor: rope::Cursor<'a>,
@@ -1218,13 +1182,6 @@ impl Buffer {
         &self.history.operations
     }
 
-    pub fn undo_history(&self) -> impl Iterator<Item = (&clock::Local, &[(clock::Local, u32)])> {
-        self.undo_map
-            .0
-            .iter()
-            .map(|(edit_id, undo_counts)| (edit_id, undo_counts.as_slice()))
-    }
-
     pub fn undo(&mut self) -> Option<(TransactionId, Operation)> {
         if let Some(entry) = self.history.pop_undo() {
             let transaction = entry.transaction.clone();

crates/text/src/undo_map.rs

@@ -0,0 +1,112 @@
+use crate::UndoOperation;
+use std::cmp;
+use sum_tree::{Bias, SumTree};
+
+#[derive(Copy, Clone, Debug)]
+struct UndoMapEntry {
+    key: UndoMapKey,
+    undo_count: u32,
+}
+
+impl sum_tree::Item for UndoMapEntry {
+    type Summary = UndoMapKey;
+
+    fn summary(&self) -> Self::Summary {
+        self.key
+    }
+}
+
+impl sum_tree::KeyedItem for UndoMapEntry {
+    type Key = UndoMapKey;
+
+    fn key(&self) -> Self::Key {
+        self.key
+    }
+}
+
+#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
+struct UndoMapKey {
+    edit_id: clock::Local,
+    undo_id: clock::Local,
+}
+
+impl sum_tree::Summary for UndoMapKey {
+    type Context = ();
+
+    fn add_summary(&mut self, summary: &Self, _: &Self::Context) {
+        *self = cmp::max(*self, *summary);
+    }
+}
+
+#[derive(Clone, Default)]
+pub struct UndoMap(SumTree<UndoMapEntry>);
+
+impl UndoMap {
+    pub fn insert(&mut self, undo: &UndoOperation) {
+        let edits = undo
+            .counts
+            .iter()
+            .map(|(edit_id, count)| {
+                sum_tree::Edit::Insert(UndoMapEntry {
+                    key: UndoMapKey {
+                        edit_id: *edit_id,
+                        undo_id: undo.id,
+                    },
+                    undo_count: *count,
+                })
+            })
+            .collect::<Vec<_>>();
+        self.0.edit(edits, &());
+    }
+
+    pub fn is_undone(&self, edit_id: clock::Local) -> bool {
+        self.undo_count(edit_id) % 2 == 1
+    }
+
+    pub fn was_undone(&self, edit_id: clock::Local, version: &clock::Global) -> bool {
+        let mut cursor = self.0.cursor::<UndoMapKey>();
+        cursor.seek(
+            &UndoMapKey {
+                edit_id,
+                undo_id: Default::default(),
+            },
+            Bias::Left,
+            &(),
+        );
+
+        let mut undo_count = 0;
+        for entry in cursor {
+            if entry.key.edit_id != edit_id {
+                break;
+            }
+
+            if version.observed(entry.key.undo_id) {
+                undo_count = cmp::max(undo_count, entry.undo_count);
+            }
+        }
+
+        undo_count % 2 == 1
+    }
+
+    pub fn undo_count(&self, edit_id: clock::Local) -> u32 {
+        let mut cursor = self.0.cursor::<UndoMapKey>();
+        cursor.seek(
+            &UndoMapKey {
+                edit_id,
+                undo_id: Default::default(),
+            },
+            Bias::Left,
+            &(),
+        );
+
+        let mut undo_count = 0;
+        for entry in cursor {
+            if entry.key.edit_id != edit_id {
+                break;
+            }
+
+            undo_count = cmp::max(undo_count, entry.undo_count);
+        }
+        undo_count
+    }
+}