Fix post-error root calculation

contracts/callcenter/src/lib.rs

@@ -92,6 +92,11 @@ impl Callcenter {
         call_with_limit(contract, "spend", &(), points_limit)?;
         Ok(())
     }
+
+    /// Just panic.
+    pub fn panik(&self) {
+        panic!("panik");
+    }
 }
 
 /// Expose `Callcenter::query_counter()` to the host
@@ -159,3 +164,9 @@ unsafe fn delegate_transaction(arg_len: u32) -> u32 {
         STATE.delegate_transaction(mod_id, rt)
     })
 }
+
+/// Expose `Callcenter::panik()` to the host
+#[no_mangle]
+unsafe fn panik(arg_len: u32) -> u32 {
+    wrap_call(arg_len, |()| STATE.panik())
+}

crumbles/Cargo.toml

@@ -15,3 +15,8 @@ license = "MPL-2.0"
 [dependencies]
 libc = "0.2"
 rangemap = "1"
+
+[dev-dependencies]
+blake3 = "1"
+hex = "0.4"
+rand = "0.8"

crumbles/src/lib.rs

@@ -831,6 +831,8 @@ unsafe fn segfault_handler(
 mod tests {
     use super::*;
 
+    use rand::rngs::StdRng;
+    use rand::{Rng, SeedableRng};
     use std::thread;
 
     const N_PAGES: usize = 65536;
@@ -969,4 +971,82 @@ mod tests {
             "Slice should be the zeros written afterwards"
         );
     }
+
+    #[test]
+    fn apply_revert_apply() {
+        const N_WRITES: usize = 64;
+        let mut rng = StdRng::seed_from_u64(0xDEAD_BEEF);
+
+        let mut mem = Mmap::new(N_PAGES, PAGE_SIZE)
+            .expect("Instantiating new memory should succeed");
+        let mut mem_alt = Mmap::new(N_PAGES, PAGE_SIZE)
+            .expect("Instantiating new memory should succeed");
+
+        // Snapshot both, make the same changes on both, and apply the changes.
+        mem.snap().expect("Snapshotting should succeed");
+        mem_alt.snap().expect("Snapshotting should succeed");
+
+        for _ in 0..N_WRITES {
+            let i = rng.gen_range(0..N_PAGES);
+            let byte = rng.gen();
+
+            mem[i] = byte;
+            mem_alt[i] = byte;
+        }
+
+        mem.apply().expect("Applying should succeed");
+        mem_alt.apply().expect("Applying should succeed");
+
+        // Snapshot one, make some changes, and revert it.
+        mem.snap().expect("Snapshotting should succeed");
+        for _ in 0..N_WRITES {
+            let i = rng.gen_range(0..N_PAGES);
+            let byte = rng.gen();
+            mem[i] = byte;
+        }
+        mem.revert().expect("Reverting should succeed");
+
+        // Snapshot both, make the same changes on both, and apply the changes.
+        mem.snap().expect("Snapshotting should succeed");
+        mem_alt.snap().expect("Snapshotting should succeed");
+
+        for _ in 0..N_WRITES {
+            let i = rng.gen_range(0..N_PAGES);
+            let byte = rng.gen();
+
+            mem[i] = byte;
+            mem_alt[i] = byte;
+        }
+
+        mem.apply().expect("Applying should succeed");
+        mem_alt.apply().expect("Applying should succeed");
+
+        mem.dirty_pages().zip(mem_alt.dirty_pages()).for_each(
+            |((dirty, clean, index), (alt_dirty, alt_clean, alt_index))| {
+                let hash_dirty = blake3::hash(dirty);
+                let hash_alt_dirty = blake3::hash(alt_dirty);
+
+                let hash_dirty = hex::encode(hash_dirty.as_bytes());
+                let hash_alt_dirty = hex::encode(hash_alt_dirty.as_bytes());
+
+                assert_eq!(
+                    hash_dirty, hash_alt_dirty,
+                    "Dirty state should be the same"
+                );
+
+                let hash_clean = blake3::hash(clean);
+                let hash_alt_clean = blake3::hash(alt_clean);
+
+                let hash_clean = hex::encode(hash_clean.as_bytes());
+                let hash_alt_clean = hex::encode(hash_alt_clean.as_bytes());
+
+                assert_eq!(
+                    hash_clean, hash_alt_clean,
+                    "Clean state should be the same"
+                );
+
+                assert_eq!(index, alt_index, "Index should be the same");
+            },
+        );
+    }
 }

