piecrust: add Session::migrate

The new `Session::migrate` function allows for swapping contract
bytecode for a different code and perform a migration of the data by
allowing a sequence of operations on a `Session` to be done whilst both
contracts are available, and, once this sequence of operations is done,
swapping the old contract for the new one.

Author: Eduardo Leegwater Simões

piecrust/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Add `Session::migrate` to allow for swapping contract code [#313]
+
 ## [0.15.0] - 2024-01-24
 
 ### Changed
@@ -345,6 +349,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#234]: https://github.com/dusk-network/piecrust/pull/234
 
 <!-- ISSUES -->
+[#301]: https://github.com/dusk-network/piecrust/issues/313
 [#301]: https://github.com/dusk-network/piecrust/issues/301
 [#296]: https://github.com/dusk-network/piecrust/issues/296
 [#287]: https://github.com/dusk-network/piecrust/issues/287

piecrust/src/session.rs

@@ -396,6 +396,47 @@ impl Session {
         })
     }
 
+    /// Migrates a `contract` to a new `bytecode`, performing modifications to
+    /// its state as specified by the closure.
+    ///
+    /// The closure takes a contract ID of where the new contract will be
+    /// available during the migration, and a mutable reference to a session,
+    /// allowing the caller to perform calls and other operations on the new
+    /// (and old) contract.
+    ///
+    /// At the end of the migration, the new contract will be available at the
+    /// given `contract` ID, and the old contract will be removed from the
+    /// state.
+    ///
+    /// # Errors
+    /// The migration may error during execution for a myriad of reasons. The
+    /// caller is encouraged to drop the `Session` should an error occur as it
+    /// will more than likely be left in an inconsistent state.
+    pub fn migrate<'a, A, D, F, const N: usize>(
+        mut self,
+        contract: ContractId,
+        bytecode: &[u8],
+        deploy_data: D,
+        deploy_gas_limit: u64,
+        closure: F,
+    ) -> Result<Self, Error>
+    where
+        A: 'a + for<'b> Serialize<StandardBufSerializer<'b>>,
+        D: Into<ContractData<'a, A, N>>,
+        F: FnOnce(ContractId, &mut Session) -> Result<(), Error>,
+    {
+        let new_contract =
+            self.deploy(bytecode, deploy_data, deploy_gas_limit)?;
+        closure(new_contract, &mut self)?;
+
+        self.inner
+            .contract_session
+            .replace(contract, new_contract)
+            .map_err(|err| PersistenceError(Arc::new(err)))?;
+
+        Ok(self)
+    }
+
     /// Execute a *feeder* call on the current state of this session.
     ///
     /// Feeder calls are used to have the contract be able to report larger

piecrust/src/store.rs

@@ -406,8 +406,13 @@ fn write_commit<P: AsRef<Path>>(
     let mut index = base
         .as_ref()
         .map_or(ContractIndex::default(), |base| base.index.clone());
-    for c in &commit_contracts {
-        index.insert(*c.0, &c.1.memory);
+
+    for (contract_id, contract_data) in &commit_contracts {
+        if contract_data.is_new {
+            index.remove_and_insert(*contract_id, &contract_data.memory);
+        } else {
+            index.insert(*contract_id, &contract_data.memory);
+        }
     }
 
     let root = *index.root();
@@ -497,6 +502,7 @@ fn write_commit_inner<P: AsRef<Path>>(
 
         let mut pages = BTreeSet::new();
 
+        // Write dirty pages and keep track of the page indices.
         for (dirty_page, _, page_index) in contract_data.memory.dirty_pages() {
             let page_path = page_path(&memory_dir, *page_index);
             fs::write(page_path, dirty_page)?;
@@ -507,13 +513,16 @@ fn write_commit_inner<P: AsRef<Path>>(
         let module_path = bytecode_path.with_extension(OBJECTCODE_EXTENSION);
         let metadata_path = bytecode_path.with_extension(METADATA_EXTENSION);
 
-        // If the contract already existed in the base commit, we hard link the
-        // bytecode, module, and metadata files to avoid duplicating them,
-        // otherwise we write them to disk.
+        // If the contract is new, we write the bytecode, module, and metadata
+        // files to disk, otherwise we hard link them to avoid duplicating them.
         //
         // Also, if there is a base commit, we hard link the pages of the
         // contracts that are not dirty.
-        if let Some(base) = &directories.base {
+        if contract_data.is_new {
+            fs::write(bytecode_path, &contract_data.bytecode)?;
+            fs::write(module_path, &contract_data.module.serialize())?;
+            fs::write(metadata_path, &contract_data.metadata)?;
+        } else if let Some(base) = &directories.base {
             if let Some(elem) = base.inner.index.get(contract) {
                 let base_bytecode_path = base.bytecode_dir.join(&contract_hex);
                 let base_module_path =
@@ -539,10 +548,6 @@ fn write_commit_inner<P: AsRef<Path>>(
                     }
                 }
             }
-        } else {
-            fs::write(bytecode_path, &contract_data.bytecode)?;
-            fs::write(module_path, &contract_data.module.serialize())?;
-            fs::write(metadata_path, &contract_data.metadata)?;
         }
     }
 

piecrust/src/store/session.rs

@@ -27,6 +27,7 @@ pub struct ContractDataEntry {
     pub module: Module,
     pub metadata: Metadata,
     pub memory: Memory,
+    pub is_new: bool,
 }
 
 /// The representation of a session with a [`ContractStore`].
@@ -231,6 +232,7 @@ impl ContractSession {
                                     module,
                                     metadata,
                                     memory,
+                                    is_new: false,
                                 })
                                 .clone();
 
@@ -296,12 +298,34 @@ impl ContractSession {
                 module,
                 metadata,
                 memory,
+                is_new: true,
             },
         );
 
         Ok(())
     }
 
+    /// Remove the `old_contract` and move the `new_contract` to the
+    /// `old_contract`, effectively replacing the `old_contract` with
+    /// `new_contract`.
+    pub fn replace(
+        &mut self,
+        old_contract: ContractId,
+        new_contract: ContractId,
+    ) -> io::Result<()> {
+        let new_contract_data =
+            self.contracts.remove(&new_contract).ok_or_else(|| {
+                io::Error::new(
+                    io::ErrorKind::Other,
+                    format!("Contract '{new_contract}' not found"),
+                )
+            })?;
+
+        self.contracts.insert(old_contract, new_contract_data);
+
+        Ok(())
+    }
+
     /// Provides metadata of the contract with a given `contract_id`.
     pub fn contract_metadata(
         &self,

piecrust/src/store/tree.rs

@@ -158,6 +158,11 @@ impl ContractIndex {
             .insert(position_from_contract(&contract), *element.tree.root());
     }
 
+    pub fn remove_and_insert(&mut self, contract: ContractId, memory: &Memory) {
+        self.contracts.remove(&contract);
+        self.insert(contract, memory);
+    }
+
     pub fn root(&self) -> Ref<Hash> {
         self.tree.root()
     }