Implement iter_arc method for List

src/error.rs

@@ -34,6 +34,7 @@ pub enum Error {
     IntraRebaseZeroHash,
     IntraRebaseZeroDepth,
     IntraRebaseRepeatVisit,
+    PackedLeavesNoArc,
 }
 
 impl Display for Error {

src/interface.rs

@@ -1,3 +1,4 @@
+use crate::ArcIter;
 use crate::level_iter::LevelIter;
 use crate::update_map::UpdateMap;
 use crate::utils::{Length, updated_length};
@@ -22,6 +23,8 @@ pub trait ImmList<T: Value> {
     fn iter_from(&self, index: usize) -> Iter<'_, T>;
 
     fn level_iter_from(&self, index: usize) -> LevelIter<'_, T>;
+
+    fn iter_arc(&self, index: usize) -> Result<ArcIter<'_, T>, Error>;
 }
 
 pub trait MutList<T: Value>: ImmList<T> {
@@ -100,6 +103,10 @@ where
         self.iter_from(0)
     }
 
+    pub fn iter_arc(&self) -> Result<ArcIter<'_, T>, Error> {
+        self.backing.iter_arc(0)
+    }
+
     pub fn iter_from(&self, index: usize) -> InterfaceIter<'_, T, U> {
         InterfaceIter {
             tree_iter: self.backing.iter_from(index),

src/iter_arc.rs

@@ -0,0 +1,157 @@
+use tree_hash::{TreeHash, TreeHashType};
+use triomphe::Arc;
+
+use crate::{
+    Error, Leaf, Tree, UpdateMap, Value,
+    utils::{Length, opt_packing_depth},
+};
+
+#[derive(Debug)]
+pub struct ArcIter<'a, T: Value> {
+    /// Stack of tree nodes corresponding to the current position.
+    stack: Vec<&'a Tree<T>>,
+    /// The list index corresponding to the current position (next element to be yielded).
+    index: usize,
+    /// The `depth` of the root tree.
+    full_depth: usize,
+    /// Cached packing depth to avoid re-calculating `opt_packing_depth`.
+    packing_depth: usize,
+    /// Number of items that will be yielded by the iterator.
+    length: Length,
+}
+
+impl<'a, T: Value> ArcIter<'a, T> {
+    pub fn from_index(
+        index: usize,
+        root: &'a Tree<T>,
+        depth: usize,
+        length: Length,
+    ) -> Result<Self, Error> {
+        if <T as TreeHash>::tree_hash_type() == TreeHashType::Basic {
+            return Err(Error::PackedLeavesNoArc);
+        }
+        let mut stack = Vec::with_capacity(depth);
+        stack.push(root);
+
+        Ok(ArcIter {
+            stack,
+            index,
+            full_depth: depth,
+            packing_depth: opt_packing_depth::<T>().unwrap_or(0),
+            length,
+        })
+    }
+}
+
+impl<'a, T: Value> ArcIter<'a, T> {
+    pub fn new(root: &'a Tree<T>, depth: usize, length: Length) -> Self {
+        let mut stack = Vec::with_capacity(depth);
+        stack.push(root);
+
+        ArcIter {
+            stack,
+            index: 0,
+            full_depth: depth,
+            packing_depth: opt_packing_depth::<T>().unwrap_or(0),
+            length,
+        }
+    }
+}
+
+impl<'a, T: Value> Iterator for ArcIter<'a, T> {
+    type Item = &'a Arc<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.length.as_usize() {
+            return None;
+        }
+
+        match self.stack.last() {
+            None | Some(Tree::Zero(_)) => None,
+            Some(Tree::Leaf(Leaf { value, .. })) => {
+                let result = Some(value);
+
+                self.index += 1;
+
+                // Backtrack to the parent node of the next subtree
+                for _ in 0..=self.index.trailing_zeros() {
+                    self.stack.pop();
+                }
+
+                result
+            }
+            Some(Tree::PackedLeaf(_)) => {
+                // Return None case of PackedLeaf
+                None
+            }
+            Some(Tree::Node { left, right, .. }) => {
+                let depth = self.full_depth - self.stack.len();
+
+                // Go left
+                if (self.index >> (depth + self.packing_depth)) & 1 == 0 {
+                    self.stack.push(left);
+                    self.next()
+                }
+                // Go right
+                else {
+                    self.stack.push(right);
+                    self.next()
+                }
+            }
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.length.as_usize().saturating_sub(self.index);
+        (remaining, Some(remaining))
+    }
+}
+
+impl<T: Value> ExactSizeIterator for ArcIter<'_, T> {}
+#[derive(Debug)]
+pub struct ArcInterfaceIter<'a, T: Value, U: UpdateMap<T>> {
+    tree_iter: ArcIter<'a, T>,
+    updates: &'a U,
+    index: usize,
+    length: usize,
+}
+
+impl<'a, T: Value, U: UpdateMap<T>> ArcInterfaceIter<'a, T, U> {
+    pub fn new(root: &'a Tree<T>, depth: usize, length: Length, updates: &'a U) -> Self {
+        ArcInterfaceIter {
+            tree_iter: ArcIter::new(root, depth, length),
+            updates,
+            index: 0,
+            length: length.as_usize(),
+        }
+    }
+}
+
+impl<'a, T: Value, U: UpdateMap<T>> Iterator for ArcInterfaceIter<'a, T, U> {
+    type Item = Arc<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index >= self.length {
+            return None;
+        }
+        let idx = self.index;
+        self.index += 1;
+
+        let backing = self.tree_iter.next();
+        if let Some(new_val) = self.updates.get(idx) {
+            Some(
+                self.updates
+                    .get_arc(idx)
+                    .unwrap_or_else(|| Arc::new(new_val.clone())),
+            )
+        } else {
+            backing.cloned()
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let rem = self.length.saturating_sub(self.index);
+        (rem, Some(rem))
+    }
+}
+impl<T: Value, U: UpdateMap<T>> ExactSizeIterator for ArcInterfaceIter<'_, T, U> {}

