Update keys iterator to iterate at the tree level

No more big vec allocation now!

automerge-wasm/src/lib.rs

@@ -104,12 +104,13 @@ impl Automerge {
     pub fn keys(&mut self, obj: JsValue, heads: Option<Array>) -> Result<Array, JsValue> {
         let obj = self.import(obj)?;
         let result = if let Some(heads) = get_heads(heads) {
-            self.0.keys_at(&obj, &heads)
+            self.0
+                .keys_at(&obj, &heads)
+                .map(|s| JsValue::from_str(&s))
+                .collect()
         } else {
-            self.0.keys(&obj)
-        }
-        .map(|s| JsValue::from_str(&s))
-        .collect();
+            self.0.keys(&obj).map(|s| JsValue::from_str(&s)).collect()
+        };
         Ok(result)
     }
 

automerge/src/autocommit.rs

@@ -1,11 +1,12 @@
 use crate::exid::ExId;
+use crate::op_set::B;
 use crate::transaction::{CommitOptions, Transactable};
 use crate::types::Patch;
 use crate::{
     change::export_change, transaction::TransactionInner, ActorId, Automerge, AutomergeError,
     Change, ChangeHash, Prop, Value,
 };
-use crate::{Keys, SyncMessage, SyncState};
+use crate::{Keys, KeysAt, SyncMessage, SyncState};
 
 /// An automerge document that automatically manages transactions.
 #[derive(Debug, Clone)]
