Fix bugs related to depth and capacity overflows (#39)

* Fix bugs in Builder

* Check capacity when pushing

* Remove type-length bound, it's too clunky

* Satisfy clippy

src/builder.rs

@@ -1,30 +1,45 @@
 use crate::utils::{opt_packing_depth, opt_packing_factor, Length, MaybeArced};
-use crate::{Arc, Error, PackedLeaf, Tree, Value};
+use crate::{Arc, Error, PackedLeaf, Tree, Value, MAX_TREE_DEPTH};
 
+#[derive(Debug)]
 pub struct Builder<T: Value> {
     stack: Vec<MaybeArced<Tree<T>>>,
+    /// The depth of the tree excluding the packing depth.
     depth: usize,
+    /// The level (depth) in the tree at which nodes
     level: usize,
     length: Length,
     /// Cached value of `opt_packing_factor`.
     packing_factor: Option<usize>,
     /// Cached value of `opt_packing_depth`.
     packing_depth: usize,
+    /// Cached value of capacity: 2^(depth + packing_depth).
+    capacity: usize,
 }
 
 impl<T: Value> Builder<T> {
-    pub fn new(depth: usize, level: usize) -> Self {
-        Self {
-            stack: Vec::with_capacity(depth),
-            depth,
-            level,
-            length: Length(0),
-            packing_factor: opt_packing_factor::<T>(),
-            packing_depth: opt_packing_depth::<T>().unwrap_or(0),
+    pub fn new(depth: usize, level: usize) -> Result<Self, Error> {
+        let packing_depth = opt_packing_depth::<T>().unwrap_or(0);
+        if depth.saturating_add(packing_depth) > MAX_TREE_DEPTH {
+            Err(Error::BuilderInvalidDepth { depth })
+        } else {
+            let capacity = 1 << (depth + packing_depth);
+            Ok(Self {
+                stack: Vec::with_capacity(depth),
+                depth,
+                level,
+                length: Length(0),
+                packing_factor: opt_packing_factor::<T>(),
+                packing_depth,
+                capacity,
+            })
         }
     }
 
     pub fn push(&mut self, value: T) -> Result<(), Error> {
+        if self.length.as_usize() == self.capacity {
+            return Err(Error::BuilderFull);
+        }
         let index = self.length.as_usize();
         let next_index = index + 1;
 
@@ -59,6 +74,10 @@ impl<T: Value> Builder<T> {
     }
 
     pub fn push_node(&mut self, node: Arc<Tree<T>>, len: usize) -> Result<(), Error> {
+        if self.length.as_usize() == self.capacity {
+            return Err(Error::BuilderFull);
+        }
+
         let index_on_level = self.length.as_usize() >> self.level;
         let next_index_on_level = index_on_level + 1;
 
@@ -93,7 +112,6 @@ impl<T: Value> Builder<T> {
             return Ok((Tree::zero(self.depth), self.depth, Length(0)));
         }
 
-        let capacity = 2usize.pow((self.depth + self.packing_depth) as u32);
         let length = self.length.as_usize();
         let level_capacity = 1 << self.level;
         let mut next_index_on_level = (length + level_capacity - 1) / level_capacity;
@@ -123,7 +141,7 @@ impl<T: Value> Builder<T> {
             }
         }
 
-        while next_index_on_level << self.level != capacity {
+        while next_index_on_level << self.level != self.capacity {
             // Push a new zero padding node on the right of the top-most stack element.
             let depth = (next_index_on_level.trailing_zeros() as usize)
                 .saturating_add(self.level)
@@ -168,3 +186,17 @@ impl<T: Value> Builder<T> {
         Ok((tree, self.depth, self.length))
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn depth_upper_limit() {
+        assert_eq!(
+            Builder::<u64>::new(62, 0).unwrap_err(),
+            Error::BuilderInvalidDepth { depth: 62 }
+        );
+        assert_eq!(Builder::<u64>::new(61, 0).unwrap().depth, 61);
+    }
+}

src/error.rs

@@ -17,6 +17,7 @@ pub enum Error {
     UpdateLeavesError,
     InvalidRebaseNode,
     InvalidRebaseLeaf,
+    BuilderInvalidDepth { depth: usize },
     BuilderExpectedLeaf,
     BuilderStackEmptyMerge,
     BuilderStackEmptyMergeLeft,
@@ -26,6 +27,7 @@ pub enum Error {
     BuilderStackEmptyFinishRight,
     BuilderStackEmptyFinalize,
     BuilderStackLeftover,
+    BuilderFull,
     BulkUpdateUnclean,
     CowMissingEntry,
     LevelIterPendingUpdates,

src/lib.rs

@@ -33,6 +33,13 @@ pub use vector::Vector;
 use ssz::{Decode, Encode};
 use tree_hash::TreeHash;
 
+/// Maximum depth for a tree.
+///
+/// We limit trees to 2^63 elements so we can avoid overflow when calculating 2^depth.
+pub const MAX_TREE_DEPTH: usize = u64::BITS as usize - 1;
+
+pub const MAX_TREE_LENGTH: usize = 1 << MAX_TREE_DEPTH;
+
 #[cfg(feature = "debug")]
 pub trait Value: Encode + Decode + TreeHash + PartialEq + Clone + std::fmt::Debug {}
 

src/list.rs

@@ -73,12 +73,12 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
         Self::try_from_iter(std::iter::repeat(elem).take(n))
     }
 
-    pub fn builder() -> Builder<T> {
+    pub fn builder() -> Result<Builder<T>, Error> {
         Builder::new(Self::depth(), 0)
     }
 
     pub fn try_from_iter(iter: impl IntoIterator<Item = T>) -> Result<Self, Error> {
-        let mut builder = Self::builder();
+        let mut builder = Self::builder()?;
 
         for item in iter.into_iter() {
             builder.push(item)?;
@@ -204,7 +204,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
         let depth = Self::depth();
         let packing_depth = opt_packing_depth::<T>().unwrap_or(0);
         let level = compute_level(n, depth, packing_depth);
-        let mut builder = Builder::new(Self::depth(), level);
+        let mut builder = Builder::new(Self::depth(), level)?;
         let mut level_iter = self.level_iter_from(n)?.peekable();
 
         while let Some(item) = level_iter.next() {