FptFoot wrapping ArcSwapOption<NodeRef>

experimental/storage/layered-map/Cargo.toml

@@ -18,6 +18,7 @@ aptos-crypto = { workspace = true }
 aptos-drop-helper = { workspace = true }
 aptos-infallible = { workspace = true }
 aptos-metrics-core = { workspace = true }
+arc-swap = { workspace = true }
 bitvec = "1.0.1"
 itertools = { workspace = true }
 once_cell = { workspace = true }

experimental/storage/layered-map/src/flatten_perfect_tree.rs

@@ -5,38 +5,71 @@ use crate::{
     node::{NodeRef, NodeStrongRef},
     utils,
 };
+use arc_swap::ArcSwapOption;
 use std::{
     fmt,
     fmt::{Debug, Formatter},
+    sync::Arc,
 };
 
-pub(crate) struct FlattenPerfectTree<K, V> {
-    leaves: Vec<NodeRef<K, V>>,
+pub(crate) struct FptFoot<K, V> {
+    /// `None` represents NodeRef::Empty, to avoid unnecessary indirection
+    node: ArcSwapOption<NodeRef<K, V>>,
 }
 
-impl<K, V> FlattenPerfectTree<K, V> {
-    pub fn new_with_empty_nodes(height: usize) -> Self {
-        let num_leaves = if height == 0 { 0 } else { 1 << (height - 1) };
-
+impl<K, V> FptFoot<K, V> {
+    pub fn empty() -> Self {
         Self {
-            leaves: vec![NodeRef::Empty; num_leaves],
+            node: ArcSwapOption::new(None),
         }
     }
 
-    pub fn get_ref(&self) -> FptRef<K, V> {
-        FptRef {
-            leaves: &self.leaves,
-        }
+    pub fn get(&self) -> NodeRef<K, V> {
+        self.node
+            .load()
+            .as_ref()
+            .map(|node_ref_arc| node_ref_arc.as_ref())
+            .cloned()
+            .unwrap_or_else(|| NodeRef::Empty)
+    }
+
+    pub fn get_strong(&self, base_layer: u64) -> NodeStrongRef<K, V> {
+        self.get().get_strong(base_layer)
+    }
+
+    pub fn set(&self, node_ref: NodeRef<K, V>) {
+        self.node.store(Self::empty_to_none(node_ref))
     }
 
-    pub fn get_mut(&mut self) -> FptRefMut<K, V> {
-        FptRefMut {
-            leaves: &mut self.leaves,
+    fn empty_to_none(node_ref: NodeRef<K, V>) -> Option<Arc<NodeRef<K, V>>> {
+        if let NodeRef::Empty = node_ref {
+            None
+        } else {
+            Some(Arc::new(node_ref.clone()))
         }
     }
+}
+
+pub(crate) struct FlattenPerfectTree<K, V> {
+    feet: Vec<FptFoot<K, V>>,
+}
+
+impl<K, V> FlattenPerfectTree<K, V> {
+    pub fn new_with_empty_feet(height: usize) -> Self {
+        let num_leaves = if height == 0 { 0 } else { 1 << (height - 1) };
+
+        let mut feet = Vec::new();
+        feet.resize_with(num_leaves, FptFoot::empty);
+
+        Self { feet }
+    }
+
+    pub fn get_ref(&self) -> FptRef<K, V> {
+        FptRef { feet: &self.feet }
+    }
 
     pub(crate) fn take_for_drop(&mut self) -> Self {
-        let mut ret = Self { leaves: Vec::new() };
+        let mut ret = Self { feet: Vec::new() };
         std::mem::swap(self, &mut ret);
 
         ret
@@ -45,64 +78,43 @@ impl<K, V> FlattenPerfectTree<K, V> {
 
 impl<K, V> Debug for FlattenPerfectTree<K, V> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        write!(f, "FlattenPerfectTree({})", self.leaves.len())
+        write!(f, "FlattenPerfectTree({})", self.feet.len())
     }
 }
 
 pub(crate) struct FptRef<'a, K, V> {
-    leaves: &'a [NodeRef<K, V>],
+    feet: &'a [FptFoot<K, V>],
 }
 
 impl<'a, K, V> FptRef<'a, K, V> {
     pub fn num_leaves(&self) -> usize {
-        self.leaves.len()
+        self.feet.len()
     }
 
     pub fn expect_sub_trees(self) -> (Self, Self) {
         assert!(!self.is_single_node());
-        let (left, right) = self.leaves.split_at(self.leaves.len() / 2);
-        (Self { leaves: left }, Self { leaves: right })
+        let (left, right) = self.feet.split_at(self.feet.len() / 2);
+        (Self { feet: left }, Self { feet: right })
     }
 
     pub fn is_single_node(&self) -> bool {
-        self.leaves.len() == 1
+        self.feet.len() == 1
     }
 
-    pub fn expect_single_node(&self, base_layer: u64) -> NodeStrongRef<K, V> {
+    pub fn expect_single_node(&self) -> &'a FptFoot<K, V> {
         assert!(self.is_single_node());
-        self.leaves[0].get_strong(base_layer)
+        &self.feet[0]
     }
 
-    pub fn expect_foot(&self, foot: usize, base_layer: u64) -> NodeStrongRef<K, V> {
-        self.leaves[foot].get_strong(base_layer)
+    pub fn expect_foot(&self, foot: usize) -> &'a FptFoot<K, V> {
+        &self.feet[foot]
     }
 
     pub fn height(&self) -> usize {
-        utils::binary_tree_height(self.leaves.len())
+        utils::binary_tree_height(self.feet.len())
     }
 
-    pub fn into_feet_iter(self) -> impl Iterator<Item = &'a NodeRef<K, V>> {
-        self.leaves.iter()
-    }
-}
-
-pub(crate) struct FptRefMut<'a, K, V> {
-    leaves: &'a mut [NodeRef<K, V>],
-}
-
-impl<'a, K, V> FptRefMut<'a, K, V> {
-    pub fn is_single_node(&self) -> bool {
-        self.leaves.len() == 1
-    }
-
-    pub fn expect_into_single_node_mut(self) -> &'a mut NodeRef<K, V> {
-        assert!(self.is_single_node());
-        &mut self.leaves[0]
-    }
-
-    pub fn expect_into_sub_trees(self) -> (Self, Self) {
-        assert!(!self.is_single_node());
-        let (left, right) = self.leaves.split_at_mut(self.leaves.len() / 2);
-        (Self { leaves: left }, Self { leaves: right })
+    pub fn into_feet_iter(self) -> impl 'a + Iterator<Item = NodeRef<K, V>> {
+        self.feet.iter().map(|foot| foot.get())
     }
 }

experimental/storage/layered-map/src/iterator.rs

@@ -6,7 +6,7 @@ use aptos_drop_helper::ArcAsyncDrop;
 use std::sync::Arc;
 
 pub(crate) struct DescendantIterator<'a, K: ArcAsyncDrop, V: ArcAsyncDrop> {
-    root: Option<&'a NodeRef<K, V>>,
+    root: Option<NodeRef<K, V>>,
     base_layer: u64,
     current_leaf: Option<Box<dyn 'a + Iterator<Item = (K, V)>>>,
     ancestors: Vec<Arc<InternalNode<K, V>>>,
@@ -17,7 +17,7 @@ where
     K: ArcAsyncDrop + Clone,
     V: ArcAsyncDrop + Clone,
 {
-    pub fn new(root: &'a NodeRef<K, V>, base_layer: u64) -> Self {
+    pub fn new(root: NodeRef<K, V>, base_layer: u64) -> Self {
         Self {
             root: Some(root),
             base_layer,
@@ -58,7 +58,7 @@ where
             match &mut self.current_leaf {
                 None => {
                     if let Some(root) = self.root.take() {
-                        // Iterater not started yet, consume root and go down.
+                        // Iterator not started yet, consume root and go down.
                         self.current_leaf =
                             self.find_next_leaf(Some(root.get_strong(self.base_layer)));
                     } else if self.ancestors.is_empty() {

experimental/storage/layered-map/src/layer.rs

@@ -49,7 +49,7 @@ impl<K: ArcAsyncDrop, V: ArcAsyncDrop> LayerInner<K, V> {
     fn new_family(use_case: &'static str) -> Arc<Self> {
         let family = HashValue::random();
         Arc::new(Self {
-            peak: FlattenPerfectTree::new_with_empty_nodes(1),
+            peak: FlattenPerfectTree::new_with_empty_feet(1),
             children: Mutex::new(Vec::new()),
             use_case,
             family,

experimental/storage/layered-map/src/map/mod.rs

@@ -66,7 +66,7 @@ where
             foot = foot << 1 | bits.next().expect("bits exhausted") as usize;
         }
 
-        self.get_under_node(peak.expect_foot(foot, self.base_layer()), key, &mut bits)
+        self.get_under_node(peak.expect_foot(foot).get_strong(self.base_layer()), key, &mut bits)
     }
 
     fn get_under_node(

experimental/storage/layered-map/src/map/new_layer_impl.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    flatten_perfect_tree::{FlattenPerfectTree, FptRef, FptRefMut},
+    flatten_perfect_tree::{FlattenPerfectTree, FptFoot, FptRef},
     metrics::TIMER,
     node::{CollisionCell, LeafContent, LeafNode, NodeRef, NodeStrongRef},
     utils::binary_tree_height,
@@ -38,13 +38,13 @@ where
             .collect_vec();
 
         let height = Self::new_peak_height(self.top_layer.peak().num_leaves(), items.len());
-        let mut new_peak = FlattenPerfectTree::new_with_empty_nodes(height);
+        let new_peak = FlattenPerfectTree::new_with_empty_feet(height);
         let builder = SubTreeBuilder {
             layer: self.top_layer.layer() + 1,
             base_layer: self.base_layer(),
             depth: 0,
             position_info: PositionInfo::new(self.top_layer.peak(), self.base_layer()),
-            output_position_info: OutputPositionInfo::new(new_peak.get_mut()),
+            output_position_info: OutputPositionInfo::new(new_peak.get_ref()),
             items: &items,
         };
         builder.build().finalize();
@@ -101,7 +101,7 @@ enum PositionInfo<'a, K, V> {
 impl<'a, K, V> PositionInfo<'a, K, V> {
     fn new(peak: FptRef<'a, K, V>, base_layer: u64) -> Self {
         if peak.num_leaves() == 1 {
-            Self::PeakFootOrBelow(peak.expect_single_node(base_layer))
+            Self::PeakFootOrBelow(peak.expect_single_node().get_strong(base_layer))
         } else {
             Self::AbovePeakFeet(peak)
         }
@@ -126,8 +126,8 @@ impl<'a, K, V> PositionInfo<'a, K, V> {
                 let (left, right) = fpt.expect_sub_trees();
                 if left.is_single_node() {
                     (
-                        PeakFootOrBelow(left.expect_single_node(base_layer)),
-                        PeakFootOrBelow(right.expect_single_node(base_layer)),
+                        PeakFootOrBelow(left.expect_single_node().get_strong(base_layer)),
+                        PeakFootOrBelow(right.expect_single_node().get_strong(base_layer)),
                     )
                 } else {
                     (AbovePeakFeet(left), AbovePeakFeet(right))
@@ -142,13 +142,13 @@ impl<'a, K, V> PositionInfo<'a, K, V> {
 }
 
 enum OutputPositionInfo<'a, K, V> {
-    AboveOrAtPeakFeet(FptRefMut<'a, K, V>),
+    AboveOrAtPeakFeet(FptRef<'a, K, V>),
     BelowPeakFeet,
 }
 
 impl<'a, K, V> OutputPositionInfo<'a, K, V> {
-    pub fn new(fpt_mut: FptRefMut<'a, K, V>) -> Self {
-        Self::AboveOrAtPeakFeet(fpt_mut)
+    pub fn new(peak: FptRef<'a, K, V>) -> Self {
+        Self::AboveOrAtPeakFeet(peak)
     }
 
     pub fn is_above_peak_feet(&self) -> bool {
@@ -171,15 +171,15 @@ impl<'a, K, V> OutputPositionInfo<'a, K, V> {
         OutputPositionInfo<'a, K, V>,
     ) {
         match self {
-            OutputPositionInfo::AboveOrAtPeakFeet(fpt_mut) => {
-                if fpt_mut.is_single_node() {
+            OutputPositionInfo::AboveOrAtPeakFeet(fpt) => {
+                if fpt.is_single_node() {
                     (
-                        PendingBuild::FootOfPeak(fpt_mut.expect_into_single_node_mut()),
+                        PendingBuild::FootOfPeak(fpt.expect_single_node()),
                         OutputPositionInfo::BelowPeakFeet,
                         OutputPositionInfo::BelowPeakFeet,
                     )
                 } else {
-                    let (left, right) = fpt_mut.expect_into_sub_trees();
+                    let (left, right) = fpt.expect_sub_trees();
                     (
                         PendingBuild::AbovePeakFeet,
                         OutputPositionInfo::AboveOrAtPeakFeet(left),
@@ -198,16 +198,16 @@ impl<'a, K, V> OutputPositionInfo<'a, K, V> {
 
 enum PendingBuild<'a, K, V> {
     AbovePeakFeet,
-    FootOfPeak(&'a mut NodeRef<K, V>),
+    FootOfPeak(&'a FptFoot<K, V>),
     BelowPeakFeet,
 }
 
 impl<'a, K, V> PendingBuild<'a, K, V> {
     fn seal_with_node(&mut self, node: NodeRef<K, V>) -> BuiltSubTree<K, V> {
         match self {
             PendingBuild::AbovePeakFeet => unreachable!("Trying to put node above peak feet."),
-            PendingBuild::FootOfPeak(ref_mut) => {
-                **ref_mut = node;
+            PendingBuild::FootOfPeak(foot) => {
+                foot.set(node);
                 BuiltSubTree::InOrAtFootOfPeak
             },
             PendingBuild::BelowPeakFeet => BuiltSubTree::BelowPeak(node),