fix(collab): place keyed-list puts through an index within a batch (#1893)

* fix(collab): place keyed-list puts through an index within a batch

Each Put of a comment, style, variable or collection scanned the whole
list for its key (cloning String keys on the way) before appending, so a
batch of N new entries cost N squared comparisons on the hub every
session shares. Removing an image also rebuilt the document's image-use
set before checking the image existed.

Replica::apply now keeps a per-batch key index for each keyed list; it
is rebuilt when the list's length or the entry at a key no longer
matches, and dropped on whole-list sets, reorders and real removals.
Removing an absent entry or image returns without scanning.

* fix(collab): keep the keyed-list index current through removals

Build the index without dead code outside tests, drop the wall-clock
assertion from the large-batch test, and shift the index down after a
single-entry removal instead of discarding it, so a run of removals in
one batch does not index the list again each time. Staleness checks for
puts and removals share one helper. Tests cover removals, duplicate keys
from a whole-list set, and absent keys.

crates/collab/src/docprops.rs

@@ -54,18 +54,71 @@ struct List<T: Keyed> {
     entry: fn(T::Key) -> Entry,
     get: fn(&Document) -> &Vec<T>,
     get_mut: fn(&mut Document) -> &mut Vec<T>,
+    /// The list's slot in an [`Index`].
+    index: fn(&mut Index) -> &mut Option<KeyIndex<T::Key>>,
 }
 
-const STYLES: List<Style> =
-    List { whole: DocProp::Styles, order: DocProp::StyleOrder, entry: Entry::Style, get: |d| &d.styles, get_mut: |d| &mut d.styles };
-const COMMENTS: List<Comment> =
-    List { whole: DocProp::Comments, order: DocProp::CommentOrder, entry: Entry::Comment, get: |d| &d.comments, get_mut: |d| &mut d.comments };
+/// Where the entries of each keyed list sit, kept while one batch is applied: a batch can put
+/// hundreds of thousands of new entries, and finding each one's place by a scan of the growing
+/// list would cost the square of that on the thread every session shares.
+#[derive(Default)]
+pub(crate) struct Index {
+    styles: Option<KeyIndex<String>>,
+    comments: Option<KeyIndex<u64>>,
+    collections: Option<KeyIndex<String>>,
+    variables: Option<KeyIndex<String>>,
+}
+
+/// The position of the first entry with each key in one list, and the list's length when it was
+/// last brought up to date: a list whose length changed since is indexed again.
+struct KeyIndex<K> {
+    at: BTreeMap<K, usize>,
+    len: usize,
+}
+
+impl<K: Ord> KeyIndex<K> {
+    fn of<T: Keyed<Key = K>>(list: &[T]) -> Self {
+        let mut at = BTreeMap::new();
+        for (i, e) in list.iter().enumerate() {
+            at.entry(e.key()).or_insert(i);
+        }
+        KeyIndex { at, len: list.len() }
+    }
+
+    /// The index in `slot`, built again from `list` unless its length and its place for `key`
+    /// still match the list, so a list changed some other way is never trusted.
+    fn current<'a, T: Keyed<Key = K>>(slot: &'a mut Option<Self>, list: &[T], key: &K) -> &'a mut Self {
+        let stale = |ix: &Self| ix.len != list.len() || ix.at.get(key).is_some_and(|&i| list.get(i).is_none_or(|e| e.key() != *key));
+        if slot.as_ref().is_none_or(stale) {
+            *slot = Some(KeyIndex::of(list));
+        }
+        slot.get_or_insert_with(|| KeyIndex::of(list))
+    }
+}
+
+const STYLES: List<Style> = List {
+    whole: DocProp::Styles,
+    order: DocProp::StyleOrder,
+    entry: Entry::Style,
+    get: |d| &d.styles,
+    get_mut: |d| &mut d.styles,
+    index: |i| &mut i.styles,
+};
+const COMMENTS: List<Comment> = List {
+    whole: DocProp::Comments,
+    order: DocProp::CommentOrder,
+    entry: Entry::Comment,
+    get: |d| &d.comments,
+    get_mut: |d| &mut d.comments,
+    index: |i| &mut i.comments,
+};
 const COLLECTIONS: List<Collection> = List {
     whole: DocProp::Collections,
     order: DocProp::CollectionOrder,
     entry: Entry::Collection,
     get: |d| &d.variables.collections,
     get_mut: |d| &mut d.variables.collections,
+    index: |i| &mut i.collections,
 };
 const VARIABLES: List<Variable> = List {
     whole: DocProp::Variables,
@@ -73,6 +126,7 @@ const VARIABLES: List<Variable> = List {
     entry: Entry::Variable,
     get: |d| &d.variables.variables,
     get_mut: |d| &mut d.variables.variables,
+    index: |i| &mut i.variables,
 };
 
 impl<T: Keyed> List<T> {
@@ -111,25 +165,53 @@ impl<T: Keyed> List<T> {
         Ok(())
     }
 
-    fn put(&self, doc: &mut Document, key: T::Key, value: &Json) -> Result<(), Error> {
+    fn put(&self, doc: &mut Document, key: T::Key, value: &Json, index: &mut Index) -> Result<(), Error> {
         let t: T = value.read("entry")?;
         if t.key() != key {
             return Err(Error::Invalid("an entry's id must match the entry it is put at".into()));
         }
         let list = (self.get_mut)(doc);
-        match list.iter_mut().find(|e| e.key() == key) {
+        let ix = KeyIndex::current((self.index)(index), list, &key);
+        match ix.at.get(&key).copied().and_then(|i| list.get_mut(i)) {
             Some(e) => *e = t,
-            None => list.push(t),
+            None => {
+                list.push(t);
+                ix.at.insert(key, list.len() - 1);
+                ix.len = list.len();
+            }
         }
         Ok(())
     }
 
-    fn remove(&self, doc: &mut Document, key: &T::Key) {
-        (self.get_mut)(doc).retain(|e| e.key() != *key);
+    fn remove(&self, doc: &mut Document, key: &T::Key, index: &mut Index) {
+        let list = (self.get_mut)(doc);
+        let slot = (self.index)(index);
+        let ix = KeyIndex::current(slot, list, key);
+        // An absent entry leaves the list as it is, without a scan of it.
+        let Some(at) = ix.at.get(key).copied() else {
+            return;
+        };
+        let before = list.len();
+        list.retain(|e| e.key() != *key);
+        // One entry gone at `at`: the index is kept current by shifting the entries after it down,
+        // so the next removal in the batch need not index the list again. Duplicates of the key
+        // (more than one entry gone) leave it to be rebuilt.
+        if list.len() + 1 == before {
+            ix.at.remove(key);
+            for i in ix.at.values_mut() {
+                if *i > at {
+                    *i -= 1;
+                }
+            }
+            ix.len = list.len();
+        } else {
+            *slot = None;
+        }
     }
 
     /// Entries in `keys` order first, the rest after in their current order.
-    fn reorder(&self, doc: &mut Document, value: &Json) -> Result<(), Error> {
+    fn reorder(&self, doc: &mut Document, value: &Json, index: &mut Index) -> Result<(), Error> {
+        *(self.index)(index) = None;
         let keys: Vec<T::Key> = value.read("order")?;
         let list = (self.get_mut)(doc);
         let mut slots: Vec<Option<T>> = std::mem::take(list).into_iter().map(Some).collect();
@@ -150,7 +232,8 @@ impl<T: Keyed> List<T> {
         Ok(())
     }
 
-    fn set_whole(&self, doc: &mut Document, value: &Json) -> Result<(), Error> {
+    fn set_whole(&self, doc: &mut Document, value: &Json, index: &mut Index) -> Result<(), Error> {
+        *(self.index)(index) = None;
         *(self.get_mut)(doc) = value.read("list")?;
         Ok(())
     }
@@ -313,31 +396,38 @@ pub(crate) fn diff(a: &Document, b: &Document, out: &mut Vec<Change>) -> Result<
     diff_map(&a.votes, &b.votes, |x, y| x == y, |k| Entry::Ballot(k.clone()), Json::of, out)
 }
 
-/// Apply a document-level change. Returns whether it can change which pages are listed (the
-/// library store page is never listed).
+/// [`apply_indexed`] for one change on its own.
+#[cfg(test)]
 pub(crate) fn apply(doc: &mut Document, change: &Change) -> Result<bool, Error> {
+    apply_indexed(doc, change, &mut Index::default())
+}
+
+/// Apply a document-level change, with the keyed lists' positions kept in `index` between the
+/// changes of one batch. Returns whether it can change which pages are listed (the library store
+/// page is never listed).
+pub(crate) fn apply_indexed(doc: &mut Document, change: &Change, index: &mut Index) -> Result<bool, Error> {
     match change {
         Change::SetDoc { prop, value } => {
             if set_scalar(doc, *prop, value)? {
                 return Ok(*prop == DocProp::Libraries);
             }
             match prop {
-                DocProp::Styles => STYLES.set_whole(doc, value),
-                DocProp::StyleOrder => STYLES.reorder(doc, value),
-                DocProp::Comments => COMMENTS.set_whole(doc, value),
-                DocProp::CommentOrder => COMMENTS.reorder(doc, value),
-                DocProp::Collections => COLLECTIONS.set_whole(doc, value),
-                DocProp::CollectionOrder => COLLECTIONS.reorder(doc, value),
-                DocProp::Variables => VARIABLES.set_whole(doc, value),
-                DocProp::VariableOrder => VARIABLES.reorder(doc, value),
+                DocProp::Styles => STYLES.set_whole(doc, value, index),
+                DocProp::StyleOrder => STYLES.reorder(doc, value, index),
+                DocProp::Comments => COMMENTS.set_whole(doc, value, index),
+                DocProp::CommentOrder => COMMENTS.reorder(doc, value, index),
+                DocProp::Collections => COLLECTIONS.set_whole(doc, value, index),
+                DocProp::CollectionOrder => COLLECTIONS.reorder(doc, value, index),
+                DocProp::Variables => VARIABLES.set_whole(doc, value, index),
+                DocProp::VariableOrder => VARIABLES.reorder(doc, value, index),
                 _ => Err(Error::Invalid(format!("{prop:?} is not a list"))),
             }?;
         }
         Change::Put { entry, value } => match entry {
-            Entry::Style(k) => STYLES.put(doc, k.clone(), value)?,
-            Entry::Comment(k) => COMMENTS.put(doc, *k, value)?,
-            Entry::Collection(k) => COLLECTIONS.put(doc, k.clone(), value)?,
-            Entry::Variable(k) => VARIABLES.put(doc, k.clone(), value)?,
+            Entry::Style(k) => STYLES.put(doc, k.clone(), value, index)?,
+            Entry::Comment(k) => COMMENTS.put(doc, *k, value, index)?,
+            Entry::Collection(k) => COLLECTIONS.put(doc, k.clone(), value, index)?,
+            Entry::Variable(k) => VARIABLES.put(doc, k.clone(), value, index)?,
             Entry::DevSnapshot(k) => {
                 let snap: DevSnapshot = value.read("dev snapshot")?;
                 doc.dev_snapshots.insert(*k, Arc::new(snap));
@@ -351,14 +441,18 @@ pub(crate) fn apply(doc: &mut Document, change: &Change) -> Result<bool, Error>
             }
         },
         Change::Remove { entry } => match entry {
-            Entry::Style(k) => STYLES.remove(doc, k),
-            Entry::Comment(k) => COMMENTS.remove(doc, k),
-            Entry::Collection(k) => COLLECTIONS.remove(doc, k),
-            Entry::Variable(k) => VARIABLES.remove(doc, k),
+            Entry::Style(k) => STYLES.remove(doc, k, index),
+            Entry::Comment(k) => COMMENTS.remove(doc, k, index),
+            Entry::Collection(k) => COLLECTIONS.remove(doc, k, index),
+            Entry::Variable(k) => VARIABLES.remove(doc, k, index),
             Entry::DevSnapshot(k) => {
                 doc.dev_snapshots.remove(k);
             }
             Entry::Image(k) => {
+                // Nothing to remove: no need to look through the document for what shows it.
+                if !doc.images.contains_key(k) {
+                    return Ok(false);
+                }
                 // An image a layer or style shows stays: removing it would let other bytes go
                 // under its key, swapping the picture that layer shows ([`image_for`]).
                 if figcraft_doc::used_images(doc).contains(k) {

crates/collab/src/keyed_list_tests.rs

@@ -0,0 +1,132 @@
+//! Keyed document lists (comments, styles, variables, collections) applied in large batches: a
+//! batch of new entries is placed through an index, not a scan of the growing list per entry, and
+//! puts, removes and reorders in one batch still leave the list a scan would.
+
+use figcraft_doc::{Comment, Document, Message};
+
+use crate::change::{Change, DocProp, Entry, Json};
+use crate::{Mode, Replica};
+
+fn comment(id: u64, text: &str, doc: &Document) -> Comment {
+    let page = doc.first_page().unwrap();
+    Comment { id, page, x: 1.0, y: 2.0, anchor: None, first: Message::new("A", text, 0), resolved: false, replies: Vec::new(), read_by: None }
+}
+
+fn put(c: &Comment) -> Change {
+    Change::Put { entry: Entry::Comment(c.id), value: Json::of(c).unwrap() }
+}
+
+#[test]
+fn a_large_batch_of_new_comments_lands_whole_and_in_order() {
+    let doc = Document::new("t");
+    let changes: Vec<Change> = (1..=100_000).map(|i| put(&comment(i, "hi", &doc))).collect();
+    let mut r = Replica::new(doc);
+    // A scan per put took minutes here; indexed it takes well under a second, even unoptimised.
+    r.apply(&changes, Mode::Strict).unwrap();
+    assert_eq!(r.document().comments.len(), 100_000);
+    assert!(r.document().comments.iter().zip(1..).all(|(c, i)| c.id == i));
+}
+
+#[test]
+fn puts_removes_and_reorders_in_one_batch_match_one_by_one() {
+    let doc = Document::new("t");
+    let c = |i, t| comment(i, t, &doc);
+    let batch = vec![
+        put(&c(1, "a")),
+        put(&c(2, "b")),
+        put(&c(3, "c")),
+        // Replaced in place.
+        put(&c(2, "B")),
+        Change::Remove { entry: Entry::Comment(1) },
+        // Absent: nothing changes.
+        Change::Remove { entry: Entry::Comment(9) },
+        put(&c(4, "d")),
+        put(&c(6, "f")),
+        // Two removals in a row: the second goes through the index the first kept current.
+        Change::Remove { entry: Entry::Comment(2) },
+        Change::Remove { entry: Entry::Comment(6) },
+        put(&c(2, "B")),
+        put(&c(3, "C")),
+        Change::SetDoc { prop: DocProp::CommentOrder, value: Json::of(&[4u64, 3, 2]).unwrap() },
+        put(&c(2, "BB")),
+        put(&c(1, "again")),
+        Change::SetDoc { prop: DocProp::Comments, value: Json::of(&vec![c(7, "x"), c(8, "y")]).unwrap() },
+        put(&c(8, "Y")),
+        put(&c(5, "e")),
+    ];
+    let mut together = Replica::new(doc.clone());
+    together.apply(&batch, Mode::Strict).unwrap();
+    let mut one_by_one = Replica::new(doc.clone());
+    for ch in &batch {
+        one_by_one.apply(std::slice::from_ref(ch), Mode::Strict).unwrap();
+    }
+    assert_eq!(together.document().comments, one_by_one.document().comments);
+    let got: Vec<(u64, &str)> = together.document().comments.iter().map(|c| (c.id, c.first.text.as_str())).collect();
+    assert_eq!(got, vec![(7, "x"), (8, "Y"), (5, "e")]);
+}
+
+#[test]
+fn removing_an_absent_image_does_not_fail() {
+    let mut doc = Document::new("t");
+    let remove = Change::Remove { entry: Entry::Image("missing".into()) };
+    assert!(!crate::docprops::apply(&mut doc, &remove).unwrap());
+}
+
+#[test]
+fn removals_in_one_batch_keep_the_index_current() {
+    let doc = Document::new("t");
+    let c = |i, t| comment(i, t, &doc);
+    let batch = vec![
+        put(&c(1, "a")),
+        put(&c(2, "b")),
+        put(&c(3, "c")),
+        put(&c(4, "d")),
+        put(&c(5, "e")),
+        Change::Remove { entry: Entry::Comment(2) },
+        Change::Remove { entry: Entry::Comment(9) },
+        Change::Remove { entry: Entry::Comment(4) },
+        // Replaced where the shifted index says they now are.
+        put(&c(5, "E")),
+        put(&c(3, "C")),
+        put(&c(1, "A")),
+        Change::Remove { entry: Entry::Comment(1) },
+        put(&c(2, "B")),
+        put(&c(5, "EE")),
+    ];
+    let mut together = Replica::new(doc.clone());
+    together.apply(&batch, Mode::Strict).unwrap();
+    let mut one_by_one = Replica::new(doc.clone());
+    for ch in &batch {
+        one_by_one.apply(std::slice::from_ref(ch), Mode::Strict).unwrap();
+    }
+    assert_eq!(together.document().comments, one_by_one.document().comments);
+    let got: Vec<(u64, &str)> = together.document().comments.iter().map(|c| (c.id, c.first.text.as_str())).collect();
+    assert_eq!(got, vec![(3, "C"), (5, "EE"), (2, "B")]);
+}
+
+#[test]
+fn duplicate_keys_from_a_whole_list_match_one_by_one() {
+    let doc = Document::new("t");
+    let c = |i, t| comment(i, t, &doc);
+    let batch = vec![
+        Change::SetDoc { prop: DocProp::Comments, value: Json::of(&vec![c(1, "a"), c(2, "b"), c(2, "b2"), c(3, "c"), c(4, "d")]).unwrap() },
+        // The first entry with the key is the one replaced, as a scan finds it.
+        put(&c(2, "B")),
+        put(&c(4, "D")),
+        // Both entries with the key go: the index is built again for what follows.
+        Change::Remove { entry: Entry::Comment(2) },
+        put(&c(3, "C")),
+        Change::Remove { entry: Entry::Comment(1) },
+        put(&c(2, "again")),
+        put(&c(4, "DD")),
+    ];
+    let mut together = Replica::new(doc.clone());
+    together.apply(&batch, Mode::Lenient).unwrap();
+    let mut one_by_one = Replica::new(doc.clone());
+    for ch in &batch {
+        one_by_one.apply(std::slice::from_ref(ch), Mode::Lenient).unwrap();
+    }
+    assert_eq!(together.document().comments, one_by_one.document().comments);
+    let got: Vec<(u64, &str)> = together.document().comments.iter().map(|c| (c.id, c.first.text.as_str())).collect();
+    assert_eq!(got, vec![(3, "C"), (4, "DD"), (2, "again")]);
+}

crates/collab/src/lib.rs

@@ -148,6 +148,8 @@ mod image_tests;
 #[cfg(test)]
 mod journal_tests;
 #[cfg(test)]
+mod keyed_list_tests;
+#[cfg(test)]
 mod merge_list_tests;
 #[cfg(test)]
 mod merge_tests;

crates/collab/src/replica.rs

@@ -192,8 +192,9 @@ impl Replica {
         self.doc.take_dirty();
         let mut out = Applied::default();
         let (mut relist, mut relist_all) = (BTreeSet::new(), false);
+        let mut index = docprops::Index::default();
         for c in changes {
-            match self.apply_one(c, &mut out) {
+            match self.apply_one(c, &mut out, &mut index) {
                 Ok(Relist::No) => {}
                 Ok(Relist::Page(p)) => {
                     relist.insert(p);
@@ -214,7 +215,7 @@ impl Replica {
     }
 
     /// Apply one change. Returns what of the page list needs rebuilding.
-    fn apply_one(&mut self, c: &Change, out: &mut Applied) -> Result<Relist, Error> {
+    fn apply_one(&mut self, c: &Change, out: &mut Applied, index: &mut docprops::Index) -> Result<Relist, Error> {
         match c {
             Change::Create { id, kind, parent, position, props } => {
                 frac::validate(position).map_err(|e| Error::Invalid(e.to_string()))?;
@@ -297,7 +298,7 @@ impl Replica {
                 Ok(if page { Relist::Page(*id) } else { Relist::No })
             }
             Change::SetDoc { .. } | Change::Put { .. } | Change::Remove { .. } => {
-                Ok(if docprops::apply(&mut self.doc, c)? { Relist::All } else { Relist::No })
+                Ok(if docprops::apply_indexed(&mut self.doc, c, index)? { Relist::All } else { Relist::No })
             }
         }
     }