From 570fbd49b3c9218fc66a32a7931e922e87091e60 Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Fri, 12 Jul 2024 14:59:05 +1000
Subject: [PATCH 1/2] Interior mutability WIP

---
 Cargo.toml                       |  3 +-
 src/interface.rs                 | 69 +++++++++++++++++++--------
 src/interface_iter.rs            | 13 +++--
 src/lib.rs                       |  5 +-
 src/list.rs                      | 82 ++++++++++++++++++++++----------
 src/tests/packed.rs              |  8 ++--
 src/tests/proptest/operations.rs |  4 +-
 src/utils.rs                     | 21 ++++++++
 src/value_ref.rs                 | 22 +++++++++
 src/vector.rs                    | 81 ++++++++++++++++++++++---------
 10 files changed, 225 insertions(+), 83 deletions(-)
 create mode 100644 src/value_ref.rs

diff --git a/Cargo.toml b/Cargo.toml
index 3692bc2..b73909e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,12 +11,13 @@ keywords = ["ethereum", "functional"]
 categories = ["data-structures", "cryptography::cryptocurrencies", ]
 
 [dependencies]
+arc-swap = "1.7.1"
 derivative = "2.2.0"
 ethereum_hashing = "0.6.0"
 ethereum_ssz = "0.5.0"
 ethereum_ssz_derive = "0.5.0"
 itertools = "0.10.3"
-parking_lot = "0.12.1"
+parking_lot = { version = "0.12.1", features = ["arc_lock"] }
 rayon = "1.5.1"
 serde = { version = "1.0.0", features = ["derive"] }
 tree_hash = "0.6.0"
diff --git a/src/interface.rs b/src/interface.rs
index a4e2844..64efae8 100644
--- a/src/interface.rs
+++ b/src/interface.rs
@@ -1,12 +1,14 @@
 use crate::level_iter::LevelIter;
 use crate::update_map::UpdateMap;
-use crate::utils::{updated_length, Length};
+use crate::utils::{arb_rwlock, partial_eq_rwlock, updated_length, Length};
 use crate::{
     interface_iter::{InterfaceIter, InterfaceIterCow},
     iter::Iter,
-    Cow, Error, Value,
+    Cow, Error, Value, ValueRef,
 };
 use arbitrary::Arbitrary;
+use derivative::Derivative;
+use parking_lot::{RwLock, RwLockReadGuard};
 use std::collections::BTreeMap;
 use std::marker::PhantomData;
 use tree_hash::Hash256;
@@ -29,13 +31,14 @@ pub trait MutList<T: Value>: ImmList<T> {
     fn validate_push(current_len: usize) -> Result<(), Error>;
     fn replace(&mut self, index: usize, value: T) -> Result<(), Error>;
     fn update<U: UpdateMap<T>>(
-        &mut self,
+        &self,
         updates: U,
         hash_updates: Option<BTreeMap<(usize, usize), Hash256>>,
     ) -> Result<(), Error>;
 }
 
-#[derive(Debug, PartialEq, Clone, Arbitrary)]
+#[derive(Debug, Derivative, Arbitrary)]
+#[derivative(PartialEq)]
 pub struct Interface<T, B, U>
 where
     T: Value,
@@ -43,10 +46,27 @@ where
     U: UpdateMap<T>,
 {
     pub(crate) backing: B,
-    pub(crate) updates: U,
+    #[derivative(PartialEq(compare_with = "partial_eq_rwlock"))]
+    #[arbitrary(with = arb_rwlock)]
+    pub(crate) updates: RwLock<U>,
     pub(crate) _phantom: PhantomData<T>,
 }
 
+impl<T, B, U> Clone for Interface<T, B, U>
+where
+    T: Value,
+    B: MutList<T> + Clone,
+    U: UpdateMap<T>,
+{
+    fn clone(&self) -> Self {
+        Self {
+            backing: self.backing.clone(),
+            updates: RwLock::new(self.updates.read().clone()),
+            _phantom: PhantomData,
+        }
+    }
+}
+
 impl<T, B, U> Interface<T, B, U>
 where
     T: Value,
@@ -56,54 +76,60 @@ where
     pub fn new(backing: B) -> Self {
         Self {
             backing,
-            updates: U::default(),
+            updates: RwLock::new(U::default()),
             _phantom: PhantomData,
         }
     }
 
-    pub fn get(&self, idx: usize) -> Option<&T> {
-        self.updates.get(idx).or_else(|| self.backing.get(idx))
+    pub fn get(&self, idx: usize) -> Option<ValueRef<'_, T>> {
+        RwLockReadGuard::try_map(self.updates.read(), |updates| updates.get(idx))
+            .ok()
+            .map(ValueRef::Pending)
+            .or_else(|| self.backing.get(idx).map(ValueRef::Applied))
     }
 
     pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
         self.updates
+            .get_mut()
             .get_mut_with(idx, |idx| self.backing.get(idx).cloned())
     }
 
     pub fn get_cow(&mut self, index: usize) -> Option<Cow<T>> {
         self.updates
+            .get_mut()
             .get_cow_with(index, |idx| self.backing.get(idx))
     }
 
     pub fn push(&mut self, value: T) -> Result<(), Error> {
         let index = self.len();
         B::validate_push(index)?;
-        self.updates.insert(index, value);
+        self.updates.get_mut().insert(index, value);
 
         Ok(())
     }
 
