perf improvements

openmls/src/binary_tree/array_representation/diff.rs

@@ -222,7 +222,7 @@ impl<'a, L: Clone + Debug + Default, P: Clone + Debug + Default> AbDiff<'a, L, P
     }
 
     /// Returns the copath of a leaf node.
-    pub(crate) fn copath(&self, leaf_index: LeafNodeIndex) -> Vec<TreeNodeIndex> {
+    pub(crate) fn copath(&self, leaf_index: LeafNodeIndex) -> impl Iterator<Item = TreeNodeIndex> {
         copath(leaf_index, self.size())
     }
 

openmls/src/binary_tree/array_representation/treemath.rs

@@ -244,6 +244,10 @@ impl TreeSize {
             self.0 = 0;
         }
     }
+
+    pub(crate) fn height(&self) -> usize {
+        log2(self.0)
+    }
 }
 
 #[test]
@@ -363,33 +367,32 @@ pub(crate) fn test_sibling(index: TreeNodeIndex) -> TreeNodeIndex {
 pub(crate) fn direct_path(node_index: LeafNodeIndex, size: TreeSize) -> Vec<ParentNodeIndex> {
     let r = root(size).u32();
 
-    let mut d = vec![];
+    let mut direct_path = Vec::with_capacity(size.height());
     let mut x = node_index.to_tree_index();
     while x != r {
         let parent = parent(TreeNodeIndex::new(x));
-        d.push(parent);
+        direct_path.push(parent);
         x = parent.to_tree_index();
     }
-    d
+    direct_path
 }
 
 /// Copath of a leaf node.
-pub(crate) fn copath(leaf_index: LeafNodeIndex, size: TreeSize) -> Vec<TreeNodeIndex> {
+pub(crate) fn copath(
+    leaf_index: LeafNodeIndex,
+    size: TreeSize,
+) -> impl Iterator<Item = TreeNodeIndex> {
     // Start with leaf
-    let mut full_path = vec![TreeNodeIndex::Leaf(leaf_index)];
+    let mut full_path = Vec::with_capacity(size.height());
+    full_path.push(TreeNodeIndex::Leaf(leaf_index));
     let mut direct_path = direct_path(leaf_index, size);
     if !direct_path.is_empty() {
-        // Remove root
-        direct_path.pop();
+        direct_path.pop(); // Remove root
     }
-    full_path.append(
-        &mut direct_path
-            .iter()
-            .map(|i| TreeNodeIndex::Parent(*i))
-            .collect(),
-    );
+    let parent_direct_path = direct_path.into_iter().map(TreeNodeIndex::Parent);
+    full_path.extend(parent_direct_path);
 
-    full_path.into_iter().map(sibling).collect()
+    full_path.into_iter().map(sibling)
 }
 
 /// Common ancestor of two leaf nodes, aka the node where their direct paths
@@ -426,7 +429,7 @@ pub(crate) fn common_direct_path(
     x_path.reverse();
     y_path.reverse();
 
-    let mut common_path = vec![];
+    let mut common_path = Vec::with_capacity(size.height());
 
     for (x, y) in x_path.iter().zip(y_path.iter()) {
         if x == y {
@@ -459,7 +462,7 @@ fn test_node_in_tree() {
     for test in tests.iter() {
         assert!(is_node_in_tree(
             TreeNodeIndex::new(test.0),
-            TreeSize::new(test.1)
+            TreeSize::new(test.1),
         ));
     }
 }
@@ -470,7 +473,7 @@ fn test_node_not_in_tree() {
     for test in tests.iter() {
         assert!(!is_node_in_tree(
             TreeNodeIndex::new(test.0),
-            TreeSize::new(test.1)
+            TreeSize::new(test.1),
         ));
     }
 }

openmls/src/treesync/diff.rs

@@ -126,10 +126,9 @@ impl<'a> TreeSyncDiff<'a> {
         let copath_resolutions = self.copath_resolutions(leaf_index);
 
         // The two vectors should have the same length
-        debug_assert_eq!(copath.len(), copath_resolutions.len());
+        // debug_assert_eq!(copath.cloned().count(), copath_resolutions.len());
 
         copath
-            .into_iter()
             .zip(copath_resolutions)
             .filter_map(|(index, resolution)| {
                 // Filter out the nodes whose copath resolution is empty
@@ -198,13 +197,14 @@ impl<'a> TreeSyncDiff<'a> {
     /// Find and return the index of either the left-most blank leaf, or, if
     /// there are no blank leaves, the leaf count.
     pub(crate) fn free_leaf_index(&self) -> LeafNodeIndex {
-        let leaf_count = self.diff.leaves().count() as u32;
+        let mut leaf_count = 0;
 
         // Search for blank leaves in existing leaves
         for (leaf_index, leaf_id) in self.diff.leaves() {
             if leaf_id.node().is_none() {
                 return leaf_index;
             }
+            leaf_count += 1;
         }
 
         // Return the next free virtual blank leaf
@@ -538,7 +538,6 @@ impl<'a> TreeSyncDiff<'a> {
         // each node.
         self.diff
             .copath(leaf_index)
-            .into_iter()
             .map(|node_index| self.resolution(node_index, &HashSet::new()))
             .collect()
     }

openmls/src/treesync/mod.rs

@@ -451,10 +451,7 @@ impl TreeSync {
         // Set the leaf indices in all the leaves and convert the node types.
         for (node_index, node_option) in ratchet_tree.0.into_iter().enumerate() {
             let ts_node_option: TreeNode<TreeSyncLeafNode, TreeSyncParentNode> = match node_option {
-                Some(node) => {
-                    let node = node.clone();
-                    TreeSyncNode::from(node).into()
-                }
+                Some(node) => TreeSyncNode::from(node).into(),
                 None => {
                     if node_index % 2 == 0 {
                         TreeNode::Leaf(TreeSyncLeafNode::blank())