Migration of block state values

plt/plt-block-state/src/block_state.rs

@@ -6,6 +6,7 @@ use crate::block_state::blob_store::{
 use crate::block_state::cacheable::Cacheable;
 use crate::block_state::external::{ExternalBlockStateOperations, ExternalBlockStateQuery};
 use crate::block_state::hash::Hashable;
+use crate::block_state::migration::Migrate;
 use crate::block_state::state::protocol_level_tokens::ProtocolLevelTokens;
 use crate::block_state::types::AccountWithCanonicalAddress;
 use crate::block_state::types::protocol_level_tokens::{
@@ -31,6 +32,7 @@ pub mod cacheable;
 pub mod external;
 pub mod hash;
 pub mod lfmb_tree;
+pub mod migration;
 mod smart_contract_trie;
 mod state;
 pub mod types;
@@ -80,23 +82,6 @@ impl BlockState {
         MutableBlockState::new(self)
     }
 
-    /// Migrate the PLT block state from one blob store to another.
-    ///
-    /// # Arguments
-    ///
-    /// - `from_loader` Blob store loader for the blob store we are migrating from.
-    /// - `to_storer` Blob store storer for the blob store we are migrating to.
-    /// - `to_protocol_version` Protocol version for the block state to migrate to.
-    pub fn migrate(
-        &self,
-        _from_loader: impl BlobStoreLoad,
-        _to_storer: impl BlobStoreStore,
-        _to_protocol_version: ProtocolVersion,
-    ) -> Self {
-        // todo ar
-        todo!()
-    }
-
     /// See [`blob_store::load_from_store`]. This function only differs by taking
     /// protocol version as argument.
     pub fn load_from_store(
@@ -130,6 +115,25 @@ impl BlockState {
             data: BlockStateData { tokens },
         })
     }
+
+    /// See [`Migrate::migrate`]. This function only differs by taking
+    /// protocol version as argument.
+    pub fn migrate(
+        &self,
+        from_loader: &impl BlobStoreLoad,
+        to_storer: &mut impl BlobStoreStore,
+        to_protocol_version: ProtocolVersion,
+    ) -> BlockStateResult<Self>
+    where
+        Self: Sized,
+    {
+        let new_tokens = self.data.tokens.migrate(from_loader, to_storer)?;
+
+        Ok(Self {
+            protocol_version: to_protocol_version,
+            data: BlockStateData { tokens: new_tokens },
+        })
+    }
 }
 
 impl Storable for BlockState {

plt/plt-block-state/src/block_state/blob_reference.rs

@@ -1,7 +1,15 @@
-//! Blob store reference types. A blob store reference represents the value of a block state component,
-//! that may be in the blob store, in memory, or both. A blob store reference to a block state component
-//! is stored in the blob store as a reference to another location in the blob store where the block
-//! state component is stored. This allows for
+//! Definition of blob store reference types. A blob store reference may largely be though of as
+//! a pointer to a value in the blob store, but the representation of the value may vary during
+//! the lifetime of the reference. The value may be represented in,
+//!
+//! * the blob store only (the reference is a pure blob store pointer)
+//! * memory (the value will eventually be written to the blob store),
+//! * both the blob store and memory (the reference is a pointer
+//!   that also caches the value in memory).
+//!
+//! A blob store reference to a block state component
+//! is stored in the blob store as a [blob location](super::blob_store::BlobStoreLocation)
+//! This allows for
 //!
 //! * sharing block state components pointed to by blob references between different block states
 //! * deciding when to load the referenced block state component,

plt/plt-block-state/src/block_state/blob_reference/hashed_cacheable_reference.rs

@@ -7,6 +7,7 @@ use crate::block_state::blob_store::{
 };
 use crate::block_state::cacheable::Cacheable;
 use crate::block_state::hash::Hashable;
+use crate::block_state::migration::Migrate;
 use crate::block_state::utils::OwnedOrBorrowed;
 use crate::block_state_interface::BlockStateResult;
 use concordium_base::common::{Buffer, Get, Put};
@@ -267,6 +268,22 @@ impl<V: Hashable + Loadable> Hashable for HashedCacheableRef<V> {
     }
 }
 
