[Layered Map] Expose inner layers

Previously, in order to find all changes between two `State`s, we could just "diff" the two map layers and we would get all the KVs that are changed.

This is insufficient with tiered storage. For instance, consider the following sequence of events:
1. At the end of block `N`, key `K` is read and promoted to hot and its value is `V0` in both hot state and cold state.
2. At the end of block `N+1`, its value is updated to `V1` in hot state.
3. At the end of block `N+2`, it's evicted and the value in cold state is updated to `V1`.
4. At the end of block `N+3`, it's read and promoted into hot state again.

If we simply look at the `StateDelta` in the current implementation, we'd see the value change in hot state, but miss the change from `V0` to `V1` in the cold state.

So we need to expose all the inner layers in order to accurately compute what changed between two points in time. This PR simply adds the API on `LayeredMap`.

Author: wqfish

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

@@ -11,12 +11,16 @@ use aptos_crypto::HashValue;
 use aptos_drop_helper::ArcAsyncDrop;
 use aptos_infallible::Mutex;
 use aptos_metrics_core::IntGaugeVecHelper;
-use std::{marker::PhantomData, sync::Arc};
+use std::{
+    marker::PhantomData,
+    sync::{Arc, Weak},
+};
 
 #[derive(Debug)]
 struct LayerInner<K: ArcAsyncDrop, V: ArcAsyncDrop> {
     peak: FlattenPerfectTree<K, V>,
-    children: Mutex<Vec<Arc<LayerInner<K, V>>>>,
+    parent: Weak<Self>,
+    children: Mutex<Vec<Arc<Self>>>,
     use_case: &'static str,
     family: HashValue,
     layer: u64,
@@ -50,6 +54,7 @@ impl<K: ArcAsyncDrop, V: ArcAsyncDrop> LayerInner<K, V> {
         let family = HashValue::random();
         Arc::new(Self {
             peak: FlattenPerfectTree::new_with_empty_nodes(1),
+            parent: Weak::new(),
             children: Mutex::new(Vec::new()),
             use_case,
             family,
@@ -61,6 +66,7 @@ impl<K: ArcAsyncDrop, V: ArcAsyncDrop> LayerInner<K, V> {
     fn spawn(self: &Arc<Self>, child_peak: FlattenPerfectTree<K, V>, base_layer: u64) -> Arc<Self> {
         let child = Arc::new(Self {
             peak: child_peak,
+            parent: Arc::downgrade(self),
             children: Mutex::new(Vec::new()),
             use_case: self.use_case,
             family: self.family,
@@ -108,6 +114,10 @@ impl<K: ArcAsyncDrop, V: ArcAsyncDrop> MapLayer<K, V> {
         }
     }
 
+    pub(crate) fn parent(&self) -> Option<Self> {
+        self.inner.parent.upgrade().map(Self::new)
+    }
+
     pub fn into_layers_view_after(self, base_layer: MapLayer<K, V>) -> LayeredMap<K, V> {
         assert!(base_layer.is_family(&self));
         assert!(base_layer.inner.layer >= self.inner.base_layer);

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

@@ -114,4 +114,27 @@ where
             .into_feet_iter()
             .flat_map(|node| DescendantIterator::new(node, self.base_layer()))
     }
+
+    /// For example, if `self` is `(2, 5]`, this returns `(2, 3], (3, 4], (4, 5]` as a list.
+    pub fn inner_maps(&self) -> Vec<Self> {
+        let num_layers = (self.top_layer.layer() - self.base_layer.layer()) as usize;
+        let mut ret = Vec::with_capacity(num_layers);
+
+        let mut current = self.top_layer.clone();
+        while current.layer() > self.base_layer.layer() {
+            let parent = current.parent().expect("The lower layer must exist.");
+            ret.push(Self::new(parent.clone(), current));
+            current = parent;
+        }
+        ret.reverse();
+
+        assert_eq!(
+            ret.len(),
+            num_layers,
+            "The number of inner maps ({}) does not match num_layers ({})",
+            ret.len(),
+            num_layers
+        );
+        ret
+    }
 }

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

@@ -19,6 +19,10 @@ impl Hash for HashCollide {
     }
 }
 
+fn naive_view_layer<K: Ord, V>(layer: Vec<(K, V)>) -> BTreeMap<K, V> {
+    layer.into_iter().collect()
+}
+
 fn naive_view_layers<K: Ord, V>(layers: impl Iterator<Item = Vec<(K, V)>>) -> BTreeMap<K, V> {
     layers.flat_map(|layer| layer.into_iter()).collect()
 }
@@ -89,7 +93,8 @@ proptest! {
         let layered_map = top_layer.into_layers_view_after(base_layer);
 
         // n.b. notice items_per_update doesn't have a placeholder for the root layer
-        let all = naive_view_layers(items_per_update.drain(base..top));
+        let items_per_update = items_per_update.drain(base..top).collect_vec();
+        let all = naive_view_layers(items_per_update.clone().into_iter());
 
         // get() individually
         for (k, v) in &all {
@@ -99,6 +104,15 @@ proptest! {
         // traversed via iterator
         let traversed = layered_map.iter().collect();
         prop_assert_eq!(all, traversed);
+
+        for (inner_map, items) in layered_map.inner_maps().into_iter().zip_eq(items_per_update.into_iter()) {
+            let all = naive_view_layer(items);
+            for (k, v) in &all {
+                prop_assert_eq!(inner_map.get(k), Some(*v));
+            }
+            let traversed = inner_map.iter().collect();
+            prop_assert_eq!(all, traversed);
+        }
     }
 
     #[test]