-    pub fn apply_updates(&mut self) -> Result<(), Error> {
-        if !self.updates.is_empty() {
-            let updates = std::mem::take(&mut self.updates);
-            self.backing.update(updates, None)
+    pub fn apply_updates(&self) -> Result<(), Error> {
+        let mut updates = self.updates.write();
+        if !updates.is_empty() {
+            self.backing.update(std::mem::take(&mut *updates), None)?;
+            drop(updates);
+            Ok(())
         } else {
             Ok(())
         }
     }
 
     pub fn has_pending_updates(&self) -> bool {
-        !self.updates.is_empty()
+        !self.updates.read().is_empty()
     }
 
-    pub fn iter(&self) -> InterfaceIter<T, U> {
+    pub fn iter(&self) -> InterfaceIter<T> {
         self.iter_from(0)
     }
 
-    pub fn iter_from(&self, index: usize) -> InterfaceIter<T, U> {
+    pub fn iter_from(&self, index: usize) -> InterfaceIter<T> {
         InterfaceIter {
             tree_iter: self.backing.iter_from(index),
-            updates: &self.updates,
             index,
             length: self.len(),
         }
@@ -113,7 +139,7 @@ where
         let index = 0;
         InterfaceIterCow {
             tree_iter: self.backing.iter_from(index),
-            updates: &mut self.updates,
+            updates: self.updates.get_mut(),
             index,
         }
     }
@@ -127,7 +153,7 @@ where
     }
 
     pub fn len(&self) -> usize {
-        updated_length(self.backing.len(), &self.updates).as_usize()
+        updated_length(self.backing.len(), &*self.updates.read()).as_usize()
     }
 
     pub fn is_empty(&self) -> bool {
@@ -135,10 +161,11 @@ where
     }
 
     pub fn bulk_update(&mut self, updates: U) -> Result<(), Error> {
-        if !self.updates.is_empty() {
+        let self_updates = self.updates.get_mut();
+        if !self_updates.is_empty() {
             return Err(Error::BulkUpdateUnclean);
         }
-        self.updates = updates;
+        *self_updates = updates;
         Ok(())
     }
 }
diff --git a/src/interface_iter.rs b/src/interface_iter.rs
index aff6a6b..244bd2e 100644
--- a/src/interface_iter.rs
+++ b/src/interface_iter.rs
@@ -1,15 +1,17 @@
 use crate::iter::Iter;
 use crate::{Cow, UpdateMap, Value};
+use parking_lot::RwLockWriteGuard;
 
 #[derive(Debug)]
-pub struct InterfaceIter<'a, T: Value, U: UpdateMap<T>> {
+pub struct InterfaceIter<'a, T: Value> {
     pub(crate) tree_iter: Iter<'a, T>,
-    pub(crate) updates: &'a U,
+    // FIXME(sproul): remove write guard and flush updates prior to iteration?
+    // pub(crate) updates: RwLockWriteGuard<'a, U>,
     pub(crate) index: usize,
     pub(crate) length: usize,
 }
 
-impl<'a, T: Value, U: UpdateMap<T>> Iterator for InterfaceIter<'a, T, U> {
+impl<'a, T: Value> Iterator for InterfaceIter<'a, T> {
     type Item = &'a T;
 
     fn next(&mut self) -> Option<&'a T> {
@@ -20,7 +22,8 @@ impl<'a, T: Value, U: UpdateMap<T>> Iterator for InterfaceIter<'a, T, U> {
         let backing_value = self.tree_iter.next();
 
         // Prioritise the value from the update map.
-        self.updates.get(index).or(backing_value)
+        // self.updates.get_mut().get(index).or(backing_value)
+        backing_value
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -29,7 +32,7 @@ impl<'a, T: Value, U: UpdateMap<T>> Iterator for InterfaceIter<'a, T, U> {
     }
 }
 
-impl<'a, T: Value, U: UpdateMap<T>> ExactSizeIterator for InterfaceIter<'a, T, U> {}
+impl<'a, T: Value> ExactSizeIterator for InterfaceIter<'a, T> {}
 
 #[derive(Debug)]
 pub struct InterfaceIterCow<'a, T: Value, U: UpdateMap<T>> {
diff --git a/src/lib.rs b/src/lib.rs
index f093028..8ea8d79 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -17,6 +17,7 @@ mod tests;
 pub mod tree;
 pub mod update_map;
 pub mod utils;
+pub mod value_ref;
 pub mod vector;
 
 pub use cow::Cow;
@@ -26,8 +27,10 @@ pub use leaf::Leaf;
 pub use list::List;
 pub use packed_leaf::PackedLeaf;
 pub use tree::Tree;
-pub use triomphe::Arc;
+// pub use triomphe::Arc;
+pub use std::sync::Arc; // FIXME(sproul)
 pub use update_map::UpdateMap;
+pub use value_ref::ValueRef;
 pub use vector::Vector;
 
 use ssz::{Decode, Encode};
diff --git a/src/list.rs b/src/list.rs
index f3d8676..b3c26f4 100644
--- a/src/list.rs
+++ b/src/list.rs
@@ -6,15 +6,20 @@ use crate::level_iter::{LevelIter, LevelNode};
 use crate::serde::ListVisitor;
 use crate::tree::RebaseAction;
 use crate::update_map::MaxMap;
-use crate::utils::{arb_arc, compute_level, int_log, opt_packing_depth, updated_length, Length};
-use crate::{Arc, Cow, Error, Tree, UpdateMap, Value};
+use crate::utils::{
+    arb_arc, arb_arc_swap, compute_level, int_log, opt_packing_depth, partial_eq_arc_swap,
+    updated_length, Length,
+};
+use crate::{Arc, Cow, Error, Tree, UpdateMap, Value, ValueRef};
 use arbitrary::Arbitrary;
+use arc_swap::ArcSwap;
 use derivative::Derivative;
 use itertools::process_results;
 use serde::{ser::SerializeSeq, Deserialize, Deserializer, Serialize, Serializer};
 use ssz::{Decode, Encode, SszEncoder, TryFromIter, BYTES_PER_LENGTH_OFFSET};
 use std::collections::BTreeMap;
 use std::marker::PhantomData;
+use std::sync::Arc as StdArc;
 use tree_hash::{Hash256, PackedEncoding, TreeHash};
 use typenum::Unsigned;
 use vec_map::VecMap;
@@ -27,12 +32,13 @@ pub struct List<T: Value, N: Unsigned, U: UpdateMap<T> = MaxMap<VecMap<T>>> {
     pub(crate) interface: Interface<T, ListInner<T, N>, U>,
 }
 
-#[derive(Debug, Clone, Derivative, Arbitrary)]
+#[derive(Debug, Derivative, Arbitrary)]
 #[derivative(PartialEq(bound = "T: Value, N: Unsigned"))]
 #[arbitrary(bound = "T: Arbitrary<'arbitrary> + Value, N: Unsigned")]
 pub struct ListInner<T: Value, N: Unsigned> {
-    #[arbitrary(with = arb_arc)]
-    pub(crate) tree: Arc<Tree<T>>,
+    #[derivative(PartialEq(compare_with = "partial_eq_arc_swap"))]
+    #[arbitrary(with = arb_arc_swap)]
+    pub(crate) tree: ArcSwap<Tree<T>>,
     pub(crate) length: Length,
     pub(crate) depth: usize,
     pub(crate) packing_depth: usize,
@@ -40,6 +46,18 @@ pub struct ListInner<T: Value, N: Unsigned> {
     _phantom: PhantomData<N>,
 }
 
+impl<T: Value, N: Unsigned> Clone for ListInner<T, N> {
+    fn clone(&self) -> Self {
+        Self {
+            tree: ArcSwap::new(self.tree.load_full()),
+            length: self.length,
+            depth: self.depth,
+            packing_depth: self.packing_depth,
+            _phantom: PhantomData,
+        }
+    }
+}
+
 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)
@@ -49,7 +67,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
         let packing_depth = opt_packing_depth::<T>().unwrap_or(0);
         Self {
             interface: Interface::new(ListInner {
-                tree,
+                tree: ArcSwap::new(tree),
                 length,
                 depth,
                 packing_depth,
@@ -113,11 +131,11 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
         self.iter().cloned().collect()
     }
 
-    pub fn iter(&self) -> InterfaceIter<T, U> {
+    pub fn iter(&self) -> InterfaceIter<T> {
         self.interface.iter()
     }
 
-    pub fn iter_from(&self, index: usize) -> Result<InterfaceIter<T, U>, Error> {
+    pub fn iter_from(&self, index: usize) -> Result<InterfaceIter<T>, Error> {
         // Return an empty iterator at index == length, just like slicing.
         if index > self.len() {
             return Err(Error::OutOfBoundsIterFrom {
@@ -145,7 +163,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
     }
 
     // Wrap trait methods so we present a Vec-like interface without having to import anything.
-    pub fn get(&self, index: usize) -> Option<&T> {
+    pub fn get(&self, index: usize) -> Option<ValueRef<T>> {
         self.interface.get(index)
     }
 
@@ -177,6 +195,10 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
         self.interface.apply_updates()
     }
 
+    pub fn apply_updates_immutable(&mut self) -> Result<(), Error> {
+        self.interface.apply_updates()
+    }
+
     pub fn bulk_update(&mut self, updates: U) -> Result<(), Error> {
         self.interface.bulk_update(updates)
     }
@@ -193,7 +215,8 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
     ///
     /// Errors if `n > self.len()`.
     pub fn pop_front_slow(&mut self, n: usize) -> Result<(), Error> {
-        *self = Self::try_from_iter(self.iter_from(n)?.cloned())?;
+        let updated = Self::try_from_iter(self.iter_from(n)?.cloned())?;
+        *self = updated;
         Ok(())
     }
 
@@ -240,9 +263,11 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
 
 impl<T: Value, N: Unsigned> ImmList<T> for ListInner<T, N> {
     fn get(&self, index: usize) -> Option<&T> {
+        // FIXME(sproul): this is fucked unfortunately
         if index < self.len().as_usize() {
-            self.tree
-                .get_recursive(index, self.depth, self.packing_depth)
+            ArcSwap::map(&self.tree, |tree: &Arc<Tree<T>>| {
+                tree.get_recursive(index, self.depth, self.packing_depth)
+            })
         } else {
             None
         }
@@ -253,11 +278,11 @@ impl<T: Value, N: Unsigned> ImmList<T> for ListInner<T, N> {
     }
 
     fn iter_from(&self, index: usize) -> Iter<T> {
-        Iter::from_index(index, &self.tree, self.depth, self.length)
+        Iter::from_index(index, &self.tree.load_full(), self.depth, self.length)
     }
 
     fn level_iter_from(&self, index: usize) -> LevelIter<T> {
-        LevelIter::from_index(index, &self.tree, self.depth, self.length)
+        LevelIter::from_index(index, &self.tree.load_full(), self.depth, self.length)
     }
 }
 
@@ -282,7 +307,11 @@ where
             });
         }
 
-        self.tree = self.tree.with_updated_leaf(index, value, self.depth)?;
+        let new_tree = self
+            .tree
+            .load()
+            .with_updated_leaf(index, value, self.depth)?;
+        self.tree.store(new_tree);
         if index == self.length.as_usize() {
             *self.length.as_mut() += 1;
         }
@@ -290,7 +319,7 @@ where
     }
 
     fn update<U: UpdateMap<T>>(
-        &mut self,
+        &self,
         updates: U,
         hash_updates: Option<BTreeMap<(usize, usize), Hash256>>,
     ) -> Result<(), Error> {
@@ -303,9 +332,11 @@ where
             return Ok(());
         }
         self.length = updated_length(self.length, &updates);
-        self.tree =
+        let updated =
             self.tree
+                .load()
                 .with_updated_leaves(&updates, 0, self.depth, hash_updates.as_ref())?;
+        self.tree.store(updated);
         Ok(())
     }
 }
@@ -319,16 +350,16 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
 
     pub fn rebase_on(&mut self, base: &Self) -> Result<(), Error> {
         match Tree::rebase_on(
-            &self.interface.backing.tree,
-            &base.interface.backing.tree,
+            &self.interface.backing.tree.load_full(),
+            &base.interface.backing.tree.load_full(),
             Some((self.interface.backing.length, base.interface.backing.length)),
             self.interface.backing.depth + self.interface.backing.packing_depth,
         )? {
             RebaseAction::EqualReplace(replacement) => {
-                self.interface.backing.tree = replacement.clone();
+                self.interface.backing.tree.store(replacement.clone());
             }
             RebaseAction::NotEqualReplace(replacement) => {
-                self.interface.backing.tree = replacement;
+                self.interface.backing.tree.store(replacement);
             }
             _ => (),
         }
@@ -356,17 +387,16 @@ impl<T: Value + Send + Sync, N: Unsigned> TreeHash for List<T, N> {
     }
 
     fn tree_hash_root(&self) -> Hash256 {
-        // FIXME(sproul): remove assert
-        assert!(!self.interface.has_pending_updates());
-
-        let root = self.interface.backing.tree.tree_hash();
+        self.apply_updates_immutable()
+            .expect("updates should apply");
+        let root = self.interface.backing.tree.load().tree_hash();
         tree_hash::mix_in_length(&root, self.len())
     }
 }
 
 impl<'a, T: Value, N: Unsigned, U: UpdateMap<T>> IntoIterator for &'a List<T, N, U> {
     type Item = &'a T;
-    type IntoIter = InterfaceIter<'a, T, U>;
+    type IntoIter = InterfaceIter<'a, T>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.iter()
diff --git a/src/tests/packed.rs b/src/tests/packed.rs
index 5351b38..3ee0848 100644
--- a/src/tests/packed.rs
+++ b/src/tests/packed.rs
@@ -13,7 +13,7 @@ fn u64_packed_list_build_and_iter() {
         assert_eq!(vec, from_iter);
 
         for i in 0..len as usize {
-            assert_eq!(list.get(i), vec.get(i));
+            assert_eq!(list.get(i).as_deref(), vec.get(i));
         }
     }
 }
@@ -40,7 +40,7 @@ fn u64_packed_vector_build_and_iter() {
     assert_eq!(vec, from_iter);
 
     for i in 0..len as usize {
-        assert_eq!(vector.get(i), vec.get(i));
+        assert_eq!(vector.get(i).as_deref(), vec.get(i));
     }
 }
 
@@ -74,11 +74,11 @@ fn out_of_order_mutations() {
     for (i, v) in mutations {
         *list.get_mut(i).unwrap() = v;
         vec[i] = v;
-        assert_eq!(list.get(i), Some(&v));
+        assert_eq!(list.get(i).as_deref(), Some(&v));
 
         list.apply_updates().unwrap();
 
-        assert_eq!(list.get(i), Some(&v));
+        assert_eq!(list.get(i).as_deref(), Some(&v));
     }
 
     assert_eq!(list.to_vec(), vec);
diff --git a/src/tests/proptest/operations.rs b/src/tests/proptest/operations.rs
index f24504f..6f60e84 100644
--- a/src/tests/proptest/operations.rs
+++ b/src/tests/proptest/operations.rs
@@ -183,7 +183,7 @@ where
                 assert_eq!(list.len(), spec.len())
             }
             Op::Get(index) => {
-                assert_eq!(list.get(index), spec.get(index));
+                assert_eq!(list.get(index).as_deref(), spec.get(index));
             }
             Op::Set(index, value) => {
                 let res = list.get_mut(index).map(|elem| *elem = value.clone());
@@ -269,7 +269,7 @@ where
                 assert_eq!(vect.len(), spec.len())
             }
             Op::Get(index) => {
-                assert_eq!(vect.get(index), spec.get(index));
+                assert_eq!(vect.get(index).as_deref(), spec.get(index));
             }
             Op::Set(index, value) => {
                 let res = vect.get_mut(index).map(|elem| *elem = value.clone());
diff --git a/src/utils.rs b/src/utils.rs
index 765385a..841340d 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,5 +1,6 @@
 use crate::{Arc, UpdateMap};
 use arbitrary::Arbitrary;
+use arc_swap::ArcSwap;
 use parking_lot::RwLock;
 use std::collections::BTreeMap;
 use tree_hash::{Hash256, TreeHash, TreeHashType};
@@ -110,6 +111,26 @@ pub fn arb_rwlock<'a, T: Arbitrary<'a>>(
     T::arbitrary(u).map(RwLock::new)
 }
 
+pub fn arb_arc_rwlock<'a, T: Arbitrary<'a>>(
+    u: &mut arbitrary::Unstructured<'a>,
+) -> arbitrary::Result<std::sync::Arc<RwLock<T>>> {
+    T::arbitrary(u).map(RwLock::new).map(std::sync::Arc::new)
+}
+
+pub fn arb_arc_swap<'a, T: Arbitrary<'a>>(
+    u: &mut arbitrary::Unstructured<'a>,
+) -> arbitrary::Result<ArcSwap<T>> {
+    T::arbitrary(u).map(std::sync::Arc::new).map(ArcSwap::new)
+}
+
+pub fn partial_eq_rwlock<'a, T: PartialEq>(x: &RwLock<T>, y: &RwLock<T>) -> bool {
+    *x.read() == *y.read()
+}
+
+pub fn partial_eq_arc_swap<'a, T: PartialEq>(x: &ArcSwap<T>, y: &ArcSwap<T>) -> bool {
+    *x.load() == *y.load()
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
diff --git a/src/value_ref.rs b/src/value_ref.rs
new file mode 100644
index 0000000..12bcd61
--- /dev/null
+++ b/src/value_ref.rs
@@ -0,0 +1,22 @@
+use parking_lot::MappedRwLockReadGuard;
+use std::ops::Deref;
+
+/// Reference to a value within a List or Vector.
+#[derive(Debug)]
+pub enum ValueRef<'a, T> {
+    /// The value is present in the `updates` map and is waiting to be applied.
+    Pending(MappedRwLockReadGuard<'a, T>),
+    /// The value is present in the tree.
+    Applied(&'a T),
+}
+
+impl<'a, T> Deref for ValueRef<'a, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            Self::Pending(guard) => guard.deref(),
+            Self::Applied(reference) => reference,
+        }
+    }
+}
diff --git a/src/vector.rs b/src/vector.rs
index f5b6711..a5bd86e 100644
--- a/src/vector.rs
+++ b/src/vector.rs
@@ -4,9 +4,10 @@ use crate::iter::Iter;
 use crate::level_iter::LevelIter;
 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::utils::{arb_arc_swap, partial_eq_arc_swap, Length};
+use crate::{Arc, Cow, Error, List, Tree, UpdateMap, Value, ValueRef};
 use arbitrary::Arbitrary;
+use arc_swap::ArcSwap;
 use derivative::Derivative;
 use serde::{Deserialize, Serialize};
 use ssz::{Decode, Encode, SszEncoder, TryFromIter, BYTES_PER_LENGTH_OFFSET};
@@ -29,18 +30,30 @@ pub struct Vector<T: Value, N: Unsigned, U: UpdateMap<T> = MaxMap<VecMap<T>>> {
     pub(crate) interface: Interface<T, VectorInner<T, N>, U>,
 }
 
-#[derive(Debug, Derivative, Clone, Arbitrary)]
+#[derive(Debug, Derivative, Arbitrary)]
 #[derivative(PartialEq(bound = "T: Value, N: Unsigned"))]
 #[arbitrary(bound = "T: Arbitrary<'arbitrary> + Value, N: Unsigned")]
 pub struct VectorInner<T: Value, N: Unsigned> {
-    #[arbitrary(with = arb_arc)]
-    pub(crate) tree: Arc<Tree<T>>,
+    #[arbitrary(with = arb_arc_swap)]
+    #[derivative(PartialEq(compare_with = "partial_eq_arc_swap"))]
+    pub(crate) tree: ArcSwap<Tree<T>>,
     pub(crate) depth: usize,
     packing_depth: usize,
     #[arbitrary(default)]
     _phantom: PhantomData<N>,
 }
 
+impl<T: Value, N: Unsigned> Clone for VectorInner<T, N> {
+    fn clone(&self) -> Self {
+        Self {
+            tree: ArcSwap::new(self.tree.load_full()),
+            depth: self.depth,
+            packing_depth: self.packing_depth,
+            _phantom: PhantomData,
+        }
+    }
+}
+
 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() {
@@ -65,11 +78,11 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
         self.iter().cloned().collect()
     }
 
-    pub fn iter(&self) -> InterfaceIter<T, U> {
+    pub fn iter(&self) -> InterfaceIter<T> {
         self.interface.iter()
     }
 
-    pub fn iter_from(&self, index: usize) -> Result<InterfaceIter<T, U>, Error> {
+    pub fn iter_from(&self, index: usize) -> Result<InterfaceIter<T>, Error> {
         if index > self.len() {
             return Err(Error::OutOfBoundsIterFrom {
                 index,
@@ -80,7 +93,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
     }
 
     // Wrap trait methods so we present a Vec-like interface without having to import anything.
-    pub fn get(&self, index: usize) -> Option<&T> {
+    pub fn get(&self, index: usize) -> Option<ValueRef<T>> {
         self.interface.get(index)
     }
 
@@ -107,6 +120,10 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
     pub fn apply_updates(&mut self) -> Result<(), Error> {
         self.interface.apply_updates()
     }
+
+    pub fn apply_updates_immutable(&self) -> Result<(), Error> {
+        self.interface.apply_updates()
+    }
 }
 
 impl<T: Value, N: Unsigned, U: UpdateMap<T>> TryFrom<List<T, N, U>> for Vector<T, N, U> {
@@ -116,7 +133,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> TryFrom<List<T, N, U>> for Vector<T
         if list.len() == N::to_usize() {
             let updates = list.interface.updates;
             let backing = VectorInner {
-                tree: list.interface.backing.tree,
+                tree: ArcSwap::new(list.interface.backing.tree.load_full()),
                 depth: list.interface.backing.depth,
                 packing_depth: list.interface.backing.packing_depth,
                 _phantom: PhantomData,
@@ -146,16 +163,16 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
 
     pub fn rebase_on(&mut self, base: &Self) -> Result<(), Error> {
         match Tree::rebase_on(
-            &self.interface.backing.tree,
-            &base.interface.backing.tree,
+            &self.interface.backing.tree.load_full(),
+            &base.interface.backing.tree.load_full(),
             None,
             self.interface.backing.depth + self.interface.backing.packing_depth,
         )? {
             RebaseAction::EqualReplace(replacement) => {
-                self.interface.backing.tree = replacement.clone();
+                self.interface.backing.tree.store(replacement.clone());
             }
             RebaseAction::NotEqualReplace(replacement) => {
-                self.interface.backing.tree = replacement;
+                self.interface.backing.tree.store(replacement);
             }
             _ => (),
         }
@@ -166,7 +183,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
 impl<T: Value, N: Unsigned, U: UpdateMap<T>> From<Vector<T, N, U>> for List<T, N, U> {
     fn from(vector: Vector<T, N, U>) -> Self {
         let mut list = List::from_parts(
-            vector.interface.backing.tree,
+            vector.interface.backing.tree.load_full(),
             vector.interface.backing.depth,
             Length(N::to_usize()),
         );
@@ -179,6 +196,7 @@ impl<T: Value, N: Unsigned> ImmList<T> for VectorInner<T, N> {
     fn get(&self, index: usize) -> Option<&T> {
         if index < self.len().as_usize() {
             self.tree
+                .load()
                 .get_recursive(index, self.depth, self.packing_depth)
         } else {
             None
@@ -190,11 +208,21 @@ impl<T: Value, N: Unsigned> ImmList<T> for VectorInner<T, N> {
     }
 
     fn iter_from(&self, index: usize) -> Iter<T> {
-        Iter::from_index(index, &self.tree, self.depth, Length(N::to_usize()))
+        Iter::from_index(
+            index,
+            &self.tree.load_full(),
+            self.depth,
+            Length(N::to_usize()),
+        )
     }
 
     fn level_iter_from(&self, index: usize) -> LevelIter<T> {
-        LevelIter::from_index(index, &self.tree, self.depth, Length(N::to_usize()))
+        LevelIter::from_index(
+            index,
+            &self.tree.load_full(),
+            self.depth,
+            Length(N::to_usize()),
+        )
     }
 }
 
@@ -214,12 +242,17 @@ where
                 len: self.len().as_usize(),
             });
         }
-        self.tree = self.tree.with_updated_leaf(index, value, self.depth)?;
+        let updated_tree = self
+            .tree
+            .load()
+            .with_updated_leaf(index, value, self.depth)?;
+        self.tree.store(updated_tree);
+
         Ok(())
     }
 
     fn update<U: UpdateMap<T>>(
-        &mut self,
+        &self,
         updates: U,
         hash_updates: Option<BTreeMap<(usize, usize), Hash256>>,
     ) -> Result<(), Error> {
@@ -231,9 +264,11 @@ where
             // Nothing to do.
             return Ok(());
         }
-        self.tree =
+        let updated_tree =
             self.tree
+                .load()
                 .with_updated_leaves(&updates, 0, self.depth, hash_updates.as_ref())?;
+        self.tree.store(updated_tree);
         Ok(())
     }
 }
@@ -264,9 +299,9 @@ impl<T: Value + Send + Sync, N: Unsigned> tree_hash::TreeHash for Vector<T, N> {
     }
 
     fn tree_hash_root(&self) -> Hash256 {
-        // FIXME(sproul): remove assert
-        assert!(!self.interface.has_pending_updates());
-        self.interface.backing.tree.tree_hash()
+        self.apply_updates_immutable()
+            .expect("updates should apply");
+        self.interface.backing.tree.load().tree_hash()
     }
 }
 
@@ -287,7 +322,7 @@ where
 
 impl<'a, T: Value, N: Unsigned, U: UpdateMap<T>> IntoIterator for &'a Vector<T, N, U> {
     type Item = &'a T;
-    type IntoIter = InterfaceIter<'a, T, U>;
+    type IntoIter = InterfaceIter<'a, T>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.iter()

From 98e0bcb3ab15452300801b8cf771fe18b6800885 Mon Sep 17 00:00:00 2001
From: Michael Sproul <michael@sigmaprime.io>
Date: Fri, 18 Oct 2024 10:32:14 +1100
Subject: [PATCH 2/2] Failed attempt

---
 src/interface.rs  | 33 ++++++++++++++++++--------
 src/lib.rs        |  1 +
 src/list.rs       | 16 +++++++++----
 src/map_option.rs | 60 +++++++++++++++++++++++++++++++++++++++++++++++
 src/vector.rs     | 11 +++++----
 5 files changed, 102 insertions(+), 19 deletions(-)
 create mode 100644 src/map_option.rs

diff --git a/src/interface.rs b/src/interface.rs
index 64efae8..89209be 100644
--- a/src/interface.rs
+++ b/src/interface.rs
@@ -1,20 +1,28 @@
 use crate::level_iter::LevelIter;
 use crate::update_map::UpdateMap;
-use crate::utils::{arb_rwlock, partial_eq_rwlock, updated_length, Length};
+use crate::utils::{arb_arc_swap, partial_eq_arc_swap, updated_length, Length};
 use crate::{
     interface_iter::{InterfaceIter, InterfaceIterCow},
     iter::Iter,
-    Cow, Error, Value, ValueRef,
+    Cow, Error, Tree, Value, ValueRef,
 };
 use arbitrary::Arbitrary;
+use arc_swap::ArcSwap;
 use derivative::Derivative;
-use parking_lot::{RwLock, RwLockReadGuard};
 use std::collections::BTreeMap;
 use std::marker::PhantomData;
+use std::sync::Arc;
 use tree_hash::Hash256;
 
+pub type GetResult<'a, T> = arc_swap::access::MapGuard<
+    &'a ArcSwap<Arc<T>>,
+    (),
+    fn(&'a Arc<Tree<T>>, usize, usize, usize) -> &'a T,
+    &'a T,
+>;
+
 pub trait ImmList<T: Value> {
-    fn get(&self, idx: usize) -> Option<&T>;
+    fn get(&self, idx: usize) -> Option<GetResult<T>>;
 
     fn len(&self) -> Length;
 
@@ -46,9 +54,9 @@ where
     U: UpdateMap<T>,
 {
     pub(crate) backing: B,
-    #[derivative(PartialEq(compare_with = "partial_eq_rwlock"))]
-    #[arbitrary(with = arb_rwlock)]
-    pub(crate) updates: RwLock<U>,
+    #[derivative(PartialEq(compare_with = "partial_eq_arc_swap"))]
+    #[arbitrary(with = arb_arc_swap)]
+    pub(crate) updates: ArcSwap<Arc<U>>,
     pub(crate) _phantom: PhantomData<T>,
 }
 
@@ -61,7 +69,7 @@ where
     fn clone(&self) -> Self {
         Self {
             backing: self.backing.clone(),
-            updates: RwLock::new(self.updates.read().clone()),
+            updates: ArcSwap::new(self.updates.load_full()),
             _phantom: PhantomData,
         }
     }
@@ -76,16 +84,21 @@ where
     pub fn new(backing: B) -> Self {
         Self {
             backing,
-            updates: RwLock::new(U::default()),
+            updates: ArcSwap::new(U::default()),
             _phantom: PhantomData,
         }
     }
 
-    pub fn get(&self, idx: usize) -> Option<ValueRef<'_, T>> {
+    pub fn get(&self, idx: usize) -> Option<GetResult<T>> {
+        panic!()
+
+        // let values_in_updates = ArcSwap::map(&self.updates, || )
+        /*
         RwLockReadGuard::try_map(self.updates.read(), |updates| updates.get(idx))
             .ok()
             .map(ValueRef::Pending)
             .or_else(|| self.backing.get(idx).map(ValueRef::Applied))
+        */
     }
 
     pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
diff --git a/src/lib.rs b/src/lib.rs
index 8ea8d79..62a24ae 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,6 +10,7 @@ pub mod iter;
 pub mod leaf;
 pub mod level_iter;
 pub mod list;
+pub mod map_option;
 pub mod packed_leaf;
 mod repeat;
 pub mod serde;
diff --git a/src/list.rs b/src/list.rs
index b3c26f4..0c45fa0 100644
--- a/src/list.rs
+++ b/src/list.rs
@@ -1,5 +1,5 @@
 use crate::builder::Builder;
-use crate::interface::{ImmList, Interface, MutList};
+use crate::interface::{GetResult, ImmList, Interface, MutList};
 use crate::interface_iter::{InterfaceIter, InterfaceIterCow};
 use crate::iter::Iter;
 use crate::level_iter::{LevelIter, LevelNode};
@@ -261,13 +261,19 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> List<T, N, U> {
     }
 }
 
+use arc_swap::access::Access;
+
 impl<T: Value, N: Unsigned> ImmList<T> for ListInner<T, N> {
-    fn get(&self, index: usize) -> Option<&T> {
+    fn get(&self, index: usize) -> Option<GetResult<T>> {
         // FIXME(sproul): this is fucked unfortunately
         if index < self.len().as_usize() {
-            ArcSwap::map(&self.tree, |tree: &Arc<Tree<T>>| {
-                tree.get_recursive(index, self.depth, self.packing_depth)
-            })
+            Some(
+                &*ArcSwap::map(&self.tree, move |tree: &Arc<Tree<T>>| {
+                    tree.get_recursive(index, self.depth, self.packing_depth)
+                        .unwrap()
+                })
+                .load(),
+            )
         } else {
             None
         }
diff --git a/src/map_option.rs b/src/map_option.rs
new file mode 100644
index 0000000..c4df796
--- /dev/null
+++ b/src/map_option.rs
@@ -0,0 +1,60 @@
+use arc_swap::{access::Access, ArcSwap};
+use std::ops::Deref;
+use std::sync::Arc;
+
+pub struct MapOption<'a, F: Fn(Arc<T>) -> Option<&'a R>, T, R: 'a> {
+    accessor: F,
+    whole_value: &'a ArcSwap<T>,
+}
+
+impl<'a, F, T, R> MapOption<'a, F, T, R>
+where
+    F: Fn(Arc<T>) -> Option<&'a R>,
+{
+    pub fn new(value: &'a ArcSwap<T>, f: F) -> Self {
+        MapOption {
+            accessor: f,
+            whole_value: value,
+        }
+    }
+}
+
+pub struct MapOptionGuard<'a, T, R> {
+    value: Option<&'a R>,
+    parent: Arc<T>,
+}
+
+impl<'a, T, R> Deref for MapOptionGuard<'a, T, R> {
+    type Target = Option<&'a R>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.value
+    }
+}
+
+/*
+pub trait Access<T> {
+    type Guard: Deref<Target = T>;
+
+    // Required method
+    fn load(&self) -> Self::Guard;
+}
+*/
+
+impl<'a, F, T, R> Access<Option<&'a R>> for MapOption<'a, F, T, R>
+where
+    F: Fn(&'a Arc<T>) -> Option<&'a R>,
+{
+    type Guard = MapOptionGuard<'a, T, R>;
+
+    fn load(&self) -> Self::Guard {
+        let loaded_value: Arc<T> = self.whole_value.load_full();
+        let mut guard = MapOptionGuard {
+            value: None,
+            parent: loaded_value,
+        };
+        let value = (self.accessor)(&guard.parent);
+        guard.value = value;
+        guard
+    }
+}
diff --git a/src/vector.rs b/src/vector.rs
index a5bd86e..0b6bd27 100644
--- a/src/vector.rs
+++ b/src/vector.rs
@@ -1,11 +1,11 @@
-use crate::interface::{ImmList, Interface, MutList};
+use crate::interface::{GetResult, ImmList, Interface, MutList};
 use crate::interface_iter::InterfaceIter;
 use crate::iter::Iter;
 use crate::level_iter::LevelIter;
 use crate::tree::RebaseAction;
 use crate::update_map::MaxMap;
 use crate::utils::{arb_arc_swap, partial_eq_arc_swap, Length};
-use crate::{Arc, Cow, Error, List, Tree, UpdateMap, Value, ValueRef};
+use crate::{Arc, Cow, Error, List, Tree, UpdateMap, Value};
 use arbitrary::Arbitrary;
 use arc_swap::ArcSwap;
 use derivative::Derivative;
@@ -93,7 +93,7 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> Vector<T, N, U> {
     }
 
     // Wrap trait methods so we present a Vec-like interface without having to import anything.
-    pub fn get(&self, index: usize) -> Option<ValueRef<T>> {
+    pub fn get(&self, index: usize) -> Option<GetResult<T>> {
         self.interface.get(index)
     }
 
@@ -193,7 +193,9 @@ impl<T: Value, N: Unsigned, U: UpdateMap<T>> From<Vector<T, N, U>> for List<T, N
 }
 
 impl<T: Value, N: Unsigned> ImmList<T> for VectorInner<T, N> {
-    fn get(&self, index: usize) -> Option<&T> {
+    fn get(&self, index: usize) -> Option<GetResult<T>> {
+        None
+        /* FIXME(sproul)
         if index < self.len().as_usize() {
             self.tree
                 .load()
@@ -201,6 +203,7 @@ impl<T: Value, N: Unsigned> ImmList<T> for VectorInner<T, N> {
         } else {
             None
         }
+        */
     }
 
     fn len(&self) -> Length {