piecrust/

session.rs

// 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.

use std::borrow::Cow;
use std::collections::btree_set::Iter;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::{Debug, Formatter};
use std::mem;
use std::ptr::NonNull;
use std::sync::{Arc, mpsc};

use bytecheck::CheckBytes;
use dusk_wasmtime::{
    Engine, LinearMemory, MemoryCreator, MemoryType, Module as WasmtimeModule,
};
use piecrust_uplink::{
    ARGBUF_LEN, CONTRACT_ID_BYTES, ContractId, Event, SCRATCH_BUF_BYTES,
};
use rkyv::ser::Serializer;
use rkyv::ser::serializers::{
    BufferScratch, BufferSerializer, CompositeSerializer,
};
use rkyv::validation::validators::DefaultValidator;
use rkyv::{Archive, Deserialize, Infallible, Serialize, check_archived_root};

use crate::call_tree::{CallTree, CallTreeElem};
use crate::contract::{ContractData, ContractMetadata, WrappedContract};
use crate::error::Error::{self, InitalizationError, PersistenceError};
use crate::instance::WrappedInstance;
use crate::store::{ContractSession, PAGE_SIZE, PageOpening};
use crate::types::StandardBufSerializer;
use crate::vm::{HostQueries, HostQuery};

const MAX_META_SIZE: usize = ARGBUF_LEN;
// Host stack limits in our current runtime effectively allow around 64 nested
// inter-contract calls; we cap at 48 to keep safety margin. Observed mainnet
// depth is currently <= 6.
pub(crate) const MAX_CALL_DEPTH: usize = 48;
const _: () = assert!(MAX_CALL_DEPTH <= ARGBUF_LEN / CONTRACT_ID_BYTES);
pub const INIT_METHOD: &str = "init";

/// A running mutation to a state.
///
/// `Session`s are spawned using a [`VM`] instance, and can be used to [`call`]
/// contracts with to modify their state. A sequence of these calls may then be
/// [`commit`]ed to, or discarded by simply allowing the session to drop.
///
/// New contracts are to be `deploy`ed in the context of a session.
///
/// [`VM`]: crate::VM
/// [`call`]: Session::call
/// [`commit`]: Session::commit
pub struct Session {
    engine: Engine,
    inner: NonNull<SessionInner>,
    original: bool,
}

unsafe impl Send for Session {}

impl Debug for Session {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Session")
            .field("inner", &self.inner)
            .field("original", &self.original)
            .finish()
    }
}

/// A session is created by leaking an using `Box::leak` on a `SessionInner`.
/// Therefore, the memory needs to be recovered.
impl Drop for Session {
    fn drop(&mut self) {
        if self.original {
            // ensure the stack is cleared and all instances are removed and
            // reclaimed on the drop of a session.
            self.clear_call_tree_and_instances();

            // SAFETY: this is safe since we guarantee that there is no aliasing
            // when a session drops.
            unsafe {
                let _ = Box::from_raw(self.inner.as_ptr());
            }
        }
    }
}

/// A hook called before each inter-contract call.
///
/// Receives the callee contract ID, the function name, and the raw argument
/// bytes. Returns `Ok(())` to allow the call, or `Err(reason)` to reject it
/// with a descriptive message.
#[cfg(feature = "call-hook")]
pub type CallHook =
    Box<dyn Fn(&ContractId, &str, &[u8]) -> Result<(), String> + Send + Sync>;

struct SessionInner {
    current: ContractId,

    call_tree: CallTree,
    instances: BTreeMap<ContractId, Box<WrappedInstance>>,
    compiled_modules: BTreeMap<ContractId, WasmtimeModule>,
    debug: Vec<String>,
    data: SessionData,

    contract_session: ContractSession,
    host_queries: HostQueries,
    buffer: Vec<u8>,

    feeder: Option<mpsc::Sender<Vec<u8>>>,
    events: Vec<Event>,

    #[cfg(feature = "call-hook")]
    call_hook: Option<CallHook>,
}

impl Debug for SessionInner {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SessionInner")
            .field("current", &self.current)
            .field("call_tree", &self.call_tree)
            .field("instances_len", &self.instances.len())
            .field("compiled_modules_len", &self.compiled_modules.len())
            .field("debug_len", &self.debug.len())
            .field("data", &self.data)
            .field("buffer_len", &self.buffer.len())
            .field("events_len", &self.events.len())
            .finish()
    }
}