src/lib.rs

@@ -7,6 +7,7 @@ pub mod error;
 pub mod interface;
 pub mod interface_iter;
 pub mod iter;
+pub mod iter_arc;
 pub mod leaf;
 pub mod level_iter;
 pub mod list;
@@ -23,6 +24,7 @@ pub mod vector;
 pub use cow::Cow;
 pub use error::Error;
 pub use interface::ImmList;
+pub use iter_arc::ArcIter;
 pub use leaf::Leaf;
 pub use list::List;
 pub use packed_leaf::PackedLeaf;

src/list.rs

@@ -2,6 +2,7 @@ use crate::builder::Builder;
 use crate::interface::{ImmList, Interface, MutList};
 use crate::interface_iter::{InterfaceIter, InterfaceIterCow};
 use crate::iter::Iter;
+use crate::iter_arc::{ArcInterfaceIter, ArcIter};
 use crate::level_iter::{LevelIter, LevelNode};
 use crate::serde::ListVisitor;
 use crate::tree::{IntraRebaseAction, RebaseAction};
@@ -16,7 +17,7 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer, ser::SerializeSeq}
 use ssz::{BYTES_PER_LENGTH_OFFSET, Decode, Encode, SszEncoder, TryFromIter};
 use std::collections::{BTreeMap, HashMap};
 use std::marker::PhantomData;
-use tree_hash::{Hash256, PackedEncoding, TreeHash};
+use tree_hash::{Hash256, PackedEncoding, TreeHash, TreeHashType};
 use typenum::Unsigned;
 use vec_map::VecMap;
 #[derive(Debug, Clone, Educe)]
@@ -136,6 +137,15 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
         Ok(self.interface.iter_from(index))
     }
 
+    pub fn iter_arc(&self) -> Result<impl Iterator<Item = Arc<T>>, Error> {
+        Ok(ArcInterfaceIter::new(
+            &self.interface.backing.tree,
+            self.interface.backing.depth,
+            Length(self.len()),
+            &self.interface.updates,
+        ))
+    }
+
     /// Iterate all internal nodes on the same level as `index`.
     pub fn level_iter_from(&self, index: usize) -> Result<LevelIter<'_, T>, Error> {
         // Return an empty iterator at index == length, just like slicing.
