piecrust: remove CallStack in favor of CallTree

piecrust/CHANGELOG.md

@@ -19,6 +19,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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

piecrust/src/session.rs

@@ -32,7 +32,7 @@ use crate::vm::HostQueries;
 use crate::Error;
 use crate::Error::{InitalizationError, PersistenceError};
 
-use call_stack::{CallStack, CallTreeElem};
+use call_stack::{CallTree, CallTreeElem};
 
 const MAX_META_SIZE: usize = ARGBUF_LEN;
 pub const INIT_METHOD: &str = "init";
@@ -87,7 +87,7 @@ impl Drop for Session {
 
 #[derive(Debug)]
 struct SessionInner {
-    call_stack: CallStack,
+    call_stack: CallTree,
     instance_map: BTreeMap<ContractId, (*mut WrappedInstance, u64)>,
     debug: Vec<String>,
     data: SessionData,
@@ -107,7 +107,7 @@ impl Session {
         data: SessionData,
     ) -> Self {
         Self::from(SessionInner {
-            call_stack: CallStack::new(),
+            call_stack: CallTree::new(),
             instance_map: BTreeMap::new(),
             debug: vec![],
             data,
@@ -489,7 +489,7 @@ impl Session {
     }
 
     pub(crate) fn nth_from_top(&self, n: usize) -> Option<CallTreeElem> {
-        self.inner.call_stack.nth_from_top(n)
+        self.inner.call_stack.nth_up(n)
     }
 
     /// Creates a new instance of the given contract, returning its memory
@@ -541,24 +541,24 @@ impl Session {
         Ok(self
             .inner
             .call_stack
-            .nth_from_top(0)
+            .nth_up(0)
             .expect("We just pushed an element to the stack"))
     }
 
     pub(crate) fn pop_callstack(&mut self) {
-        if let Some(element) = self.inner.call_stack.pop() {
+        if let Some(element) = self.inner.call_stack.move_up() {
             self.update_instance_count(element.contract_id, false);
         }
     }
 
     pub(crate) fn pop_callstack_prune(&mut self) {
-        if let Some(element) = self.inner.call_stack.pop_prune() {
+        if let Some(element) = self.inner.call_stack.move_up_prune() {
             self.update_instance_count(element.contract_id, false);
         }
     }
 
     pub(crate) fn revert_callstack(&mut self) -> Result<(), std::io::Error> {
-        for elem in self.inner.call_stack.iter_tree() {
+        for elem in self.inner.call_stack.iter() {
             let instance = self
                 .instance(&elem.contract_id)
                 .expect("instance should exist");
@@ -655,7 +655,7 @@ impl Session {
                 .get_remaining_points()
                 .expect("there should be remaining points");
 
-        for elem in self.inner.call_stack.iter_tree() {
+        for elem in self.inner.call_stack.iter() {
             let instance = self
                 .instance(&elem.contract_id)
                 .expect("instance should exist");

piecrust/src/session/call_stack.rs

@@ -16,77 +16,20 @@ pub struct CallTreeElem {
     pub mem_len: usize,
 }
 
-/// A stack of contract calls.
+/// The tree of contract calls.
 #[derive(Debug, Default)]
-pub struct CallStack {
-    stack: Vec<CallTreeElem>,
-    tree: CallTree,
-}
-
-impl CallStack {
-    pub const fn new() -> Self {
-        Self {
-            stack: Vec::new(),
-            tree: CallTree::new(),
-        }
-    }
-
-    /// Push an element to the call stack.
-    pub fn push(&mut self, elem: CallTreeElem) {
-        self.tree.push(elem);
-        self.stack.push(elem);
-    }
-
-    /// Pops an element from the callstack.
-    pub fn pop(&mut self) -> Option<CallTreeElem> {
-        self.tree.pop();
-        self.stack.pop()
-    }
-
-    /// Pops an element from the callstack and prunes the call tree.
-    pub fn pop_prune(&mut self) -> Option<CallTreeElem> {
-        self.tree.pop_prune();
-        self.stack.pop()
-    }
-
-    /// Returns a view of the stack to the `n`th element from the top.
-    pub fn nth_from_top(&self, n: usize) -> Option<CallTreeElem> {
-        let len = self.stack.len();
-
-        if len > n {
-            Some(self.stack[len - (n + 1)])
-        } else {
-            None
-        }
-    }
-
-    /// Clear the call stack of all elements.
-    pub fn clear(&mut self) {
-        self.tree.clear();
-        self.stack.clear();
-    }
-
-    /// 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_tree(&self) -> CallTreeIter {
-        CallTreeIter {
-            node: self.tree.0,
-            _marker: PhantomData,
-        }
-    }
-}
-
-#[derive(Debug, Default)]
-struct CallTree(Option<*mut CallTreeInner>);
+pub struct CallTree(Option<*mut CallTreeInner>);
 
 impl CallTree {
     /// Creates a new empty call tree, starting with at the given contract.
-    const fn new() -> Self {
+    pub const fn new() -> Self {
         Self(None)
     }
 
-    /// Pushes a new child to the current node, and advances to it.
-    fn push(&mut self, elem: CallTreeElem) {
+    /// Push an element to the call tree.
+    ///
+    /// This pushes a new child to the current node, and advances to it.
+    pub fn push(&mut self, elem: CallTreeElem) {
         match self.0 {
             None => self.0 = Some(CallTreeInner::new(elem)),
             Some(inner) => unsafe {
@@ -97,32 +40,67 @@ impl CallTree {
         }
     }
 
-    /// Moves to the previous node.
-    fn pop(&mut self) {
-        self.0 = self.0.and_then(|inner| unsafe {
+    /// Moves to the previous node, returning the current element.
+    pub fn move_up(&mut self) -> Option<CallTreeElem> {
+        self.0.map(|inner| unsafe {
+            let elem = (*inner).elem;
+
             let prev = (*inner).prev;
             if prev.is_none() {
                 free_tree(inner);
             }
-            prev
-        });
+            self.0 = prev;
+
+            elem
+        })
     }
 
-    /// Clears the tree under the current node, and moves to the previous node.
-    fn pop_prune(&mut self) {
-        self.0 = self.0.and_then(|inner| unsafe {
+    /// Moves to the previous node, returning the current element, while
+    /// clearing the tree under it.
+    pub fn move_up_prune(&mut self) -> Option<CallTreeElem> {
+        self.0.map(|inner| unsafe {
+            let elem = (*inner).elem;
+
             let prev = (*inner).prev;
             if let Some(prev) = prev {
                 (*prev).children.pop();
             }
             free_tree(inner);
-            prev
-        });
+            self.0 = prev;
+
+            elem
+        })
+    }
+
+    /// Returns the `n`th previous element from the counting from the current
+    /// node.
+    ///
+    /// The zeroth previous element is the current node.
+    pub fn nth_up(&self, n: usize) -> Option<CallTreeElem> {
+        let mut current = self.0;
+
+        let mut i = 0;
+        while i < n {
+            current = current.and_then(|inner| unsafe { (*inner).prev });
+            i += 1;
+        }
+
+        current.map(|inner| unsafe { (*inner).elem })
     }
 
-    fn clear(&mut self) {
+    /// Clears the call tree of all elements.
+    pub fn clear(&mut self) {
         while self.0.is_some() {
-            self.pop();
+            self.move_up();
+        }
+    }
+
+    /// 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) -> CallTreeIter {
+        CallTreeIter {
+            node: self.0,
+            _marker: PhantomData,
         }
     }
 }