Merge pull request #320 from dusk-network/migrate-bytecode-313

Add bytecode migrations

Author: Eduardo Leegwater Simões

contracts/Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "constructor",
     "counter",
     "debugger",
+    "double_counter",
     "eventer",
     "everest",
     "fallible_counter",

contracts/double_counter/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "double_counter"
+version = "0.1.0"
+authors = [
+    "Eduardo Leegwater Simões <[email protected]>",
+]
+edition = "2021"
+
+license = "MPL-2.0"
+
+[dependencies]
+piecrust-uplink = { path = "../../piecrust-uplink", features = ["abi", "dlmalloc"] }
+
+[lib]
+crate-type = ["cdylib", "rlib"]

contracts/double_counter/src/lib.rs

@@ -0,0 +1,61 @@
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+//
+// Copyright (c) DUSK NETWORK. All rights reserved.
+
+//! Contract to implement a simple double counter that can be read and
+//! have either value incremented by one count.
+
+#![no_std]
+
+use piecrust_uplink as uplink;
+
+/// Struct that describes the state of the DoubleCounter contract
+pub struct DoubleCounter {
+    left_value: i64,
+    right_value: i64,
+}
+
+/// State of the DoubleCounter contract
+static mut STATE: DoubleCounter = DoubleCounter {
+    left_value: 0xfc,
+    right_value: 0xcf,
+};
+
+impl DoubleCounter {
+    /// Read the value of the counter
+    pub fn read_values(&self) -> (i64, i64) {
+        (self.left_value, self.right_value)
+    }
+
+    /// Increment the value of the left counter by 1
+    pub fn increment_left(&mut self) {
+        let value = self.left_value + 1;
+        self.left_value = value;
+    }
+
+    /// Increment the value of the right counter by 1
+    pub fn increment_right(&mut self) {
+        let value = self.right_value + 1;
+        self.right_value = value;
+    }
+}
+
+/// Expose `Counter::read_value()` to the host
+#[no_mangle]
+unsafe fn read_values(arg_len: u32) -> u32 {
+    uplink::wrap_call(arg_len, |_: ()| STATE.read_values())
+}
+
+/// Expose `Counter::increment_left()` to the host
+#[no_mangle]
+unsafe fn increment_left(arg_len: u32) -> u32 {
+    uplink::wrap_call(arg_len, |_: ()| STATE.increment_left())
+}
+
+/// Expose `Counter::increment_right()` to the host
+#[no_mangle]
+unsafe fn increment_right(arg_len: u32) -> u32 {
+    uplink::wrap_call(arg_len, |_: ()| STATE.increment_right())
+}

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()
     }

piecrust/tests/persistence.rs

@@ -148,3 +148,63 @@ fn contracts_persistence() -> Result<(), Error> {
     );
     Ok(())
 }
+
+#[test]
+fn migration() -> Result<(), Error> {
+    let vm = VM::ephemeral()?;
+    let mut session = vm.session(SessionData::builder())?;
+
+    let contract = session.deploy(
+        contract_bytecode!("counter"),
+        ContractData::builder(OWNER),
+        LIMIT,
+    )?;
+
+    session.call::<_, ()>(contract, "increment", &(), LIMIT)?;
+    session.call::<_, ()>(contract, "increment", &(), LIMIT)?;
+
+    let root = session.commit()?;
+
+    let mut session = vm.session(SessionData::builder().base(root))?;
+
+    session = session.migrate(
+        contract,
+        contract_bytecode!("double_counter"),
+        ContractData::builder(OWNER),
+        LIMIT,
+        |new_contract, session| {
+            let old_counter_value = session
+                .call::<_, i64>(contract, "read_value", &(), LIMIT)?
+                .data;
+
+            let (left_counter_value, _) = session
+                .call::<_, (i64, i64)>(new_contract, "read_values", &(), LIMIT)?
+                .data;
+
+            let diff = old_counter_value - left_counter_value;
+            for _ in 0..diff {
+                session.call::<_, ()>(
+                    new_contract,
+                    "increment_left",
+                    &(),
+                    LIMIT,
+                )?;
+            }
+
+            Ok(())
+        },
+    )?;
+
+    let root = session.commit()?;
+
+    let mut session = vm.session(SessionData::builder().base(root))?;
+
+    let (left_counter, right_counter) = session
+        .call::<_, (i64, i64)>(contract, "read_values", &(), LIMIT)?
+        .data;
+
+    assert_eq!(left_counter, 0xfe);
+    assert_eq!(right_counter, 0xcf);
+
+    Ok(())
+}