+impl<V: Migrate + Storable + Loadable> Migrate for HashedCacheableRef<V> {
+    fn migrate(
+        &self,
+        from_loader: &impl BlobStoreLoad,
+        to_storer: &mut impl BlobStoreStore,
+    ) -> BlockStateResult<Self>
+    where
+        Self: Sized,
+    {
+        let migrated_value = self.value(from_loader)?.migrate(from_loader, to_storer)?;
+        let new_hcr = Self::new(migrated_value);
+        new_hcr.inner.repr.get_reference_or_store(to_storer);
+        Ok(new_hcr)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -340,6 +357,35 @@ mod tests {
         assert_eq!(val_ref_tmp, val_ref);
     }
 
+    /// Test migrate a value:
+    ///
+    /// * create value in source blob store
+    /// * migrate to new blob store
+    /// * load the new value from the new blob store
+    #[test]
+    fn test_migrate() {
+        let mut from_store = BlobStoreStub::default();
+        let mut to_store = BlobStoreStub::default();
+
+        // Create new value and store it in the source blob store
+        let val = TestRef::new(StoreSerialized(1u64));
+        blob_store::store_to_store(&mut from_store, &val);
+
+        // Migrate to destination blob store
+        let new_val = val.migrate(&from_store, &mut to_store).unwrap();
+        let new_blob_loc = blob_store::store_to_store(&mut to_store, &new_val);
+        drop(val);
+
+        // Assert migrated reference
+        assert_cached_repr(&new_val);
+        assert_eq!(*new_val.value(&to_store).unwrap(), StoreSerialized(1));
+        drop(new_val);
+
+        // Load migrated reference and assert
+        let new_val2: TestRef = blob_store::load_from_store(&to_store, new_blob_loc).unwrap();
+        assert_eq!(*new_val2.value(&to_store).unwrap(), StoreSerialized(1));
+    }
+
     /// Test storing cached value.
     #[test]
     fn test_store_cached_value() {
@@ -515,11 +561,11 @@ mod tests {
         );
     }
 
+    type NestedTestRef = HashedCacheableRef<HashedCacheableRef<StoreSerialized<u64>>>;
+
     /// Test store, load and cache a reference with a nested reference.
     #[test]
     fn test_nested_reference_store_load_and_cache() {
-        type NestedTestRef = HashedCacheableRef<HashedCacheableRef<StoreSerialized<u64>>>;
-
         let mut store = BlobStoreStub::default();
         let val1 = HashedCacheableRef::new(HashedCacheableRef::new(StoreSerialized(1u64)));
 
@@ -542,4 +588,34 @@ mod tests {
         assert_eq!(*val_tmp, StoreSerialized(1));
         assert_eq!(val_ref_tmp, val_nested_blob_ref);
     }
+
+    /// Test migrate a reference with a nested reference.
+    #[test]
+    fn test_nested_reference_migrate() {
+        let mut from_store = BlobStoreStub::default();
+        let mut to_store = BlobStoreStub::default();
+
+        // Create new value and store it in the source blob store
+        let val = HashedCacheableRef::new(HashedCacheableRef::new(StoreSerialized(1u64)));
+        blob_store::store_to_store(&mut from_store, &val);
+
+        // Migrate to destination blob store
+        let new_val = val.migrate(&from_store, &mut to_store).unwrap();
+        let new_blob_loc = blob_store::store_to_store(&mut to_store, &new_val);
+        drop(val);
+
+        // Assert migrated reference
+        assert_eq!(
+            *new_val.value(&to_store).unwrap().value(&to_store).unwrap(),
+            StoreSerialized(1)
+        );
+        drop(new_val);
+
+        // Load migrated reference and assert
+        let new_val2: NestedTestRef = blob_store::load_from_store(&to_store, new_blob_loc).unwrap();
+        assert_eq!(
+            *new_val2.value(&to_store).unwrap().value(&to_store).unwrap(),
+            StoreSerialized(1)
+        );
+    }
 }

plt/plt-block-state/src/block_state/blob_store.rs

@@ -151,7 +151,7 @@ pub mod test_stub {
     use super::*;
 
     /// In-memory blob store stub implemented via a `Vec`.
-    #[derive(Default, Debug)]
+    #[derive(Default, Debug, Clone)]
     pub struct BlobStoreStub(pub Vec<u8>);
 
     impl BlobStoreStore for BlobStoreStub {

plt/plt-block-state/src/block_state/lfmb_tree.rs

@@ -9,6 +9,7 @@ use crate::block_state::blob_store::{
 use crate::block_state::cacheable::Cacheable;
 use crate::block_state::hash;
 use crate::block_state::hash::Hashable;
+use crate::block_state::migration::Migrate;
 use crate::block_state::utils::OwnedOrBorrowed;
 use crate::block_state_interface::{BlockStateFailure, BlockStateResult};
 use concordium_base::common::{Buffer, Get, Put};
@@ -873,6 +874,52 @@ impl<V: Cacheable + Loadable> Cacheable for Subtree<V> {
     }
 }
 
+impl<K, V: Migrate + Storable + Loadable> Migrate for LfmbTree<K, V> {
+    fn migrate(
+        &self,
+        from_loader: &impl BlobStoreLoad,
+        to_storer: &mut impl BlobStoreStore,
+    ) -> BlockStateResult<Self>
+    where
+        Self: Sized,
+    {
+        let new_inner = match &self.inner {
+            LfmbTreeInner::Empty => LfmbTreeInner::Empty,
+            LfmbTreeInner::NonEmpty(size, subtree) => {
+                LfmbTreeInner::NonEmpty(*size, subtree.migrate(from_loader, to_storer)?)
+            }
+        };
+
+        Ok(Self {
+            inner: new_inner,
+            _key_type: PhantomData,
+        })
+    }
+}
+
+impl<V: Migrate + Storable + Loadable> Migrate for Subtree<V> {
+    fn migrate(
+        &self,
+        from_loader: &impl BlobStoreLoad,
+        to_storer: &mut impl BlobStoreStore,
+    ) -> BlockStateResult<Self>
+    where
+        Self: Sized,
+    {
+        Ok(match self {
+            Subtree::Leaf(value_ref) => {
+                let new_value_ref = value_ref.migrate(from_loader, to_storer)?;
+                Subtree::Leaf(new_value_ref)
+            }
+            Subtree::Node(height, left_ref, right_ref) => {
+                let new_left_ref = left_ref.migrate(from_loader, to_storer)?;
+                let new_right_ref = right_ref.migrate(from_loader, to_storer)?;
+                Subtree::Node(*height, new_left_ref, new_right_ref)
+            }
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -1067,7 +1114,52 @@ mod tests {
             let tree2: TestTree = blob_store::load_from_store(&store, blob_ref).unwrap();
 
             // Assert loaded tree is equal to the tree we started with
-            assert_trees_eq(&store, &tree1, &tree2, format!("loaded tree of size {}", i));
+            assert_trees_eq(
+                &store,
+                &store,
+                &tree1,
+                &tree2,
+                format!("loaded tree of size {}", i),
+            );
+        }
+    }
+
+    /// Tests migrating tree into new blob store
+    #[test]
+    fn prop_test_migrate() {
+        for i in 0..100 {
+            let mut from_store = BlobStoreStub::default();
+            let mut to_store = BlobStoreStub::default();
+
+            // Create tree and store it
+            let tree = create_tree(&mut from_store, i);
+            blob_store::store_to_store(&mut from_store, &tree);
+
+            // Migrate the tree and store it
+            let new_tree = tree.migrate(&from_store, &mut to_store).unwrap();
+            let new_blob_loc = blob_store::store_to_store(&mut to_store, &new_tree);
+
+            // Assert migrated tree is equal to the tree we started with
+            assert_trees_eq(
+                &from_store,
+                &to_store,
+                &tree,
+                &new_tree,
+                format!("loaded tree of size {}", i),
+            );
+            drop(new_tree);
+
+            // Load migrated tree from destination store
+            let new_tree2: TestTree = blob_store::load_from_store(&to_store, new_blob_loc).unwrap();
+
+            // Assert tree loaded from destination store is equal to the tree we started with
+            assert_trees_eq(
+                &from_store,
+                &to_store,
+                &tree,
+                &new_tree2,
+                format!("loaded tree of size {}", i),
+            );
         }
     }
 
@@ -1084,7 +1176,13 @@ mod tests {
             tree2.cache_reference_values(&store).expect("cache");
 
             // Assert cached tree is identical to the tree with started with
-            assert_trees_eq(&store, &tree1, &tree2, format!("cached tree of size {}", i));
+            assert_trees_eq(
+                &store,
+                &store,
+                &tree1,
+                &tree2,
+                format!("cached tree of size {}", i),
+            );
 
             // Assert that when caching again or looking up entries, we don't need to read from the blob store again.
             // We assert that by using UnreachableBlobStore.
@@ -1198,7 +1296,8 @@ mod tests {
 
     /// Assert node structure and values in tree are equal.
     fn assert_trees_eq<K: Debug, V: Loadable + Clone + PartialEq + Debug>(
-        loader: &impl BlobStoreLoad,
+        loader1: &impl BlobStoreLoad,
+        loader2: &impl BlobStoreLoad,
         tree1: &LfmbTree<K, V>,
         tree2: &LfmbTree<K, V>,
         context: String,
@@ -1212,7 +1311,7 @@ mod tests {
                 LfmbTreeInner::NonEmpty(size2, subtree2),
             ) => {
                 assert_eq!(size1, size2);
-                assert_subtrees_eq(loader, subtree1, subtree2, context.clone());
+                assert_subtrees_eq(loader1, loader2, subtree1, subtree2, context.clone());
             }
             (_, _) => {
                 panic!("{}: trees not equal: {:?}, {:?}", context, tree1, tree2);
@@ -1222,28 +1321,29 @@ mod tests {
 
     /// Assert node structure and values in subtree are equal.
     fn assert_subtrees_eq<V: Loadable + Clone + PartialEq + Debug>(
-        loader: &impl BlobStoreLoad,
+        loader1: &impl BlobStoreLoad,
+        loader2: &impl BlobStoreLoad,
         subtree1: &Subtree<V>,
         subtree2: &Subtree<V>,
         context: String,
     ) {
         match (subtree1, subtree2) {
             (Subtree::Leaf(val_ref1), Subtree::Leaf(val_ref2)) => {
-                let val1 = &*val_ref1.value(loader).unwrap();
-                let val2 = &*val_ref2.value(loader).unwrap();
+                let val1 = &*val_ref1.value(loader1).unwrap();
+                let val2 = &*val_ref2.value(loader2).unwrap();
                 assert_eq!(val1, val2, "{}: leaf value", context);
             }
             (
                 Subtree::Node(height1, left_ref1, right_ref1),
                 Subtree::Node(height2, left_ref2, right_ref2),
             ) => {
                 assert_eq!(height1, height2);
-                let left1 = &*left_ref1.value(loader).unwrap();
-                let right1 = &*right_ref1.value(loader).unwrap();
-                let left2 = &*left_ref2.value(loader).unwrap();
-                let right2 = &*right_ref2.value(loader).unwrap();
-                assert_subtrees_eq(loader, left1, left2, context.clone());
-                assert_subtrees_eq(loader, right1, right2, context.clone());
+                let left1 = &*left_ref1.value(loader1).unwrap();
+                let right1 = &*right_ref1.value(loader1).unwrap();
+                let left2 = &*left_ref2.value(loader2).unwrap();
+                let right2 = &*right_ref2.value(loader2).unwrap();
+                assert_subtrees_eq(loader1, loader2, left1, left2, context.clone());
+                assert_subtrees_eq(loader1, loader2, right1, right2, context.clone());
             }
             (_, _) => {
                 panic!("subtrees not equal: {:?}, {:?}", subtree1, subtree2);