diff --git a/src/builder.rs b/src/builder.rs index 5e3e620..61741a0 100644 --- a/src/builder.rs +++ b/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 { stack: Vec>>, + /// 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, /// Cached value of `opt_packing_depth`. packing_depth: usize, + /// Cached value of capacity: 2^(depth + packing_depth). + capacity: usize, } impl Builder { - pub fn new(depth: usize, level: usize) -> Self { - Self { - stack: Vec::with_capacity(depth), - depth, - level, - length: Length(0), - packing_factor: opt_packing_factor::(), - packing_depth: opt_packing_depth::().unwrap_or(0), + pub fn new(depth: usize, level: usize) -> Result { + let packing_depth = opt_packing_depth::().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::(), + 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 Builder { } pub fn push_node(&mut self, node: Arc>, 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 Builder { 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 Builder { } } - 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 Builder { Ok((tree, self.depth, self.length)) } } + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn depth_upper_limit() { + assert_eq!( + Builder::::new(62, 0).unwrap_err(), + Error::BuilderInvalidDepth { depth: 62 } + ); + assert_eq!(Builder::::new(61, 0).unwrap().depth, 61); + } +} diff --git a/src/error.rs b/src/error.rs index 8d88d39..1757bac 100644 --- a/src/error.rs +++ b/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, diff --git a/src/lib.rs b/src/lib.rs index 33a167d..f093028 100644 --- a/src/lib.rs +++ b/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 {} diff --git a/src/list.rs b/src/list.rs index 5a35b45..bb20dc6 100644 --- a/src/list.rs +++ b/src/list.rs @@ -73,12 +73,12 @@ impl> List { Self::try_from_iter(std::iter::repeat(elem).take(n)) } - pub fn builder() -> Builder { + pub fn builder() -> Result, Error> { Builder::new(Self::depth(), 0) } pub fn try_from_iter(iter: impl IntoIterator) -> Result { - let mut builder = Self::builder(); + let mut builder = Self::builder()?; for item in iter.into_iter() { builder.push(item)?; @@ -204,7 +204,7 @@ impl> List { let depth = Self::depth(); let packing_depth = opt_packing_depth::().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() {