struct SessionMemoryCreator {
    inner: NonNull<SessionInner>,
}

unsafe impl Send for SessionMemoryCreator {}

unsafe impl Sync for SessionMemoryCreator {}

unsafe impl MemoryCreator for SessionMemoryCreator {
    /// This new memory is created for the contract currently at the top of the
    /// call tree.
    fn new_memory(
        &self,
        _ty: MemoryType,
        minimum: usize,
        _maximum: Option<usize>,
        _reserved_size_in_bytes: Option<usize>,
        _guard_size_in_bytes: usize,
    ) -> Result<Box<dyn LinearMemory>, String> {
        // SAFETY: wasmtime calls this synchronously during instance creation.
        // No other path concurrently accesses `SessionInner` in this callback.
        let inner = unsafe { &mut *self.inner.as_ptr() };
        let contract = inner.current;

        let contract_data =
            inner.contract_session.contract(contract).map_err(|err| {
                format!("Failed to get contract from session: {err:?}")
            })?;

        let mut memory = contract_data
            .expect("Contract data should exist at this point")
            .memory;

        if memory.is_new() {
            memory.set_current_len(minimum);
        }

        Ok(Box::new(memory))
    }
}

impl Session {
    fn inner(&self) -> &SessionInner {
        unsafe { self.inner.as_ref() }
    }

    fn inner_mut(&mut self) -> &mut SessionInner {
        unsafe { self.inner.as_mut() }
    }

    pub(crate) fn new(
        engine: Engine,
        contract_session: ContractSession,
        host_queries: HostQueries,
        data: SessionData,
    ) -> Self {
        let inner = SessionInner {
            current: ContractId::from_bytes([0; CONTRACT_ID_BYTES]),
            call_tree: CallTree::new(),
            instances: BTreeMap::new(),
            compiled_modules: BTreeMap::new(),
            debug: vec![],
            data,
            contract_session,
            host_queries,
            buffer: vec![0; PAGE_SIZE],
            feeder: None,
            events: vec![],
            #[cfg(feature = "call-hook")]
            call_hook: None,
        };

        // This implementation purposefully boxes and leaks the `SessionInner`.
        let inner = Box::leak(Box::new(inner));

        let mut session = Self {
            engine: engine.clone(),
            inner: NonNull::from(inner),
            original: true,
        };

        let mut config = engine.config().clone();
        config.with_host_memory(Arc::new(SessionMemoryCreator {
            inner: session.inner,
        }));

        session.engine = Engine::new(&config)
            .expect("Engine configuration is set at compile time");

        session
    }

    /// Clone the given session. We explicitly **do not** implement the
    /// [`Clone`] trait here, since we don't want allow the user to clone a
    /// session.
    ///
    /// This keeps cloning internal and preserves ownership/drop invariants
    /// around the shared `SessionInner` pointer.
    pub(crate) fn clone(&self) -> Self {
        Self {
            engine: self.engine.clone(),
            inner: self.inner,
            original: false,
        }
    }

    /// Return a reference to the engine used in this session.
    pub(crate) fn engine(&self) -> &Engine {
        &self.engine
    }

