Receive call tree after execution

Cargo.toml

@@ -4,3 +4,4 @@ members = [
     "piecrust-uplink",
     "piecrust",
 ]
+resolver = "2"

contracts/Cargo.toml

@@ -21,3 +21,4 @@ members = [
     "stack",
     "vector",
 ]
+resolver = "2"

piecrust/CHANGELOG.md

@@ -9,14 +9,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add `spent` field to `CallTreeElem` [#206]
+- Add `call_tree` to `CallReceipt` [#206]
+- Expose `CallTree` and `CallTreeElem` in the public API [#206]
+- Add `CallTreeIter` to improve iteration over call tree [#206]
 - Add `panic` import implementation [#271]
 - Add `Error::ContractPanic` variant [#271]
 
 ### Changed
 
+- Rename `StackElement` to `CallTreeElem` [#206]
 - Allow for multiple initializations on a new memory [#271]
 - Downcast `Error::RuntimeError` on each call boundary [#271]
 
+### Removed
+
+- Remove `CallStack` in favor of `CallTree` [#206]
+
 ## [0.10.0] - 2023-09-13
 
 ### Added
@@ -250,6 +259,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [#216]: https://github.com/dusk-network/piecrust/issues/216
 [#213]: https://github.com/dusk-network/piecrust/issues/213
 [#207]: https://github.com/dusk-network/piecrust/issues/207
+[#206]: https://github.com/dusk-network/piecrust/issues/206
 [#202]: https://github.com/dusk-network/piecrust/issues/202
 [#201]: https://github.com/dusk-network/piecrust/issues/201
 [#181]: https://github.com/dusk-network/piecrust/issues/181

piecrust/src/call_tree.rs

@@ -0,0 +1,249 @@
+// 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::marker::PhantomData;
+use std::mem;
+
+use piecrust_uplink::ContractId;
+
+/// An element of the call tree.
+#[derive(Debug, Clone, Copy)]
+pub struct CallTreeElem {
+    pub contract_id: ContractId,
+    pub limit: u64,
+    pub spent: u64,
+    pub mem_len: usize,
+}
+
+/// The tree of contract calls.
+#[derive(Debug, Default)]
+pub struct CallTree(Option<*mut CallTreeNode>);
+
+impl CallTree {
+    /// Creates a new empty call tree, starting with the given contract.
+    pub(crate) const fn new() -> Self {
+        Self(None)
+    }
+
+    /// Push an element to the call tree.
+    ///
+    /// This pushes a new child to the current node, and advances to it.
+    pub(crate) fn push(&mut self, elem: CallTreeElem) {
+        match self.0 {
+            None => self.0 = Some(CallTreeNode::new(elem)),
+            Some(inner) => unsafe {
+                let node = CallTreeNode::with_parent(elem, inner);
+                (*inner).children.push(node);
+                self.0 = Some(node)
+            },
+        }
+    }
+
+    /// Moves to the parent node and set the gas spent of the current element,
+    /// returning it.
+    pub(crate) fn move_up(&mut self, spent: u64) -> Option<CallTreeElem> {
+        self.0.map(|inner| unsafe {
+            (*inner).elem.spent = spent;
+            let elem = (*inner).elem;
+
+            let parent = (*inner).parent;
+            if parent.is_none() {
+                free_tree(inner);
+            }
+            self.0 = parent;
+
+            elem
+        })
+    }
+
+    /// Moves to the parent node, clearing the tree under it, and returns the
+    /// current element.
+    pub(crate) fn move_up_prune(&mut self) -> Option<CallTreeElem> {
+        self.0.map(|inner| unsafe {
+            let elem = (*inner).elem;
+
+            let parent = (*inner).parent;
+            if let Some(parent) = parent {
+                (*parent).children.pop();
+            }
+            free_tree(inner);
+            self.0 = parent;
+
+            elem
+        })
+    }
+
+    /// Give the current node the amount spent and recusively update amount
+    /// spent to accurately reflect what each node spent during each call.
+    pub(crate) fn update_spent(&mut self, spent: u64) {
+        if let Some(inner) = self.0 {
+            unsafe {
+                (*inner).elem.spent = spent;
+                update_spent(inner);
+            }
+        }
+    }
+
+    /// Returns the `n`th parent element counting from the current node. The
+    /// zeroth parent element is the current node.
+    pub(crate) fn nth_parent(&self, n: usize) -> Option<CallTreeElem> {
+        let mut current = self.0;
+
+        let mut i = 0;
+        while i < n {
+            current = current.and_then(|inner| unsafe { (*inner).parent });
+            i += 1;
+        }
+
+        current.map(|inner| unsafe { (*inner).elem })
+    }
+
+    /// Clears the call tree of all elements.
+    pub(crate) fn clear(&mut self) {
+        unsafe {
+            if let Some(inner) = self.0 {
+                let mut root = inner;
+
+                while let Some(parent) = (*root).parent {
+                    root = parent;
+                }
+
+                self.0 = None;
+                free_tree(root);
+            }
+        }
+    }
+
+    /// Returns an iterator over the call tree, starting from the rightmost
+    /// leaf, and proceeding to the top of the current position of the tree.
+    pub fn iter(&self) -> impl Iterator<Item = &CallTreeElem> {
+        CallTreeIter {
+            tree: self.0.map(|root| unsafe {
+                let mut node = root;
+
+                while !(*node).children.is_empty() {
+                    let last_index = (*node).children.len() - 1;
+                    node = (*node).children[last_index];
+                }
+
+                Subtree { root, node }
+            }),
+            _marker: PhantomData,
+        }
+    }
+}
+
+struct Subtree {
+    root: *mut CallTreeNode,
+    node: *mut CallTreeNode,
+}
+
+impl Drop for CallTree {
+    fn drop(&mut self) {
+        self.clear()
+    }
+}
+
+unsafe fn update_spent(node: *mut CallTreeNode) {
+    let node = &mut *node;
+    node.children.iter_mut().for_each(|&mut child| unsafe {
+        node.elem.spent -= (*child).elem.spent;
+        update_spent(child);
+    });
+}
+
+unsafe fn free_tree(root: *mut CallTreeNode) {
+    let mut node = Box::from_raw(root);
+
+    let mut children = Vec::new();
+    mem::swap(&mut node.children, &mut children);
+
+    for child in children {
+        free_tree(child);
+    }
+}
+
+struct CallTreeNode {
+    elem: CallTreeElem,
+    children: Vec<*mut Self>,
+    parent: Option<*mut Self>,
+}
+
+impl CallTreeNode {
+    fn new(elem: CallTreeElem) -> *mut Self {
+        Box::leak(Box::new(Self {
+            elem,
+            children: Vec::new(),
+            parent: None,
+        }))
+    }
+
+    fn with_parent(elem: CallTreeElem, parent: *mut Self) -> *mut Self {
+        Box::leak(Box::new(Self {
+            elem,
+            children: Vec::new(),
+            parent: Some(parent),
+        }))
+    }
+}
+
+/// An iterator over a [`CallTree`].
+///
+/// It starts at the righmost node and proceeds leftward towards its siblings,
+/// up toward the root.
+struct CallTreeIter<'a> {
+    tree: Option<Subtree>,
+    _marker: PhantomData<&'a ()>,
+}
+
+impl<'a> Iterator for CallTreeIter<'a> {
+    type Item = &'a CallTreeElem;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // SAFETY: This is safe since we guarantee that the tree exists between
+        // the root and the current node. This is done by ensuring the iterator
+        // can only exist during the lifetime of the tree.
+        unsafe {
+            let tree = self.tree.as_mut()?;
+
+            let node = tree.node;
+            let elem = &(*node).elem;
+
+            if node == tree.root {
+                self.tree = None;
+                return Some(elem);
+            }
+
+            let parent = (*node).parent.expect(
+                "There should always be a parent in a subtree before the root",
+            );
+
+            tree.node = {
+                let node_index = (*parent)
+                    .children
+                    .iter()
+                    .position(|&n| n == node)
+                    .expect("The child must be the its parent's child");
+
+                if node_index == 0 {
+                    parent
+                } else {
+                    let sibling_index = node_index - 1;
+                    let mut next_node = (*parent).children[sibling_index];
+
+                    while !(*next_node).children.is_empty() {
+                        let last_index = (*next_node).children.len() - 1;
+                        next_node = (*next_node).children[last_index];
+                    }
+
+                    next_node
+                }
+            };
+
+            Some(elem)
+        }
+    }
+}

piecrust/src/imports.rs

@@ -182,7 +182,7 @@ fn c(
 
     let ret = match instance.with_memory_mut(with_memory) {
         Ok((ret_len, callee_spent)) => {
-            env.pop_callstack();
+            env.move_up_call_tree(callee_spent);
             instance.set_remaining_points(caller_remaining - callee_spent);
             ret_len
         }
@@ -193,7 +193,7 @@ fn c(
                     io: Arc::new(io_err),
                 };
             }
-            env.pop_callstack_prune();
+            env.move_up_prune_call_tree();
             instance.set_remaining_points(caller_remaining - callee_limit);
 
             ContractError::from(err).into()

piecrust/src/lib.rs

@@ -105,6 +105,7 @@
 
 #[macro_use]
 mod bytecode_macro;
+mod call_tree;
 mod contract;
 mod error;
 mod imports;
@@ -114,6 +115,7 @@ mod store;
 mod types;
 mod vm;
 
+pub use call_tree::{CallTree, CallTreeElem};
 pub use contract::{ContractData, ContractDataBuilder};
 pub use error::Error;
 pub use session::{CallReceipt, Session, SessionData};