@@ -267,6 +277,10 @@ impl<T: Value, N: Unsigned> ImmList<T> for ListInner<T, N> {
     fn level_iter_from(&self, index: usize) -> LevelIter<'_, T> {
         LevelIter::from_index(index, &self.tree, self.depth, self.length)
     }
+
+    fn iter_arc(&self, index: usize) -> Result<ArcIter<'_, T>, Error> {
+        ArcIter::from_index(index, &self.tree, self.depth, self.length)
+    }
 }
 
 impl<T, N> MutList<T> for ListInner<T, N>
@@ -370,8 +384,8 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Default for List<T, N, U> {
 }
 
 impl<T: Value + Send + Sync, N: Unsigned, U: UpdateMap<T>> TreeHash for List<T, N, U> {
-    fn tree_hash_type() -> tree_hash::TreeHashType {
-        tree_hash::TreeHashType::List
+    fn tree_hash_type() -> TreeHashType {
+        TreeHashType::List
     }
 
     fn tree_hash_packed_encoding(&self) -> PackedEncoding {

src/tests/proptest/operations.rs

@@ -4,7 +4,7 @@ use proptest::prelude::*;
 use ssz::{Decode, Encode};
 use std::fmt::Debug;
 use std::marker::PhantomData;
-use tree_hash::{Hash256, TreeHash};
+use tree_hash::{Hash256, TreeHash, TreeHashType};
 use typenum::{U1, U2, U3, U4, U7, U8, U9, U32, U33, U1024, Unsigned};
 
 const OP_LIMIT: usize = 128;
@@ -107,6 +107,8 @@ pub enum Op<T> {
     Push(T),
     /// Check the `iter` method.
     Iter,
+    /// Cheeck the `iter_arc` method for non basic types.
+    IterArc,
     /// Check the `iter_from` method.
     IterFrom(usize),
     /// Check the `pop_front` method.
@@ -154,10 +156,11 @@ where
         Just(Op::Debase),
         Just(Op::FromIntoRoundtrip),
         Just(Op::IntraRebase),
+        Just(Op::IterArc),
     ];
     prop_oneof![
         10 => a_block,
-        6 => b_block
+        7 => b_block
     ]
 }
 
@@ -210,6 +213,16 @@ where
             Op::Iter => {
                 assert!(list.iter().eq(spec.iter()));
             }
+            Op::IterArc => {
+                if <T as TreeHash>::tree_hash_type() != TreeHashType::Basic {
+                    let actual: Vec<T> =
+                        list.iter_arc().unwrap().map(|arc| (*arc).clone()).collect();
+                    let expected: Vec<T> = spec.iter().cloned().collect();
+
+                    assert_eq!(actual, expected);
+                }
+            }
+
             Op::IterFrom(index) => match (list.iter_from(index), spec.iter_from(index)) {
                 (Ok(iter1), Ok(iter2)) => assert!(iter1.eq(iter2)),
                 (Err(e1), Err(e2)) => assert_eq!(e1, e2),
@@ -306,6 +319,16 @@ where
             Op::Iter => {
                 assert!(vect.iter().eq(spec.iter()));
             }
+            Op::IterArc => {
+                if <T as TreeHash>::tree_hash_type() != TreeHashType::Basic {
+                    let actual: Vec<T> =
+                        vect.iter_arc().unwrap().map(|arc| (*arc).clone()).collect();
+                    let expected: Vec<T> = spec.iter().cloned().collect();
+
+                    assert!(actual.eq(&expected));
+                }
+            }
+
             Op::IterFrom(index) => match (vect.iter_from(index), spec.iter_from(index)) {
                 (Ok(iter1), Ok(iter2)) => assert!(iter1.eq(iter2)),
                 (Err(e1), Err(e2)) => assert_eq!(e1, e2),