    /// Deploy a contract, returning its [`ContractId`] and an optional
    /// [`CallReceipt`] for the `init` function execution. The ID is computed
    /// using a `blake3` hash of the `bytecode`. Contracts using the `memory64`
    /// proposal are accepted in just the same way as 32-bit contracts, and
    /// their handling is totally transparent.
    ///
    /// Since a deployment may execute some contract initialization code, that
    /// code will be metered and executed with the given `gas_limit`. If the
    /// contract exports an `init` function, the returned receipt contains
    /// the gas spent, events emitted, and call tree produced by that call.
    ///
    /// # Errors
    /// It is possible that a collision between contract IDs occurs, even for
    /// different contract IDs. This is due to the fact that all contracts have
    /// to fit into a sparse merkle tree with `2^32` positions, and as such
    /// a 256-bit number has to be mapped into a 32-bit number.
    ///
    /// If such a collision occurs, [`PersistenceError`] will be returned.
    ///
    /// [`ContractId`]: ContractId
    /// [`CallReceipt`]: CallReceipt
    /// [`PersistenceError`]: PersistenceError
    ///
    /// # Panics
    /// If `deploy_data` does not specify an owner, this will panic.
    pub fn deploy<'a, A, R, D>(
        &mut self,
        bytecode: &[u8],
        deploy_data: D,
        gas_limit: u64,
    ) -> Result<(ContractId, Option<CallReceipt<R>>), Error>
    where
        A: 'a + for<'b> Serialize<StandardBufSerializer<'b>>,
        R: Archive,
        R::Archived: Deserialize<R, Infallible>
            + for<'b> CheckBytes<DefaultValidator<'b>>,
        D: Into<ContractData<'a, A>>,
    {
        let deploy_data = deploy_data.into();

        let mut init_arg = None;
        if let Some(arg) = deploy_data.init_arg {
            let mut sbuf = [0u8; SCRATCH_BUF_BYTES];
            let scratch = BufferScratch::new(&mut sbuf);
            let ser = BufferSerializer::new(&mut self.inner_mut().buffer[..]);
            let mut ser = CompositeSerializer::new(ser, scratch, Infallible);

            ser.serialize_value(arg)?;
            let pos = ser.pos();

            init_arg = Some(self.inner().buffer[0..pos].to_vec());
        }

        let (contract_id, receipt) = self.deploy_raw(
            deploy_data.contract_id,
            bytecode,
            init_arg,
            deploy_data
                .owner
                .expect("Owner must be specified when deploying a contract"),
            gas_limit,
        )?;

        let receipt = receipt.map(|r| r.deserialize()).transpose()?;
        Ok((contract_id, receipt))
    }

    /// Deploy a contract, returning its [`ContractId`] and an optional
    /// [`CallReceipt`] for the `init` function execution. If ID is not
    /// provided, it is computed using a `blake3` hash of the `bytecode`.
    /// Contracts using the `memory64` proposal are accepted in just the same
    /// way as 32-bit contracts, and their handling is totally transparent.
    ///
    /// Since a deployment may execute some contract initialization code, that
    /// code will be metered and executed with the given `gas_limit`. If the
    /// contract exports an `init` function, the returned receipt contains
    /// the gas spent, events emitted, and call tree produced by that call.
    ///
    /// # Errors
    /// It is possible that a collision between contract IDs occurs, even for
    /// different contract IDs. This is due to the fact that all contracts have
    /// to fit into a sparse merkle tree with `2^32` positions, and as such
    /// a 256-bit number has to be mapped into a 32-bit number.
    ///
    /// If such a collision occurs, [`PersistenceError`] will be returned.
    ///
    /// [`ContractId`]: ContractId
    /// [`CallReceipt`]: CallReceipt
    /// [`PersistenceError`]: PersistenceError
    #[allow(clippy::type_complexity)]
    pub fn deploy_raw(
        &mut self,
        contract_id: Option<ContractId>,
        bytecode: &[u8],
        init_arg: Option<Vec<u8>>,
        owner: Vec<u8>,
        gas_limit: u64,
    ) -> Result<(ContractId, Option<CallReceipt<Vec<u8>>>), Error> {
        let contract_id = contract_id.unwrap_or({
            let hash = blake3::hash(bytecode);
            ContractId::from_bytes(hash.into())
        });
        let receipt =
            self.do_deploy(contract_id, bytecode, init_arg, owner, gas_limit)?;

        Ok((contract_id, receipt))
    }

    #[allow(clippy::too_many_arguments)]
    fn do_deploy(
        &mut self,
        contract_id: ContractId,
        bytecode: &[u8],
        arg: Option<Vec<u8>>,
        owner: Vec<u8>,
        gas_limit: u64,
    ) -> Result<Option<CallReceipt<Vec<u8>>>, Error> {
        if self
            .inner_mut()
            .contract_session
            .contract_deployed(contract_id)
        {
            return Err(InitalizationError(
                "Deployed error already exists".into(),
            ));
        }

        let wrapped_contract =
            WrappedContract::new(&self.engine, bytecode, None::<&[u8]>)?;
        let contract_metadata = ContractMetadata { contract_id, owner };
        let metadata_bytes = Self::serialize_data(&contract_metadata)?;

        self.inner_mut()
            .contract_session
            .deploy(
                contract_id,
                bytecode,
                wrapped_contract.as_bytes(),
                contract_metadata,
                metadata_bytes.as_slice(),
            )
            .map_err(|err| PersistenceError(Arc::new(err)))?;
        self.inner_mut().compiled_modules.remove(&contract_id);

        let instantiate = || {
            self.create_instance(contract_id)?;
            let has_init = self
                .instance(&contract_id)
                .expect("instance should exist")
                .is_function_exported(INIT_METHOD);

            if has_init {
                // If no argument was provided, we call the init method anyway,
                // but with an empty argument. The alternative is to panic, but
                // that assumes that the caller of `deploy` knows that the
                // contract has an init method in the first place, which might
                // not be the case, such as when ingesting untrusted bytecode.
                let arg = arg.unwrap_or_default();
                let (data, gas_spent, call_tree) =
                    self.call_inner(contract_id, INIT_METHOD, arg, gas_limit)?;
                let events = mem::take(&mut self.inner_mut().events);
                return Ok(Some(CallReceipt {
                    gas_limit,
                    gas_spent,
                    events,
                    call_tree,
                    data,
                }));
            }

            Ok(None)
        };

        instantiate().inspect_err(|_| {
            self.inner_mut()
                .contract_session
                .remove_contract(&contract_id);
            self.inner_mut().compiled_modules.remove(&contract_id);
        })
    }

    /// Execute a call on the current state of this session.
    ///
    /// Calls are atomic, meaning that on failure their execution doesn't modify
    /// the state. They are also metered, and will execute with the given
    /// `gas_limit`. This value should never be 0.
    ///
    /// # Errors
    /// The call may error during execution for a wide array of reasons, the
    /// most common ones being running against the gas limit and a contract
    /// panic. Calling the 'init' method is not allowed except for when called
    /// from the deploy method.
    pub fn call<A, R>(
        &mut self,
        contract: ContractId,
        fn_name: &str,
        fn_arg: &A,
        gas_limit: u64,
    ) -> Result<CallReceipt<R>, Error>
    where
        A: for<'b> Serialize<StandardBufSerializer<'b>>,
        A::Archived: for<'b> CheckBytes<DefaultValidator<'b>>,
        R: Archive,
        R::Archived: Deserialize<R, Infallible>
            + for<'b> CheckBytes<DefaultValidator<'b>>,
    {
        if fn_name == INIT_METHOD {
            return Err(InitalizationError("init call not allowed".into()));
        }

        let mut sbuf = [0u8; SCRATCH_BUF_BYTES];
        let scratch = BufferScratch::new(&mut sbuf);
        let ser = BufferSerializer::new(&mut self.inner_mut().buffer[..]);
        let mut ser = CompositeSerializer::new(ser, scratch, Infallible);

        ser.serialize_value(fn_arg)?;
        let pos = ser.pos();

        let receipt = self.call_raw(
            contract,
            fn_name,
            self.inner().buffer[..pos].to_vec(),
            gas_limit,
        )?;

        receipt.deserialize()
    }

    /// Execute a raw call on the current state of this session.
    ///
    /// Raw calls do not specify the type of the argument or of the return. The
    /// caller is responsible for serializing the argument as the target
    /// `contract` expects.
    ///
    /// For more information about calls see [`call`].
    ///
    /// [`call`]: Session::call
    pub fn call_raw<V: Into<Vec<u8>>>(
        &mut self,
        contract: ContractId,
        fn_name: &str,
        fn_arg: V,
        gas_limit: u64,
    ) -> Result<CallReceipt<Vec<u8>>, Error> {
        if fn_name == INIT_METHOD {
            return Err(InitalizationError("init call not allowed".into()));
        }

        let (data, gas_spent, call_tree) =
            self.call_inner(contract, fn_name, fn_arg.into(), gas_limit)?;
        let events = mem::take(&mut self.inner_mut().events);

        Ok(CallReceipt {
            gas_limit,
            gas_spent,
            events,
            call_tree,
            data,
        })
    }

    /// 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.
    ///
    /// If the `owner` of a contract is not set, it will be set to the owner of
    /// the contract being replaced. If it is set, then it will be used as the
    /// new owner.
    ///
    /// # 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.
    ///
    /// # Panics
    /// If the owner of the new contract is not set in `deploy_data`, and the
    /// contract being replaced does not exist, this will panic.
    pub fn migrate<'a, A, D, F>(
        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>>,
        F: FnOnce(ContractId, &mut Session) -> Result<(), Error>,
    {
        let mut new_contract_data = deploy_data.into();

        // If the contract being replaced exists, and the caller did not specify
        // an owner, set the owner to the owner of the contract being replaced.
        if let Some(old_contract_data) = self
            .inner_mut()
            .contract_session
            .contract(contract)
            .map_err(|err| PersistenceError(Arc::new(err)))?
        {
            if new_contract_data.owner.is_none() {
                new_contract_data.owner =
                    Some(old_contract_data.metadata.data().owner.clone());
            }
        }

        let (new_contract, _init_receipt) = self.deploy::<_, (), _>(
            bytecode,
            new_contract_data,
            deploy_gas_limit,
        )?;

        closure(new_contract, &mut self)?;

        self.inner_mut()
            .contract_session
            .replace(contract, new_contract)?;
        self.inner_mut().compiled_modules.remove(&contract);
        self.inner_mut().compiled_modules.remove(&new_contract);

        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
    /// amounts of data to the host via the channel included in this call.
    ///
    /// These calls should be performed with a large amount of gas, since the
    /// contracts may spend quite a large amount in an effort to report data.
    pub fn feeder_call<A, R>(
        &mut self,
        contract: ContractId,
        fn_name: &str,
        fn_arg: &A,
        gas_limit: u64,
        feeder: mpsc::Sender<Vec<u8>>,
    ) -> Result<CallReceipt<R>, Error>
    where
        A: for<'b> Serialize<StandardBufSerializer<'b>>,
        A::Archived: for<'b> CheckBytes<DefaultValidator<'b>>,
        R: Archive,
        R::Archived: Deserialize<R, Infallible>
            + for<'b> CheckBytes<DefaultValidator<'b>>,
    {
        self.inner_mut().feeder = Some(feeder);
        let r = self.call(contract, fn_name, fn_arg, gas_limit);
        self.inner_mut().feeder = None;
        r
    }

    /// Execute a raw *feeder* call on the current state of this session.
    ///
    /// See [`feeder_call`] and [`call_raw`] for more information of this type
    /// of call.
    ///
    /// [`feeder_call`]: [`Session::feeder_call`]
    /// [`call_raw`]: [`Session::call_raw`]
    pub fn feeder_call_raw<V: Into<Vec<u8>>>(
        &mut self,
        contract: ContractId,
        fn_name: &str,
        fn_arg: V,
        gas_limit: u64,
        feeder: mpsc::Sender<Vec<u8>>,
    ) -> Result<CallReceipt<Vec<u8>>, Error> {
        self.inner_mut().feeder = Some(feeder);
        let r = self.call_raw(contract, fn_name, fn_arg, gas_limit);
        self.inner_mut().feeder = None;
        r
    }

    /// Returns the current length of the memory of the given contract.
    ///
    /// 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_mut()
            .contract_session
            .contract(contract_id)
            .map_err(|err| PersistenceError(Arc::new(err)))?
            .map(|data| data.memory.current_len()))
    }

    pub(crate) fn instance(
        &mut self,
        contract_id: &ContractId,
    ) -> Option<&mut WrappedInstance> {
        self.inner_mut()
            .instances
            .get_mut(contract_id)
            .map(Box::as_mut)
    }

    fn clear_call_tree_and_instances(&mut self) {
        self.inner_mut().call_tree.clear();
        self.inner_mut().instances.clear();
    }

    /// Return the state root of the current state of the session.
    ///
    /// The state root is the root of a merkle tree whose leaves are the hashes
    /// of the state of of each contract, ordered by their contract ID.
    ///
    /// It also doubles as the ID of a commit - the commit root.
    pub fn root(&self) -> [u8; 32] {
        self.inner().contract_session.root().into()
    }

    /// Returns an iterator over the pages (and their indices) of a contract's
    /// memory, together with a proof of their inclusion in the state.
    ///
    /// The proof is a Merkle inclusion proof, and the caller is able to verify
    /// it by using [`verify`], and matching the root with the one returned by
    /// [`root`].
    ///
    /// [`verify`]: PageOpening::verify
    /// [`root`]: Session::root
    pub fn memory_pages(
        &self,
        contract: ContractId,
    ) -> Option<impl Iterator<Item = (usize, &[u8], PageOpening)>> {
        self.inner().contract_session.memory_pages(contract)
    }

    pub(crate) fn push_event(&mut self, event: Event) {
        self.inner_mut().events.push(event);
    }

    pub(crate) fn event_checkpoint(&self) -> usize {
        self.inner().events.len()
    }

    pub(crate) fn revert_events_from(&mut self, checkpoint: usize) {
        for event in self.inner_mut().events.iter_mut().skip(checkpoint) {
            event.reverted = true;
        }
    }

    pub(crate) fn push_feed(&mut self, data: Vec<u8>) -> Result<(), Error> {
        let feed = self.inner().feeder.as_ref().ok_or(Error::MissingFeed)?;
        feed.send(data).map_err(Error::FeedPulled)
    }

    fn new_instance(
        &mut self,
        contract_id: ContractId,
    ) -> Result<WrappedInstance, Error> {
        let store_data = self
            .inner_mut()
            .contract_session
            .contract(contract_id)
            .map_err(|err| PersistenceError(Arc::new(err)))?
            .ok_or(Error::ContractDoesNotExist(contract_id))?;

        let module = if let Some(module) =
            self.inner().compiled_modules.get(&contract_id)
        {
            module.clone()
        } else {
            let module = unsafe {
                WasmtimeModule::deserialize(
                    &self.engine,
                    store_data.module.serialize(),
                )?
            };
            self.inner_mut()
                .compiled_modules
                .insert(contract_id, module.clone());
            module
        };

        self.inner_mut().current = contract_id;

        let instance = WrappedInstance::new(
            self.clone(),
            contract_id,
            &module,
            store_data.memory,
        )?;

        Ok(instance)
    }

    pub(crate) fn host_query_arc(
        &self,
        name: &str,
    ) -> Option<Arc<dyn HostQuery>> {
        self.inner().host_queries.get_arc(name)
    }

    pub(crate) fn nth_from_top(&self, n: usize) -> Option<CallTreeElem> {
        self.inner().call_tree.nth_parent(n)
    }

    pub(crate) fn call_ids(&self) -> Vec<&ContractId> {
        self.inner().call_tree.call_ids()
    }

    /// Creates a new instance of the given contract, returning its memory
    /// length.
    fn create_instance(
        &mut self,
        contract: ContractId,
    ) -> Result<usize, Error> {
        let instance = self.new_instance(contract)?;
        if self.inner().instances.contains_key(&contract) {
            panic!("Contract already in the stack: {contract:?}");
        }

        let mem_len = instance.mem_len();

        self.inner_mut()
            .instances
            .insert(contract, Box::new(instance));
        Ok(mem_len)
    }

    pub(crate) fn push_callstack(
        &mut self,
        contract_id: ContractId,
        limit: u64,
    ) -> Result<CallTreeElem, Error> {
        let current_depth = self.inner().call_tree.depth();
        if current_depth >= MAX_CALL_DEPTH {
            return Err(Error::SessionError(
                format!("Maximum call depth exceeded ({MAX_CALL_DEPTH})")
                    .into(),
            ));
        }

        let mem_len =
            if let Some(instance) = self.inner().instances.get(&contract_id) {
                instance.mem_len()
            } else {
                self.create_instance(contract_id)?
            };

        self.inner_mut().call_tree.push(CallTreeElem {
            contract_id,
            limit,
            spent: 0,
            mem_len,
        });

        Ok(self
            .inner()
            .call_tree
            .nth_parent(0)
            .expect("We just pushed an element to the stack"))
    }

    pub(crate) fn move_up_call_tree(&mut self, spent: u64) {
        self.inner_mut().call_tree.move_up(spent);
    }

    pub(crate) fn move_up_prune_call_tree(&mut self) {
        self.inner_mut().call_tree.move_up_prune();
    }

    pub(crate) fn revert_callstack(&mut self) -> Result<(), std::io::Error> {
        let call_tree: Vec<_> =
            self.inner().call_tree.iter().copied().collect();
        for elem in call_tree {
            let instance = self
                .instance(&elem.contract_id)
                .expect("instance should exist");
            instance.revert()?;
            instance.set_len(elem.mem_len);
        }

        Ok(())
    }

    /// Commits the given session to disk, consuming the session and returning
    /// its state root.
    pub fn commit(mut self) -> Result<[u8; 32], Error> {
        self.inner_mut()
            .contract_session
            .commit()
            .map(Into::into)
            .map_err(|err| PersistenceError(Arc::new(err)))
    }

    #[cfg(feature = "debug")]
    pub(crate) fn register_debug<M: Into<String>>(&mut self, msg: M) {
        self.inner_mut().debug.push(msg.into());
    }

    pub fn with_debug<C, R>(&self, c: C) -> R
    where
        C: FnOnce(&[String]) -> R,
    {
        c(&self.inner().debug)
    }

    /// Returns the value of a metadata item.
    pub fn meta(&self, name: &str) -> Option<Vec<u8>> {
        self.inner().data.get(name)
    }

    /// Set the value of a metadata item.
    ///
    /// Returns the previous value of the metadata item.
    pub fn set_meta<S, V>(
        &mut self,
        name: S,
        value: V,
    ) -> Result<Option<Vec<u8>>, Error>
    where
        S: Into<Cow<'static, str>>,
        V: for<'a> Serialize<StandardBufSerializer<'a>>,
    {
        let data = Self::serialize_data(&value)?;
        Ok(self.inner_mut().data.set(name, data))
    }

    /// Remove a metadata item.
    ///
    /// Returns the value of the removed item (if any).
    pub fn remove_meta<S>(&mut self, name: S) -> Option<Vec<u8>>
    where
        S: Into<Cow<'static, str>>,
    {
        self.inner_mut().data.remove(name)
    }

    pub fn serialize_data<V>(value: &V) -> Result<Vec<u8>, Error>
    where
        V: for<'a> Serialize<StandardBufSerializer<'a>>,
    {
        let mut buf = [0u8; MAX_META_SIZE];
        let mut sbuf = [0u8; SCRATCH_BUF_BYTES];

        let ser = BufferSerializer::new(&mut buf[..]);
        let scratch = BufferScratch::new(&mut sbuf);

        let mut serializer =
            StandardBufSerializer::new(ser, scratch, Infallible);
        serializer.serialize_value(value)?;

        let pos = serializer.pos();

        Ok(buf[..pos].to_vec())
    }

    fn call_inner(
        &mut self,
        contract: ContractId,
        fname: &str,
        fdata: Vec<u8>,
        limit: u64,
    ) -> Result<(Vec<u8>, u64, CallTree), Error> {
        let event_checkpoint = self.event_checkpoint();
        let stack_element = self.push_callstack(contract, limit)?;
        {
            let instance = self
                .instance(&stack_element.contract_id)
                .expect("instance should exist");
            instance
                .snap()
                .map_err(|err| Error::MemorySnapshotFailure {
                    reason: None,
                    io: Arc::new(err),
                })?;
        }

        let ret_len = {
            let instance = self
                .instance(&stack_element.contract_id)
                .expect("instance should exist");
            let arg_len = instance.write_bytes_to_arg_buffer(&fdata)?;
            instance
                .call(fname, arg_len, limit)
                .map_err(Error::normalize)
        };

        let ret_len = match ret_len {
            Ok(ret_len) => ret_len,
            Err(err) => {
                let err = if let Err(io_err) = self.revert_callstack() {
                    Error::MemorySnapshotFailure {
                        reason: Some(Arc::new(err)),
                        io: Arc::new(io_err),
                    }
                } else {
                    err
                };
                self.revert_events_from(event_checkpoint);
                self.move_up_prune_call_tree();
                self.clear_call_tree_and_instances();
                return Err(err);
            }
        };

        let (ret, spent) = {
            let instance = self
                .instance(&stack_element.contract_id)
                .expect("instance should exist");
            let ret = instance.read_bytes_from_arg_buffer(ret_len as u32);
            let spent = limit - instance.get_remaining_gas();
            (ret, spent)
        };

        let call_tree: Vec<_> =
            self.inner().call_tree.iter().copied().collect();
        for elem in call_tree {
            let instance = self
                .instance(&elem.contract_id)
                .expect("instance should exist");
            instance
                .apply()
                .map_err(|err| Error::MemorySnapshotFailure {
                    reason: None,
                    io: Arc::new(err),
                })?;
        }

        let mut call_tree = CallTree::new();
        mem::swap(&mut self.inner_mut().call_tree, &mut call_tree);
        call_tree.update_spent(spent);

        self.clear_call_tree_and_instances();

        Ok((ret, spent, call_tree))
    }

    pub fn contract_metadata(
        &mut self,
        contract_id: &ContractId,
    ) -> Option<&ContractMetadata> {
        self.inner_mut()
            .contract_session
            .contract_metadata(contract_id)
    }

    /// Set a hook that is called before each inter-contract call.
    ///
    /// The hook receives the callee contract ID, the function name, and the
    /// raw argument bytes.
    #[cfg(feature = "call-hook")]
    pub fn set_call_hook(&mut self, hook: CallHook) -> Option<CallHook> {
        self.inner_mut().call_hook.replace(hook)
    }

    /// Remove the call hook, if one is set.
    #[cfg(feature = "call-hook")]
    pub fn clear_call_hook(&mut self) -> Option<CallHook> {
        self.inner_mut().call_hook.take()
    }

    /// Run the call hook, if one is set.
    ///
    /// Returns `Ok(())` if the call is allowed (or no hook is set), or
    /// `Err(reason)` if the hook rejects.
    #[cfg(feature = "call-hook")]
    pub(crate) fn call_hook(
        &self,
        callee: &ContractId,
        fn_name: &str,
        arg: &[u8],
    ) -> Result<(), String> {
        if let Some(hook) = &self.inner().call_hook {
            hook(callee, fn_name, arg)
        } else {
            Ok(())
        }
    }
}

