Apply updates recursively

src/interface.rs

@@ -8,7 +8,7 @@ use crate::{
 use arbitrary::Arbitrary;
 use std::collections::BTreeMap;
 use std::marker::PhantomData;
-use tree_hash::Hash256;
+use tree_hash::{Hash256, TreeHashType};
 
 pub trait ImmList<T: Value> {
     fn get(&self, idx: usize) -> Option<&T>;
@@ -80,8 +80,22 @@ where
         Ok(())
     }
 
+    pub fn apply_recursive_updates(&mut self) -> Result<(), Error> {
+        let is_recursive = match T::tree_hash_type() {
+            TreeHashType::Basic => false,
+            TreeHashType::Container | TreeHashType::List | TreeHashType::Vector => true,
+        };
+
+        if is_recursive {
+            self.updates.for_each_mut(|item| item.apply())
+        } else {
+            Ok(())
+        }
+    }
+
     pub fn apply_updates(&mut self) -> Result<(), Error> {
         if !self.updates.is_empty() {
+            self.apply_recursive_updates()?;
             let updates = std::mem::take(&mut self.updates);
             self.backing.update(updates, None)
         } else {

src/lib.rs

@@ -32,6 +32,20 @@ pub use vector::Vector;
 use ssz::{Decode, Encode};
 use tree_hash::TreeHash;
 
-pub trait Value: Encode + Decode + TreeHash + PartialEq + Clone {}
+pub trait PendingUpdates {
+    fn apply(&mut self) -> Result<(), Error> {
+        Ok(())
+    }
+}
 
-impl<T> Value for T where T: Encode + Decode + TreeHash + PartialEq + Clone {}
+pub trait Value: Encode + Decode + TreeHash + PartialEq + Clone + PendingUpdates {}
+
+impl<T> Value for T where T: Encode + Decode + TreeHash + PartialEq + Clone + PendingUpdates {}
+
+// Default impls for known types
+impl PendingUpdates for u8 {}
+impl PendingUpdates for u16 {}
+impl PendingUpdates for u32 {}
+impl PendingUpdates for u64 {}
+impl PendingUpdates for u128 {}
+impl PendingUpdates for tree_hash::Hash256 {}

src/list.rs

@@ -6,7 +6,7 @@ use crate::serde::ListVisitor;
 use crate::tree::RebaseAction;
 use crate::update_map::MaxMap;
 use crate::utils::{arb_arc, int_log, opt_packing_depth, updated_length, Length};
-use crate::{Arc, Cow, Error, Tree, UpdateMap, Value};
+use crate::{Arc, Cow, Error, PendingUpdates, Tree, UpdateMap, Value};
 use arbitrary::Arbitrary;
 use derivative::Derivative;
 use itertools::process_results;
@@ -39,6 +39,12 @@ pub struct ListInner<T: Value, N: Unsigned> {
     _phantom: PhantomData<N>,
 }
 
+impl<T: Value, N: Unsigned, U: UpdateMap<T>> PendingUpdates for List<T, N, U> {
+    fn apply(&mut self) -> Result<(), Error> {
+        self.interface.apply_updates()
+    }
+}
+
 impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
     pub fn new(vec: Vec<T>) -> Result<Self, Error> {
         Self::try_from_iter(vec)

src/tests/mod.rs

@@ -4,5 +4,6 @@ mod builder;
 mod iterator;
 mod packed;
 mod proptest;
+mod recursion;
 mod repeat;
 mod size_of;

src/tests/proptest/mod.rs

@@ -1,4 +1,4 @@
-use crate::List;
+use crate::{List, PendingUpdates};
 use proptest::prelude::*;
 use ssz_derive::{Decode, Encode};
 use tree_hash::Hash256;
@@ -44,6 +44,14 @@ pub struct Large {
     d: List<u64, U4>,
 }
 
+impl PendingUpdates for Large {
+    fn apply(&mut self) -> Result<(), crate::Error> {
+        // TODO use macro derive
+        self.d.apply()?;
+        Ok(())
+    }
+}
+
 pub fn arb_large() -> impl Strategy<Value = Large> {
     (
         any::<u64>(),

src/tests/recursion.rs

@@ -0,0 +1,66 @@
+use crate::{List, PendingUpdates};
+use ssz_derive::{Decode, Encode};
+use tree_hash::TreeHash;
+use tree_hash_derive::TreeHash;
+use typenum::U16;
+
+#[test]
+fn recursive_list_list() {
+    let mut l = List::<List<u64, U16>, U16>::default();
+
+    l.push(<_>::default()).unwrap();
+    l.get_mut(0).unwrap().push(1).unwrap();
+    l.apply_updates().unwrap();
+    // assert it does not throw
+    let h_1 = l.tree_hash_root();
+
+    assert!(!l.has_pending_updates());
+    assert_eq!(*l.get(0).unwrap().get(0).unwrap(), 1);
+
+    // Replace value
+    *l.get_mut(0).unwrap().get_mut(0).unwrap() = 2;
+    assert_eq!(*l.get(0).unwrap().get(0).unwrap(), 2);
+
+    // Commit only top list
+    l.apply_updates().unwrap();
+    // Root should be different
+    assert_ne!(h_1, l.tree_hash_root());
+}
+
+/// Struct with multiple fields shared by multiple proptests.
+#[derive(Default, Debug, Clone, PartialEq, Encode, Decode, TreeHash)]
+pub struct ListContainer {
+    a: u8,
+    b: List<u64, U16>,
+}
+
+impl PendingUpdates for ListContainer {
+    fn apply(&mut self) -> Result<(), crate::Error> {
+        // TODO use macro derive
+        self.b.apply()?;
+        Ok(())
+    }
+}
+
+#[test]
+fn recursive_list_container_list() {
+    let mut l = List::<ListContainer, U16>::default();
+
+    l.push(<_>::default()).unwrap();
+    l.get_mut(0).unwrap().b.push(1).unwrap();
+    l.apply_updates().unwrap();
+    // assert it does not throw
+    let h_1 = l.tree_hash_root();
+
+    assert!(!l.has_pending_updates());
+    assert_eq!(*l.get(0).unwrap().b.get(0).unwrap(), 1);
+
+    // Replace value
+    *l.get_mut(0).unwrap().b.get_mut(0).unwrap() = 2;
+    assert_eq!(*l.get(0).unwrap().b.get(0).unwrap(), 2);
+
+    // Commit only top list
+    l.apply_updates().unwrap();
+    // Root should be different
+    assert_ne!(h_1, l.tree_hash_root());
+}

src/update_map.rs

@@ -25,6 +25,8 @@ pub trait UpdateMap<T>: Default + Clone {
     where
         F: FnMut(usize, &T) -> ControlFlow<(), Result<(), E>>;
 
+    fn for_each_mut<E, F: FnMut(&mut T) -> Result<(), E>>(&mut self, f: F) -> Result<(), E>;
+
     fn max_index(&self) -> Option<usize>;
 
     fn len(&self) -> usize;
@@ -90,6 +92,13 @@ impl<T: Clone> UpdateMap<T> for BTreeMap<usize, T> {
         Ok(())
     }
 
+    fn for_each_mut<E, F: FnMut(&mut T) -> Result<(), E>>(&mut self, mut f: F) -> Result<(), E> {
+        for (_, item) in self.iter_mut() {
+            f(item)?;
+        }
+        Ok(())
+    }
+
     fn max_index(&self) -> Option<usize> {
         max_btree_index(self)
     }
@@ -159,6 +168,13 @@ impl<T: Clone> UpdateMap<T> for VecMap<T> {
         Ok(())
     }
 
+    fn for_each_mut<E, F: FnMut(&mut T) -> Result<(), E>>(&mut self, mut f: F) -> Result<(), E> {
+        for (_, item) in self.iter_mut() {
+            f(item)?;
+        }
+        Ok(())
+    }
+
     fn max_index(&self) -> Option<usize> {
         // FIXME(sproul): this is slow, make a wrapper type that tracks the max index
         self.keys().next_back()
@@ -214,6 +230,10 @@ where
         self.inner.for_each_range(start, end, f)
     }
 
+    fn for_each_mut<E, F: FnMut(&mut T) -> Result<(), E>>(&mut self, f: F) -> Result<(), E> {
+        self.inner.for_each_mut(f)
+    }
+
     fn len(&self) -> usize {
         self.inner.len()
     }

src/vector.rs

@@ -4,7 +4,7 @@ use crate::iter::Iter;
 use crate::tree::RebaseAction;
 use crate::update_map::MaxMap;
 use crate::utils::{arb_arc, Length};
-use crate::{Arc, Cow, Error, List, Tree, UpdateMap, Value};
+use crate::{Arc, Cow, Error, List, PendingUpdates, Tree, UpdateMap, Value};
 use arbitrary::Arbitrary;
 use derivative::Derivative;
 use serde::{Deserialize, Serialize};
@@ -40,6 +40,12 @@ pub struct VectorInner<T: Value, N: Unsigned> {
     _phantom: PhantomData<N>,
 }
 
+impl<T: Value, N: Unsigned, U: UpdateMap<T>> PendingUpdates for Vector<T, N, U> {
+    fn apply(&mut self) -> Result<(), Error> {
+        self.interface.apply_updates()
+    }
+}
+
 impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
     pub fn new(vec: Vec<T>) -> Result<Self, Error> {
         if vec.len() == N::to_usize() {