piecrust/CHANGELOG.md

@@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Changed
+
+- De-instantiate modules after call [#296]
+- Change `Session::memory_len` to return `Result<Option<usize>>`, and not
+  require a contract to be instantiated [#296]
+
+### Fixed
+
+- Fix inconsistent state root after erroring call [#296]
+
 ## [0.13.0] - 2023-11-22
 
 ## Added
@@ -304,6 +314,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 -->
+[#296]: https://github.com/dusk-network/piecrust/issues/296
 [#287]: https://github.com/dusk-network/piecrust/issues/287
 [#281]: https://github.com/dusk-network/piecrust/issues/281
 [#273]: https://github.com/dusk-network/piecrust/issues/273

piecrust/Cargo.toml

@@ -46,6 +46,11 @@ name = "crossover"
 path = "tests/crossover.rs"
 required-features = ["debug"]
 
+[[test]]
+name = "commit"
+path = "tests/commit.rs"
+required-features = ["debug"]
+
 [[test]]
 name = "debugger"
 path = "tests/debugger.rs"

piecrust/src/instance.rs

@@ -88,6 +88,7 @@ impl WrappedInstance {
         contract: &WrappedContract,
         memory: Memory,
     ) -> Result<Self, Error> {
+        let mut memory = memory;
         let engine = session.engine().clone();
 
         let env = Env {
@@ -173,6 +174,9 @@ impl WrappedInstance {
             return Err(Error::InvalidArgumentBuffer);
         }
 
+        // A memory is no longer new after one instantiation
+        memory.is_new = false;
+
         let wrapped = WrappedInstance {
             store,
             instance,

piecrust/src/session.rs

@@ -5,7 +5,6 @@
 // Copyright (c) DUSK NETWORK. All rights reserved.
 
 use std::borrow::Cow;
-use std::collections::btree_map::Entry;
 use std::collections::BTreeMap;
 use std::fmt::{Debug, Formatter};
 use std::mem;
@@ -92,7 +91,7 @@ struct SessionInner {
     current: ContractId,
 
     call_tree: CallTree,
-    instance_map: BTreeMap<ContractId, (*mut WrappedInstance, u64)>,
+    instances: BTreeMap<ContractId, *mut WrappedInstance>,
     debug: Vec<String>,
     data: SessionData,
 
@@ -146,7 +145,7 @@ impl Session {
         let inner = SessionInner {
             current: ContractId::uninitialized(),
             call_tree: CallTree::new(),
-            instance_map: BTreeMap::new(),
+            instances: BTreeMap::new(),
             debug: vec![],
             data,
             contract_session,
@@ -451,47 +450,35 @@ impl Session {
 
     /// Returns the current length of the memory of the given contract.
     ///
-    /// If the contract does not exist, or is otherwise not instantiated in this
-    /// session, it will return `None`.
-    pub fn memory_len(&self, contract_id: &ContractId) -> Option<usize> {
-        self.instance(contract_id)
-            .map(|instance| instance.mem_len())
+    /// If the contract does not exist, it will return `None`.
+    pub fn memory_len(
+        &mut self,
+        contract_id: ContractId,
+    ) -> Result<Option<usize>, Error> {
+        Ok(self
+            .inner
+            .contract_session
+            .contract(contract_id)
+            .map_err(|err| Error::PersistenceError(Arc::new(err)))?
+            .map(|data| data.memory.current_len))
     }
 
     pub(crate) fn instance<'a>(
         &self,
         contract_id: &ContractId,
     ) -> Option<&'a mut WrappedInstance> {
-        self.inner
-            .instance_map
-            .get(contract_id)
-            .map(|(instance, _)| {
-                // SAFETY: We guarantee that the instance exists since we're in
-                // control over if it is dropped in `pop`
-                unsafe { &mut **instance }
-            })
-    }
-
-    fn update_instance_count(&mut self, contract_id: ContractId, inc: bool) {
-        match self.inner.instance_map.entry(contract_id) {
-            Entry::Occupied(mut entry) => {
-                let (_, count) = entry.get_mut();
-                if inc {
-                    *count += 1
-                } else {
-                    *count -= 1
-                };
-            }
-            _ => unreachable!("map must have an instance here"),
-        };
+        self.inner.instances.get(contract_id).map(|instance| {
+            // SAFETY: We guarantee that the instance exists since we're in
+            // control over if it is dropped with the session.
+            unsafe { &mut **instance }
+        })
     }
 
     fn clear_stack_and_instances(&mut self) {
         self.inner.call_tree.clear();
 
-        while !self.inner.instance_map.is_empty() {
-            let (_, (instance, _)) =
-                self.inner.instance_map.pop_first().unwrap();
+        while !self.inner.instances.is_empty() {
+            let (_, instance) = self.inner.instances.pop_first().unwrap();
             unsafe {
                 let _ = Box::from_raw(instance);
             };
@@ -582,7 +569,7 @@ impl Session {
         contract: ContractId,
     ) -> Result<usize, Error> {
         let instance = self.new_instance(contract)?;
-        if self.inner.instance_map.get(&contract).is_some() {
+        if self.inner.instances.get(&contract).is_some() {
             panic!("Contract already in the stack: {contract:?}");
         }
 
@@ -591,7 +578,7 @@ impl Session {
         let instance = Box::new(instance);
         let instance = Box::leak(instance) as *mut WrappedInstance;
 
-        self.inner.instance_map.insert(contract, (instance, 1));
+        self.inner.instances.insert(contract, instance);
         Ok(mem_len)
     }
 
@@ -604,7 +591,6 @@ impl Session {
 
         match instance {
             Some(instance) => {
-                self.update_instance_count(contract_id, true);
                 self.inner.call_tree.push(CallTreeElem {
                     contract_id,
                     limit,
@@ -631,15 +617,11 @@ impl Session {
     }
 
     pub(crate) fn move_up_call_tree(&mut self, spent: u64) {
-        if let Some(element) = self.inner.call_tree.move_up(spent) {
-            self.update_instance_count(element.contract_id, false);
-        }
+        self.inner.call_tree.move_up(spent);
     }
 
     pub(crate) fn move_up_prune_call_tree(&mut self) {
-        if let Some(element) = self.inner.call_tree.move_up_prune() {
-            self.update_instance_count(element.contract_id, false);
-        }
+        self.inner.call_tree.move_up_prune();
     }
 
     pub(crate) fn revert_callstack(&mut self) -> Result<(), std::io::Error> {
@@ -730,6 +712,7 @@ impl Session {
                     };
                 }
                 self.move_up_prune_call_tree();
+                self.clear_stack_and_instances();
                 err
             })
             .map_err(Error::normalize)?;
@@ -748,6 +731,7 @@ impl Session {
                     io: Arc::new(err),
                 })?;
         }
+        self.clear_stack_and_instances();
 
         let mut call_tree = CallTree::new();
         mem::swap(&mut self.inner.call_tree, &mut call_tree);

piecrust/tests/commit.rs

@@ -232,6 +232,67 @@ fn multiple_commits() -> Result<(), Error> {
     Ok(())
 }
 
+#[test]
+fn root_equal_on_err() -> Result<(), Error> {
+    let vm = VM::ephemeral()?;
+
+    let mut session = vm.session(SessionData::builder())?;
+
+    let callcenter_id = session.deploy(
+        contract_bytecode!("callcenter"),
+        ContractData::builder(OWNER),
+        LIMIT,
+    )?;
+    let counter_id = session.deploy(
+        contract_bytecode!("counter"),
+        ContractData::builder(OWNER),
+        LIMIT,
+    )?;
+
+    let root = session.commit()?;
+
+    let mut session_after = vm.session(SessionData::builder().base(root))?;
+    let mut session_after_alt =
+        vm.session(SessionData::builder().base(root))?;
+
+    assert_eq!(
+        session_after.root(),
+        session_after_alt.root(),
+        "Roots should be equal at the beginning"
+    );
+
+    session_after
+        .call::<_, ()>(callcenter_id, "panik", &counter_id, LIMIT)
+        .expect_err("Calling with too little gas should error");
+
+    assert_eq!(
+        session_after.root(),
+        session_after_alt.root(),
+        "Roots should be equal immediately after erroring call"
+    );
+
+    session_after.call::<_, ()>(
+        callcenter_id,
+        "increment_counter",
+        &counter_id,
+        LIMIT,
+    )?;
+    session_after_alt.call::<_, ()>(
+        callcenter_id,
+        "increment_counter",
+        &counter_id,
+        LIMIT,
+    )?;
+
+    assert_eq!(
+        session_after.root(),
+        session_after_alt.root(),
+        "Roots should be equal after call"
+    );
+
+    Ok(())
+}
+
 fn increment_counter_and_commit(
     mut session: Session,
     id: ContractId,