/// The receipt given for a call execution using one of either [`call`] or
/// [`call_raw`]. This receipt is also returned within a tuple when
/// a Contract is deployed and its `init` method is executed.
///
/// [`call`]: [`Session::call`]
/// [`call_raw`]: [`Session::call_raw`]
#[derive(Debug)]
pub struct CallReceipt<T> {
    /// The amount of gas spent in the execution of the call.
    pub gas_spent: u64,
    /// The limit used in during this execution.
    pub gas_limit: u64,

    /// The events emitted during the execution of the call.
    pub events: Vec<Event>,
    /// The call tree produced during the execution.
    pub call_tree: CallTree,

    /// The data returned by the called contract.
    pub data: T,
}

impl CallReceipt<Vec<u8>> {
    /// Deserializes a `CallReceipt<Vec<u8>>` into a `CallReceipt<T>` using
    /// `rkyv`.
    fn deserialize<T>(self) -> Result<CallReceipt<T>, Error>
    where
        T: Archive,
        T::Archived: Deserialize<T, Infallible>
            + for<'b> CheckBytes<DefaultValidator<'b>>,
    {
        let ta = check_archived_root::<T>(&self.data[..])?;
        let data = ta.deserialize(&mut Infallible)?;

        Ok(CallReceipt {
            gas_spent: self.gas_spent,
            gas_limit: self.gas_limit,
            events: self.events,
            call_tree: self.call_tree,
            data,
        })
    }
}

