piecrust: add spent field to CallTreeElem

piecrust/CHANGELOG.md

@@ -9,6 +9,7 @@ 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]

piecrust/src/call_tree.rs

@@ -14,12 +14,13 @@ use piecrust_uplink::ContractId;
 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 CallTreeInner>);
+pub struct CallTree(Option<*mut CallTreeNode>);
 
 impl CallTree {
     /// Creates a new empty call tree, starting with at the given contract.
@@ -32,18 +33,20 @@ impl CallTree {
     /// 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(CallTreeInner::new(elem)),
+            None => self.0 = Some(CallTreeNode::new(elem)),
             Some(inner) => unsafe {
-                let node = CallTreeInner::with_prev(elem, inner);
+                let node = CallTreeNode::with_prev(elem, inner);
                 (*inner).children.push(node);
                 self.0 = Some(node)
             },
         }
     }
 
-    /// Moves to the previous node, returning the current element.
-    pub(crate) fn move_up(&mut self) -> Option<CallTreeElem> {
+    /// Moves to the previous 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 prev = (*inner).prev;
@@ -56,8 +59,8 @@ impl CallTree {
         })
     }
 
-    /// Moves to the previous node, returning the current element, while
-    /// clearing the tree under it.
+    /// Moves to the previous 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;
@@ -73,6 +76,17 @@ impl CallTree {
         })
     }
 
+    /// 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 previous element from the counting from the current
     /// node.
     ///
@@ -91,28 +105,59 @@ impl CallTree {
 
     /// Clears the call tree of all elements.
     pub(crate) fn clear(&mut self) {
-        while self.0.is_some() {
-            self.move_up();
+        unsafe {
+            if let Some(inner) = self.0 {
+                let mut root = inner;
+
+                while let Some(prev) = (*root).prev {
+                    root = prev;
+                }
+
+                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 {
-            node: self.0,
+            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 free_tree(root: *mut CallTreeInner) {
+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();
@@ -123,13 +168,13 @@ unsafe fn free_tree(root: *mut CallTreeInner) {
     }
 }
 
-struct CallTreeInner {
+struct CallTreeNode {
     elem: CallTreeElem,
     children: Vec<*mut Self>,
     prev: Option<*mut Self>,
 }
 
-impl CallTreeInner {
+impl CallTreeNode {
     fn new(elem: CallTreeElem) -> *mut Self {
         Box::leak(Box::new(Self {
             elem,
@@ -152,7 +197,7 @@ impl CallTreeInner {
 /// It starts at the righmost node and proceeds leftward towards its siblings,
 /// up toward the root.
 struct CallTreeIter<'a> {
-    node: Option<*mut CallTreeInner>,
+    tree: Option<Subtree>,
     _marker: PhantomData<&'a ()>,
 }
 
@@ -161,14 +206,21 @@ impl<'a> Iterator for CallTreeIter<'a> {
 
     fn next(&mut self) -> Option<Self::Item> {
         unsafe {
-            let node = match self.node {
-                Some(node) => node,
-                None => return None,
-            };
+            let tree = self.tree.as_mut()?;
 
+            let node = tree.node;
             let elem = &(*node).elem;
 
-            self.node = (*node).prev.map(|prev_node| {
+            if node == tree.root {
+                self.tree = None;
+                return Some(elem);
+            }
+
+            let prev_node = (*node).prev.expect(
+                "There should always be a prev in a subtree before the root",
+            );
+
+            tree.node = {
                 let node_index = (*prev_node)
                     .children
                     .iter()
@@ -188,7 +240,7 @@ impl<'a> Iterator for CallTreeIter<'a> {
 
                     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.move_up_call_tree();
+            env.move_up_call_tree(callee_spent);
             instance.set_remaining_points(caller_remaining - callee_spent);
             ret_len
         }

piecrust/src/session.rs

@@ -522,6 +522,7 @@ impl Session {
                 self.inner.call_tree.push(CallTreeElem {
                     contract_id,
                     limit,
+                    spent: 0,
                     mem_len: instance.mem_len(),
                 });
             }
@@ -530,6 +531,7 @@ impl Session {
                 self.inner.call_tree.push(CallTreeElem {
                     contract_id,
                     limit,
+                    spent: 0,
                     mem_len,
                 });
             }
@@ -542,8 +544,8 @@ impl Session {
             .expect("We just pushed an element to the stack"))
     }
 
-    pub(crate) fn move_up_call_tree(&mut self) {
-        if let Some(element) = self.inner.call_tree.move_up() {
+    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);
         }
     }
@@ -666,6 +668,7 @@ impl Session {
 
         let mut call_tree = CallTree::new();
         mem::swap(&mut self.inner.call_tree, &mut call_tree);
+        call_tree.update_spent(spent);
 
         Ok((ret, spent, call_tree))
     }