shardtree: Use `Tree` constructors instead of struct creation

This makes it easier for us to track where new tree nodes are created,
and how the node kinds are used within the `ShardTree` data model.

shardtree/src/batch.rs

@@ -1,6 +1,6 @@
 //! Helpers for inserting many leaves into a tree at once.
 
-use std::{collections::BTreeMap, fmt, ops::Range, sync::Arc};
+use std::{collections::BTreeMap, fmt, ops::Range};
 
 use incrementalmerkletree::{Address, Hashable, Level, Position, Retention};
 use tracing::trace;
@@ -8,7 +8,7 @@ use tracing::trace;
 use crate::{
     error::{InsertionError, ShardTreeError},
     store::{Checkpoint, ShardStore},
-    IncompleteAt, LocatedPrunableTree, LocatedTree, Node, RetentionFlags, ShardTree, Tree,
+    IncompleteAt, LocatedPrunableTree, LocatedTree, RetentionFlags, ShardTree, Tree,
 };
 
 impl<
@@ -203,9 +203,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                 let rflags = RetentionFlags::from(retention);
                 let mut subtree = LocatedTree {
                     root_addr: Address::from(position),
-                    root: Tree(Node::Leaf {
+                    root: Tree::leaf((value.clone(), rflags)),
-                        value: (value.clone(), rflags),
-                    }),
                 };
 
                 if position.is_right_child() {
@@ -268,11 +266,7 @@ fn unite<H: Hashable + Clone + PartialEq>(
         root: if lroot.root_addr.level() < prune_below {
             Tree::unite(lroot.root_addr.level(), None, lroot.root, rroot.root)
         } else {
-            Tree(Node::Parent {
+            Tree::parent(None, lroot.root, rroot.root)
-                ann: None,
-                left: Arc::new(lroot.root),
-                right: Arc::new(rroot.root),
-            })
         },
     }
 }
@@ -297,7 +291,7 @@ fn combine_with_empty<H: Hashable + Clone + PartialEq>(
     });
     let sibling = LocatedTree {
         root_addr: sibling_addr,
-        root: Tree(Node::Nil),
+        root: Tree::empty(),
     };
     let (lroot, rroot) = if root.root_addr.is_left_child() {
         (root, sibling)

shardtree/src/lib.rs

@@ -424,7 +424,7 @@ impl<
                                         l_addr.level(),
                                         ann.clone(),
                                         new_left,
-                                        Tree(Node::Nil),
+                                        Tree::empty(),
                                     ),
                                     pos,
                                 )
@@ -444,9 +444,7 @@ impl<
                     )
                 }
                 Node::Leaf { value: (h, r) } => Some((
-                    Tree(Node::Leaf {
+                    Tree::leaf((h.clone(), *r | RetentionFlags::CHECKPOINT)),
-                        value: (h.clone(), *r | RetentionFlags::CHECKPOINT),
-                    }),
                     root_addr.max_position(),
                 )),
                 Node::Nil | Node::Pruned => None,
@@ -822,6 +820,10 @@ impl<
                             (None, None) => unreachable!(),
                         };
 