#[derive(Debug, Default)]
pub struct SessionData {
    data: BTreeMap<Cow<'static, str>, Vec<u8>>,
    pub base: Option<[u8; 32]>,
    excluded_host_queries: BTreeSet<String>,
}

impl SessionData {
    pub fn builder() -> SessionDataBuilder {
        SessionDataBuilder {
            data: BTreeMap::new(),
            base: None,
            excluded_host_queries: BTreeSet::new(),
        }
    }

    fn get(&self, name: &str) -> Option<Vec<u8>> {
        self.data.get(name).cloned()
    }

    fn set<S>(&mut self, name: S, data: Vec<u8>) -> Option<Vec<u8>>
    where
        S: Into<Cow<'static, str>>,
    {
        self.data.insert(name.into(), data)
    }

    fn remove<S>(&mut self, name: S) -> Option<Vec<u8>>
    where
        S: Into<Cow<'static, str>>,
    {
        self.data.remove(&name.into())
    }

    pub fn excluded_host_queries(&self) -> Iter<String> {
        self.excluded_host_queries.iter()
    }
}

impl From<SessionDataBuilder> for SessionData {
    fn from(builder: SessionDataBuilder) -> Self {
        builder.build()
    }
}

pub struct SessionDataBuilder {
    data: BTreeMap<Cow<'static, str>, Vec<u8>>,
    base: Option<[u8; 32]>,
    excluded_host_queries: BTreeSet<String>,
}

impl SessionDataBuilder {
    pub fn insert<S, V>(mut self, name: S, value: V) -> Result<Self, Error>
    where
        S: Into<Cow<'static, str>>,
        V: for<'a> Serialize<StandardBufSerializer<'a>>,
    {
        let data = Session::serialize_data(&value)?;
        self.data.insert(name.into(), data);
        Ok(self)
    }

    pub fn base(mut self, base: [u8; 32]) -> Self {
        self.base = Some(base);
        self
    }

    pub fn exclude_hq(mut self, name: String) -> Self {
        self.excluded_host_queries.insert(name);
        self
    }

    fn build(&self) -> SessionData {
        SessionData {
            data: self.data.clone(),
            base: self.base,
            excluded_host_queries: self.excluded_host_queries.clone(),
        }
    }
}