test: add tree mem bench

crates/dev-utils/src/lib.rs

@@ -71,7 +71,7 @@ unsafe impl GlobalAlloc for Counter {
 static A: Counter = Counter;
 
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
-pub struct ByteSize(usize);
+pub struct ByteSize(pub usize);
 
 impl Debug for ByteSize {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

crates/examples/examples/outliner.rs

@@ -0,0 +1,103 @@
+use dev_utils::{get_mem_usage, ByteSize};
+use loro::LoroDoc;
+
+pub fn main() {
+    // Number of nodes
+    let n = 100_000;
+    // Number of moves per node
+    let k = 5;
+    // Max depth of the tree
+    let max_depth = 10;
+    let avg_peer_edits = 1000;
+    let peers = (n * k / avg_peer_edits).max(1);
+    println!("Number of nodes: {}", n);
+    println!("Number of moves per node: {}", k);
+    println!("Average number of peer edits: {}", avg_peer_edits);
+    println!("Number of peers: {}", peers);
+
+    let doc = LoroDoc::new();
+    let tree = doc.get_tree("tree");
+    let mut nodes = vec![];
+    let mut depth = vec![0; n];
+    for _ in 0..n {
+        let node = tree.create(None).unwrap();
+        nodes.push(node);
+    }
+
+    doc.commit();
+    println!(
+        "Memory usage after creating {} nodes: {}",
+        n,
+        get_mem_usage()
+    );
+
+    doc.compact_change_store();
+    println!(
+        "Memory usage after compacting change store: {}",
+        get_mem_usage()
+    );
+
+    println!(
+        "Updates size: {}",
+        ByteSize(doc.export_from(&Default::default()).len())
+    );
+    let snapshot = doc.export_snapshot();
+    println!("Snapshot size: {}", ByteSize(snapshot.len()));
+    doc.with_oplog(|oplog| {
+        println!(
+            "Change store kv size: {}",
+            ByteSize(oplog.change_store_kv_size())
+        );
+    });
+
+    // Move nodes around
+    for _ in (0..n * k).step_by(avg_peer_edits) {
+        let new_doc = doc.fork();
+        let new_tree = new_doc.get_tree("tree");
+
+        for _ in 0..avg_peer_edits {
+            let (mut i, mut j) = rand::random::<(usize, usize)>();
+            i %= n;
+            j %= n;
+            while depth[j] > max_depth {
+                j = rand::random::<usize>() % n;
+            }
+
+            if new_tree.mov_to(nodes[i], nodes[j], 0).is_ok() {
+                depth[i] = depth[j] + 1;
+            }
+        }
+
+        doc.import(&new_doc.export_from(&doc.oplog_vv())).unwrap();
+    }
+
+    let mem = get_mem_usage();
+    println!("Memory usage after moving {} nodes: {}", n, mem);
+
+    doc.compact_change_store();
+    let mem_after_compact = get_mem_usage();
+    println!(
+        "Memory usage after compacting change store: {}",
+        mem_after_compact
+    );
+
+    doc.free_diff_calculator();
+    println!(
+        "Memory usage after freeing diff calculator: {}",
+        get_mem_usage()
+    );
+
+    println!(
+        "Updates size: {}",
+        ByteSize(doc.export_from(&Default::default()).len())
+    );
+    let snapshot = doc.export_snapshot();
+    println!("Snapshot size: {}", ByteSize(snapshot.len()));
+    doc.compact_change_store();
+    doc.with_oplog(|oplog| {
+        println!(
+            "Change store kv size: {}",
+            ByteSize(oplog.change_store_kv_size())
+        );
+    });
+}

crates/loro-internal/src/oplog.rs

@@ -751,6 +751,11 @@ impl OpLog {
     pub fn compact_change_store(&mut self) {
         self.change_store.flush_and_compact();
     }
+
+    #[inline]
+    pub fn change_store_kv_size(&self) -> usize {
+        self.change_store.kv_size()
+    }
 }
 
 #[derive(Debug)]

crates/loro-internal/src/oplog/change_store.rs

@@ -435,6 +435,15 @@ impl ChangeStore {
             merge_interval,
         }
     }
+
+    pub fn kv_size(&self) -> usize {
+        self.external_kv
+            .lock()
+            .unwrap()
+            .scan(Bound::Unbounded, Bound::Unbounded)
+            .map(|(k, v)| k.len() + v.len())
+            .sum()
+    }
 }
 
 #[derive(Clone, Debug)]

crates/loro/src/lib.rs

@@ -1247,6 +1247,8 @@ impl Default for LoroText {
 /// LoroTree container. It's used to model movable trees.
 ///
 /// You may use it to model directories, outline or other movable hierarchical data.
+///
+/// Learn more at https://loro.dev/docs/tutorial/tree
 #[derive(Clone, Debug)]
 pub struct LoroTree {
     handler: InnerTreeHandler,
@@ -1323,6 +1325,11 @@ impl LoroTree {
         self.handler.create_at(parent, index)
     }
 
+    /// Get the root nodes of the forest.
+    pub fn roots(&self) -> Vec<TreeID> {
+        self.handler.roots()
+    }
+
     /// Create a new tree node at the given index and return the [`TreeID`].
     ///
     /// If the `parent` is `None`, the created node is the root of a tree.
@@ -1472,8 +1479,8 @@ impl LoroTree {
     ///
     /// - If the target node does not exist, return `None`.
     /// - If the target node is a root node, return `Some(None)`.
-    pub fn parent(&self, target: &TreeID) -> Option<Option<TreeID>> {
-        self.handler.get_node_parent(target)
+    pub fn parent(&self, target: TreeID) -> Option<Option<TreeID>> {
+        self.handler.get_node_parent(&target)
     }
 
     /// Return whether target node exists.
@@ -1504,9 +1511,9 @@ impl LoroTree {
     }
 
     /// Return the fractional index of the target node with hex format.
-    pub fn fractional_index(&self, target: &TreeID) -> Option<String> {
+    pub fn fractional_index(&self, target: TreeID) -> Option<String> {
         self.handler
-            .get_position_by_tree_id(target)
+            .get_position_by_tree_id(&target)
             .map(|x| x.to_string())
     }