+                        // We don't use the `Tree::parent` constructor here, because it
+                        // creates `Arc`s for the child nodes internally, but if we don't
+                        // have a new child then we want to use the `Arc` for the existing
+                        // child.
                         let new_parent = Tree(Node::Parent {
                             ann: new_left
                                 .as_ref()

shardtree/src/prunable.rs

@@ -246,9 +246,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
                 (Tree(Node::Leaf { value: vl }), Tree(Node::Leaf { value: vr })) => {
                     if vl.0 == vr.0 {
                         // Merge the flags together.
-                        Ok(Tree(Node::Leaf {
+                        Ok(Tree::leaf((vl.0, vl.1 | vr.1)))
-                            value: (vl.0, vl.1 | vr.1),
-                        }))
                     } else {
                         trace!(left = ?vl.0, right = ?vr.0, "Merge conflict for leaves");
                         Err(addr)
@@ -314,7 +312,7 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
     /// `level` must be the level of the two nodes that are being joined.
     pub(crate) fn unite(level: Level, ann: Option<Arc<H>>, left: Self, right: Self) -> Self {
         match (left, right) {
-            (Tree(Node::Nil), Tree(Node::Nil)) => Tree(Node::Nil),
+            (Tree(Node::Nil), Tree(Node::Nil)) => Tree::empty(),
             (Tree(Node::Leaf { value: lv }), Tree(Node::Leaf { value: rv }))
                 // we can prune right-hand leaves that are not marked or reference leaves; if a
                 // leaf is a checkpoint then that information will be propagated to the replacement
@@ -322,18 +320,12 @@ impl<H: Hashable + Clone + PartialEq> PrunableTree<H> {
                 if lv.1 == RetentionFlags::EPHEMERAL &&
                     (rv.1 & (RetentionFlags::MARKED | RetentionFlags::REFERENCE)) == RetentionFlags::EPHEMERAL =>
             {
-                Tree(
+                Tree::leaf((H::combine(level, &lv.0, &rv.0), rv.1))
-                    Node::Leaf {
-                        value: (H::combine(level, &lv.0, &rv.0), rv.1),
-                    },
-                )
             }
-            (left, right) => Tree(
+            (left, right) => Tree::parent(
-                Node::Parent {
                     ann,
-                    left: Arc::new(left),
+                    left,
-                    right: Arc::new(right),
+                    right,
-                },
             ),
         }
     }
@@ -514,7 +506,7 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                         // to the left to truncate the left child and then reconstruct the
                         // node with `Nil` as the right sibling
                         go(position, l_child, left.as_ref()).map(|left| {
-                            Tree::unite(l_child.level(), ann.clone(), left, Tree(Node::Nil))
+                            Tree::unite(l_child.level(), ann.clone(), left, Tree::empty())
                         })
                     } else {
                         // we are truncating within the range of the right node, so recurse
@@ -584,9 +576,8 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                 // In the case that we are replacing a node entirely, we need to extend the
                 // subtree up to the level of the node being replaced, adding Nil siblings
                 // and recording the presence of those incomplete nodes when necessary
-                let replacement = |ann: Option<Arc<H>>,
+                let replacement =
-                                   mut node: LocatedPrunableTree<H>,
+                    |ann: Option<Arc<H>>, mut node: LocatedPrunableTree<H>, pruned: bool| {
-                                   pruned: bool| {
                         // construct the replacement node bottom-up
                         let mut incomplete = vec![];
                         while node.root_addr.level() < root_addr.level() {
@@ -594,22 +585,18 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                                 address: node.root_addr.sibling(),
                                 required_for_witness: contains_marked,
                             });
-                        let empty = Arc::new(Tree(if pruned { Node::Pruned } else { Node::Nil }));
+                            let empty = if pruned {
-                        let full = Arc::new(node.root);
+                                Tree::empty_pruned()
+                            } else {
+                                Tree::empty()
+                            };
+                            let full = node.root;
                             node = LocatedTree {
                                 root_addr: node.root_addr.parent(),
                                 root: if node.root_addr.is_right_child() {
-                                Tree(Node::Parent {
+                                    Tree::parent(None, empty, full)
-                                    ann: None,
-                                    left: empty,
-                                    right: full,
-                                })
                                 } else {
-                                Tree(Node::Parent {
+                                    Tree::parent(None, full, empty)
-                                    ann: None,
-                                    left: full,
-                                    right: empty,
-                                })
                                 },
                             };
                         }
@@ -936,16 +923,14 @@ impl<H: Hashable + Clone + PartialEq> LocatedPrunableTree<H> {
                         // a partially-pruned branch, and if it's a marked node then it will
                         // be a level-0 leaf.
                         match to_clear {
-                            [(_, flags)] => Tree(Node::Leaf {
+                            [(_, flags)] => Tree::leaf((h.clone(), *r & !*flags)),
-                                value: (h.clone(), *r & !*flags),
-                            }),
                             _ => {
                                 panic!("Tree state inconsistent with checkpoints.");
                             }
                         }
                     }
-                    Node::Nil => Tree(Node::Nil),
+                    Node::Nil => Tree::empty(),
-                    Node::Pruned => Tree(Node::Pruned),
+                    Node::Pruned => Tree::empty_pruned(),
                 }
             }
         }

shardtree/src/store/memory.rs

@@ -6,7 +6,7 @@ use std::convert::{Infallible, TryFrom};
 use incrementalmerkletree::Address;
 
 use super::{Checkpoint, ShardStore};
-use crate::{LocatedPrunableTree, LocatedTree, Node, PrunableTree, Tree};
+use crate::{LocatedPrunableTree, LocatedTree, PrunableTree, Tree};
 
 /// An implementation of [`ShardStore`] that stores all state in memory.
 ///
@@ -54,7 +54,7 @@ impl<H: Clone, C: Clone + Ord> ShardStore for MemoryShardStore<H, C> {
         {
             self.shards.push(LocatedTree {
                 root_addr: Address::from_parts(subtree_addr.level(), subtree_idx),
-                root: Tree(Node::Nil),
+                root: Tree::empty(),
             })
         }
 

shardtree/src/testing.rs

@@ -30,24 +30,17 @@ where
     A::Value: Clone + 'static,
     V::Value: Clone + 'static,
 {
-    let leaf = prop_oneof![
+    let leaf = prop_oneof![Just(Tree::empty()), arb_leaf.prop_map(Tree::leaf)];
-        Just(Tree(Node::Nil)),
-        arb_leaf.prop_map(|value| Tree(Node::Leaf { value }))
-    ];
 
     leaf.prop_recursive(depth, size, 2, move |inner| {
         (arb_annotation.clone(), inner.clone(), inner).prop_map(|(ann, left, right)| {
-            Tree(if left.is_nil() && right.is_nil() {
+            if left.is_nil() && right.is_nil() {
-                Node::Nil
+                Tree::empty()
             } else {
-                Node::Parent {
+                Tree::parent(ann, left, right)
-                    ann,
-                    left: Arc::new(left),
-                    right: Arc::new(right),
             }
         })
     })
-    })
 }
 
 pub fn arb_prunable_tree<H: Strategy + Clone + 'static>(

shardtree/src/tree.rs

@@ -315,7 +315,7 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
     pub fn empty(root_addr: Address) -> Self {
         Self {
             root_addr,
-            root: Tree(Node::Nil),
+            root: Tree::empty(),
         }
     }
 
@@ -324,7 +324,7 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
     pub fn with_root_value(root_addr: Address, value: V) -> Self {
         Self {
             root_addr,
-            root: Tree(Node::Leaf { value }),
+            root: Tree::leaf(value),
         }
     }
 
@@ -412,12 +412,10 @@ impl<A: Default + Clone, V: Clone> LocatedTree<A, V> {
 pub(crate) mod tests {
     use incrementalmerkletree::{Address, Level};
 
-    use super::{LocatedTree, Node, Tree};
+    use super::{LocatedTree, Tree};
 
     pub(crate) fn str_leaf<A>(c: &str) -> Tree<A, String> {
-        Tree(Node::Leaf {
+        Tree::leaf(c.to_string())
-            value: c.to_string(),
-        })
     }
 
     pub(crate) fn nil<A, B>() -> Tree<A, B> {