@@ -291,11 +292,11 @@ impl Transactable for AutoCommit {
     // PropAt::()
     // NthAt::()
 
-    fn keys(&self, obj: &ExId) -> Keys {
+    fn keys(&self, obj: &ExId) -> Keys<{ B }> {
         self.doc.keys(obj)
     }
 
-    fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> Keys {
+    fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> KeysAt<{ B }> {
         self.doc.keys_at(obj, heads)
     }
 

automerge/src/automerge.rs

@@ -4,6 +4,7 @@ use crate::change::encode_document;
 use crate::exid::ExId;
 use crate::keys::Keys;
 use crate::op_set::OpSet;
+use crate::op_set::B;
 use crate::transaction::{
     CommitOptions, Transaction, TransactionFailure, TransactionInner, TransactionResult,
     TransactionSuccess,
@@ -12,6 +13,7 @@ use crate::types::{
     ActorId, ChangeHash, Clock, ElemId, Export, Exportable, Key, ObjId, Op, OpId, OpType, Patch,
     ScalarValue, Value,
 };
+use crate::KeysAt;
 use crate::{legacy, query, types, ObjType};
 use crate::{AutomergeError, Change, Prop};
 use serde::Serialize;
@@ -190,30 +192,29 @@ impl Automerge {
     ///
     /// For a map this returns the keys of the map.
     /// For a list this returns the element ids (opids) encoded as strings.
-    pub fn keys(&self, obj: &ExId) -> Keys {
+    pub fn keys(&self, obj: &ExId) -> Keys<{ B }> {
         if let Ok(obj) = self.exid_to_obj(obj) {
-            let q = self.ops.search(obj, query::Keys::new());
-            Keys::new(self, q.keys)
+            let iter_keys = self.ops.keys(obj);
+            Keys::new(self, iter_keys)
         } else {
-            Keys::new(self, vec![])
+            Keys::new(self, None)
         }
     }
 
     /// Historical version of [`keys`](Self::keys).
-    pub fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> Keys {
+    pub fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> KeysAt<{ B }> {
         if let Ok(obj) = self.exid_to_obj(obj) {
             let clock = self.clock_at(heads);
-            let q = self.ops.search(obj, query::KeysAt::new(clock));
-            Keys::new(self, q.keys)
+            KeysAt::new(self, self.ops.keys_at(obj, clock))
         } else {
-            Keys::new(self, vec![])
+            KeysAt::new(self, None)
         }
     }
 
     pub fn length(&self, obj: &ExId) -> usize {
         if let Ok(inner_obj) = self.exid_to_obj(obj) {
             match self.ops.object_type(&inner_obj) {
-                Some(ObjType::Map) | Some(ObjType::Table) => self.keys(obj).len(),
+                Some(ObjType::Map) | Some(ObjType::Table) => self.keys(obj).count(),
                 Some(ObjType::List) | Some(ObjType::Text) => {
                     self.ops.search(inner_obj, query::Len::new()).len
                 }
@@ -228,7 +229,7 @@ impl Automerge {
         if let Ok(inner_obj) = self.exid_to_obj(obj) {
             let clock = self.clock_at(heads);
             match self.ops.object_type(&inner_obj) {
-                Some(ObjType::Map) | Some(ObjType::Table) => self.keys_at(obj, heads).len(),
+                Some(ObjType::Map) | Some(ObjType::Table) => self.keys_at(obj, heads).count(),
                 Some(ObjType::List) | Some(ObjType::Text) => {
                     self.ops.search(inner_obj, query::LenAt::new(clock)).len
                 }
@@ -1086,7 +1087,7 @@ mod tests {
         assert!(doc.value_at(&ROOT, "prop2", &heads5)?.unwrap().0 == Value::str("val3"));
         assert!(doc.value_at(&ROOT, "prop3", &heads5)?.unwrap().0 == Value::str("val4"));
 
-        assert_eq!(doc.keys_at(&ROOT, &[]).len(), 0);
+        assert_eq!(doc.keys_at(&ROOT, &[]).count(), 0);
         assert_eq!(doc.length_at(&ROOT, &[]), 0);
         assert!(doc.value_at(&ROOT, "prop1", &[])? == None);
         assert!(doc.value_at(&ROOT, "prop2", &[])? == None);

automerge/src/keys.rs

@@ -1,34 +1,21 @@
-use crate::{types::Key, Automerge};
+use crate::{query::IterKeys, Automerge};
 
-pub struct Keys<'a> {
-    index: usize,
-    keys: Vec<Key>,
+pub struct Keys<'a, 'k, const B: usize> {
+    keys: Option<IterKeys<'k, B>>,
     doc: &'a Automerge,
 }
 
-impl<'a> Keys<'a> {
-    pub(crate) fn new(doc: &'a Automerge, keys: Vec<Key>) -> Self {
-        Self {
-            index: 0,
-            keys,
-            doc,
-        }
+impl<'a, 'k, const B: usize> Keys<'a, 'k, B> {
+    pub(crate) fn new(doc: &'a Automerge, keys: Option<IterKeys<'k, B>>) -> Self {
+        Self { keys, doc }
     }
 }
 
-impl<'a> Iterator for Keys<'a> {
+impl<'a, 'k, const B: usize> Iterator for Keys<'a, 'k, B> {
     type Item = String;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let n = self.keys.get(self.index)?;
-        self.index += 1;
-        Some(self.doc.to_string(*n))
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let num_left = self.keys.len() - self.index;
-        (num_left, Some(num_left))
+        let key = self.keys.as_mut()?.next()?;
+        Some(self.doc.to_string(key))
     }
 }
-
-impl<'a> ExactSizeIterator for Keys<'a> {}

automerge/src/keys_at.rs

@@ -0,0 +1,21 @@
+use crate::{query::IterKeysAt, Automerge};
+
+pub struct KeysAt<'a, 'k, const B: usize> {
+    keys: Option<IterKeysAt<'k, B>>,
+    doc: &'a Automerge,
+}
+
+impl<'a, 'k, const B: usize> KeysAt<'a, 'k, B> {
+    pub(crate) fn new(doc: &'a Automerge, keys: Option<IterKeysAt<'k, B>>) -> Self {
+        Self { keys, doc }
+    }
+}
+
+impl<'a, 'k, const B: usize> Iterator for KeysAt<'a, 'k, B> {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let key = self.keys.as_mut()?.next()?;
+        Some(self.doc.to_string(key))
+    }
+}

automerge/src/lib.rs

@@ -35,6 +35,7 @@ mod error;
 mod exid;
 mod indexed_cache;
 mod keys;
+mod keys_at;
 mod legacy;
 mod op_set;
 mod op_tree;
@@ -52,6 +53,7 @@ pub use change::{decode_change, Change};
 pub use error::AutomergeError;
 pub use exid::ExId as ObjId;
 pub use keys::Keys;
+pub use keys_at::KeysAt;
 pub use legacy::Change as ExpandedChange;
 pub use sync::{BloomFilter, SyncHave, SyncMessage, SyncState};
 pub use types::{ActorId, ChangeHash, ObjType, OpType, Prop};

automerge/src/op_set.rs

@@ -1,13 +1,15 @@
+use crate::clock::Clock;
 use crate::indexed_cache::IndexedCache;
 use crate::op_tree::OpTreeInternal;
-use crate::query::TreeQuery;
+use crate::query::{self, TreeQuery};
 use crate::types::{ActorId, Key, ObjId, Op, OpId, OpType};
 use crate::ObjType;
 use fxhash::FxBuildHasher;
 use std::cmp::Ordering;
 use std::collections::HashMap;
 
-pub(crate) type OpSet = OpSetInternal<16>;
+pub(crate) const B: usize = 16;
+pub(crate) type OpSet = OpSetInternal<B>;
 
 #[derive(Debug, Clone)]
 pub(crate) struct OpSetInternal<const B: usize> {
@@ -41,6 +43,22 @@ impl<const B: usize> OpSetInternal<B> {
         }
     }
 
+    pub fn keys(&self, obj: ObjId) -> Option<query::IterKeys<B>> {
+        if let Some((_typ, tree)) = self.trees.get(&obj) {
+            tree.keys()
+        } else {
+            None
+        }
+    }
+
+    pub fn keys_at(&self, obj: ObjId, clock: Clock) -> Option<query::IterKeysAt<B>> {
+        if let Some((_typ, tree)) = self.trees.get(&obj) {
+            tree.keys_at(clock)
+        } else {
+            None
+        }
+    }
+
     pub fn search<Q>(&self, obj: ObjId, query: Q) -> Q
     where
         Q: TreeQuery<B>,

automerge/src/op_tree.rs

@@ -5,8 +5,11 @@ use std::{
 };
 
 pub(crate) use crate::op_set::OpSetMetadata;
-use crate::query::{Index, QueryResult, TreeQuery};
 use crate::types::{Op, OpId};
+use crate::{
+    clock::Clock,
+    query::{self, Index, QueryResult, TreeQuery},
+};
 use std::collections::HashSet;
 
 #[allow(dead_code)]
@@ -36,6 +39,16 @@ impl<const B: usize> OpTreeInternal<B> {
         self.root_node.as_ref().map_or(0, |n| n.len())
     }
 
+    pub fn keys(&self) -> Option<query::IterKeys<B>> {
+        self.root_node.as_ref().map(query::IterKeys::new)
+    }
+
+    pub fn keys_at(&self, clock: Clock) -> Option<query::IterKeysAt<B>> {
+        self.root_node
+            .as_ref()
+            .map(|root| query::IterKeysAt::new(root, clock))
+    }
+
     pub fn search<Q>(&self, mut query: Q, m: &OpSetMetadata) -> Q
     where
         Q: TreeQuery<B>,

automerge/src/query.rs

@@ -20,8 +20,8 @@ mod prop_at;
 mod seek_op;
 
 pub(crate) use insert::InsertNth;
-pub(crate) use keys::Keys;
-pub(crate) use keys_at::KeysAt;
+pub(crate) use keys::IterKeys;
+pub(crate) use keys_at::IterKeysAt;
 pub(crate) use len::Len;
 pub(crate) use len_at::LenAt;
 pub(crate) use list_vals::ListVals;

automerge/src/query/keys.rs

@@ -1,29 +1,36 @@
 use crate::op_tree::OpTreeNode;
-use crate::query::{QueryResult, TreeQuery};
 use crate::types::Key;
 use std::fmt::Debug;
 
-#[derive(Debug, Clone, PartialEq)]
-pub(crate) struct Keys<const B: usize> {
-    pub keys: Vec<Key>,
+#[derive(Debug)]
+pub(crate) struct IterKeys<'a, const B: usize> {
+    index: usize,
+    last_key: Option<Key>,
+    root_child: &'a OpTreeNode<B>,
 }
 
-impl<const B: usize> Keys<B> {
-    pub fn new() -> Self {
-        Keys { keys: vec![] }
+impl<'a, const B: usize> IterKeys<'a, B> {
+    pub(crate) fn new(root_child: &'a OpTreeNode<B>) -> Self {
+        Self {
+            index: 0,
+            last_key: None,
+            root_child,
+        }
     }
 }
 
-impl<const B: usize> TreeQuery<B> for Keys<B> {
-    fn query_node(&mut self, child: &OpTreeNode<B>) -> QueryResult {
-        let mut last = None;
-        for i in 0..child.len() {
-            let op = child.get(i).unwrap();
-            if Some(op.key) != last && op.visible() {
-                self.keys.push(op.key);
-                last = Some(op.key);
+impl<'a, const B: usize> Iterator for IterKeys<'a, B> {
+    type Item = Key;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        for i in self.index..self.root_child.len() {
+            let op = self.root_child.get(i)?;
+            self.index += 1;
+            if Some(op.key) != self.last_key && op.visible() {
+                self.last_key = Some(op.key);
+                return Some(op.key);
             }
         }
-        QueryResult::Finish
+        None
     }
 }

automerge/src/query/keys_at.rs

@@ -1,36 +1,42 @@
-use crate::query::{QueryResult, TreeQuery, VisWindow};
-use crate::types::{Clock, Key, Op};
+use crate::op_tree::OpTreeNode;
+use crate::query::VisWindow;
+use crate::types::{Clock, Key};
 use std::fmt::Debug;
 
-#[derive(Debug, Clone, PartialEq)]
-pub(crate) struct KeysAt<const B: usize> {
+#[derive(Debug)]
+pub(crate) struct IterKeysAt<'a, const B: usize> {
     clock: Clock,
-    pub keys: Vec<Key>,
-    last: Option<Key>,
     window: VisWindow,
-    pos: usize,
+    index: usize,
+    last_key: Option<Key>,
+    root_child: &'a OpTreeNode<B>,
 }
 
-impl<const B: usize> KeysAt<B> {
-    pub fn new(clock: Clock) -> Self {
-        KeysAt {
+impl<'a, const B: usize> IterKeysAt<'a, B> {
+    pub(crate) fn new(root_child: &'a OpTreeNode<B>, clock: Clock) -> Self {
+        Self {
             clock,
-            pos: 0,
-            last: None,
-            keys: vec![],
-            window: Default::default(),
+            window: VisWindow::default(),
+            index: 0,
+            last_key: None,
+            root_child,
         }
     }
 }
 
-impl<const B: usize> TreeQuery<B> for KeysAt<B> {
-    fn query_element(&mut self, op: &Op) -> QueryResult {
-        let visible = self.window.visible_at(op, self.pos, &self.clock);
-        if Some(op.key) != self.last && visible {
-            self.keys.push(op.key);
-            self.last = Some(op.key);
+impl<'a, const B: usize> Iterator for IterKeysAt<'a, B> {
+    type Item = Key;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        for i in self.index..self.root_child.len() {
+            let op = self.root_child.get(i)?;
+            let visible = self.window.visible_at(op, self.index, &self.clock);
+            self.index += 1;
+            if Some(op.key) != self.last_key && visible {
+                self.last_key = Some(op.key);
+                return Some(op.key);
+            }
         }
-        self.pos += 1;
-        QueryResult::Next
+        None
     }
 }

automerge/src/transaction/manual_transaction.rs

@@ -1,5 +1,6 @@
 use crate::exid::ExId;
-use crate::{Automerge, ChangeHash, Prop, Value};
+use crate::op_set::B;
+use crate::{Automerge, ChangeHash, KeysAt, Prop, Value};
 use crate::{AutomergeError, Keys};
 
 use super::{CommitOptions, Transactable, TransactionInner};
@@ -134,11 +135,11 @@ impl<'a> Transactable for Transaction<'a> {
             .splice(self.doc, obj, pos, del, vals)
     }
 
-    fn keys(&self, obj: &ExId) -> Keys {
+    fn keys(&self, obj: &ExId) -> Keys<{ B }> {
         self.doc.keys(obj)
     }
 
-    fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> Keys {
+    fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> KeysAt<{ B }> {
         self.doc.keys_at(obj, heads)
     }
 

automerge/src/transaction/transactable.rs

@@ -1,5 +1,6 @@
 use crate::exid::ExId;
-use crate::{AutomergeError, ChangeHash, Keys, Prop, Value};
+use crate::op_set::B;
+use crate::{AutomergeError, ChangeHash, Keys, KeysAt, Prop, Value};
 use unicode_segmentation::UnicodeSegmentation;
 
 /// A way of mutating a document within a single change.
@@ -68,10 +69,10 @@ pub trait Transactable {
     }
 
     /// Get the keys of the given object, it should be a map.
-    fn keys(&self, obj: &ExId) -> Keys;
+    fn keys(&self, obj: &ExId) -> Keys<{ B }>;
 
     /// Get the keys of the given object at a point in history.
-    fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> Keys;
+    fn keys_at(&self, obj: &ExId, heads: &[ChangeHash]) -> KeysAt<{ B }>;
 
     /// Get the length of the given object.
     fn length(&self, obj: &ExId) -> usize;