optionally read from machine stack in stackful pre-order iterator (#1812)

src/heap_iter.rs

@@ -1,8 +1,9 @@
 #[cfg(test)]
 pub(crate) use crate::machine::gc::{IteratorUMP, StacklessPreOrderHeapIter};
-use crate::machine::heap::*;
 
 use crate::atom_table::*;
+use crate::machine::heap::*;
+use crate::machine::stack::*;
 use crate::types::*;
 
 use modular_bitfield::prelude::*;
@@ -72,6 +73,7 @@ fn forward_if_referent_marked(heap: &mut [HeapCellValue], h: usize) {
 #[derive(Debug)]
 pub struct StackfulPreOrderHeapIter<'a> {
     pub heap: &'a mut Vec<HeapCellValue>,
+    machine_stack: Option<&'a Stack>,
     stack: Vec<IterStackLoc>,
     h: usize,
 }
@@ -109,10 +111,15 @@ impl<'a> StackfulPreOrderHeapIter<'a> {
         Self {
             heap,
             h,
+            machine_stack: None,
             stack: vec![IterStackLoc::iterable_heap_loc(h)],
         }
     }
 
+    pub fn iterate_over_machine_stack(&mut self, stack: &'a Stack) {
+        self.machine_stack = Some(stack);
+    }
+
     #[inline]
     pub fn push_stack(&mut self, h: usize) {
         self.stack.push(IterStackLoc::iterable_heap_loc(h));
@@ -166,6 +173,26 @@ impl<'a> StackfulPreOrderHeapIter<'a> {
         }
     }
 
+    fn stack_deref(&self, s: usize) -> Option<HeapCellValue> {
+        if let Some(stack) = &self.machine_stack {
+            let mut cell = stack[s];
+
+            while cell.is_stack_var() {
+                let s = cell.get_value();
+
+                if cell == stack[s] {
+                    break;
+                }
+
+                cell = stack[s];
+            }
+
+            return Some(cell);
+        }
+
+        None
+    }
+
     fn follow(&mut self) -> Option<HeapCellValue> {
         while let Some(h) = self.stack.pop() {
             if h.is_pending_mark() {
@@ -193,7 +220,14 @@ impl<'a> StackfulPreOrderHeapIter<'a> {
                 continue;
             }
 
-            read_heap_cell!(*cell,
+            let cell = if cell.get_tag() == HeapCellValueTag::StackVar {
+                let cell = *cell;
+                self.stack_deref(cell.get_value()).unwrap_or(cell)
+            } else {
+                *cell
+            };
+
+            read_heap_cell!(cell,
                (HeapCellValueTag::Str | HeapCellValueTag::PStrLoc, vh) => {
                    self.push_if_unmarked(vh);
                    self.stack.push(IterStackLoc::mark_heap_loc(vh));
@@ -241,7 +275,7 @@ impl<'a> StackfulPreOrderHeapIter<'a> {
                    return Some(self.heap[h]);
                }
                _ => {
-                   return Some(*cell);
+                   return Some(cell);
                }
             )
         }

src/heap_print.rs

@@ -14,6 +14,7 @@ use crate::machine::heap::*;
 use crate::machine::machine_indices::*;
 use crate::machine::machine_state::pstr_loc_and_offset;
 use crate::machine::partial_string::*;
+use crate::machine::stack::*;
 use crate::machine::streams::*;
 use crate::types::*;
 
@@ -474,6 +475,7 @@ pub struct HCPrinter<'a, Outputter> {
     outputter: Outputter,
     iter: StackfulPreOrderHeapIter<'a>,
     atom_tbl: &'a mut AtomTable,
+    stack: &'a Stack,
     op_dir: &'a OpDir,
     state_stack: Vec<TokenOrRedirect>,
     toplevel_spec: Option<DirectedOp>,
@@ -539,6 +541,7 @@ impl<'a, Outputter: HCValueOutputter> HCPrinter<'a, Outputter> {
     pub fn new(
         heap: &'a mut Heap,
         atom_tbl: &'a mut AtomTable,
+        stack: &'a Stack,
         op_dir: &'a OpDir,
         output: Outputter,
         cell: HeapCellValue,
@@ -547,6 +550,7 @@ impl<'a, Outputter: HCValueOutputter> HCPrinter<'a, Outputter> {
             outputter: output,
             iter: stackful_preorder_iter(heap, cell),
             atom_tbl,
+            stack,
             op_dir,
             state_stack: vec![],
             toplevel_spec: None,
@@ -1441,12 +1445,7 @@ impl<'a, Outputter: HCValueOutputter> HCPrinter<'a, Outputter> {
     ) {
         let negated_operand = negated_op_needs_bracketing(&self.iter, self.op_dir, &op);
 
-        let addr = match self.check_for_seen() {
+        let print_struct = |printer: &mut Self, name: Atom, arity: usize| {
-            Some(addr) => addr,
-            None => return,
-        };
-
-        let print_atom = |printer: &mut Self, name: Atom, arity: usize| {
             if name == atom!("[]") && arity == 0 {
                 if !printer.at_cdr("") {
                     append_str!(printer, "[]");
@@ -1494,14 +1493,18 @@ impl<'a, Outputter: HCValueOutputter> HCPrinter<'a, Outputter> {
             }
         };
 
+        let addr = match self.check_for_seen() {
+            Some(addr) => addr,
+            None => return,
+        };
+
         read_heap_cell!(addr,
             (HeapCellValueTag::Atom, (name, arity)) => {
-                print_atom(self, name, arity);
+                print_struct(self, name, arity);
             }
             (HeapCellValueTag::Char, c) => {
                 let name = self.atom_tbl.build_with(&String::from(c));
-                print_atom(self, name, 0);
+                print_struct(self, name, 0);
-                // print_char!(self, self.quoted, c);
             }
             (HeapCellValueTag::Str, s) => {
                 let (name, arity) = cell_as_atom_cell!(self.iter.heap[s])
@@ -1594,6 +1597,8 @@ impl<'a, Outputter: HCValueOutputter> HCPrinter<'a, Outputter> {
 
     pub fn print(mut self) -> Outputter {
         let spec = self.toplevel_spec.take();
+
+        self.iter.iterate_over_machine_stack(self.stack);
         self.handle_heap_term(spec, false, self.max_depth);
 
         while let Some(loc_data) = self.state_stack.pop() {
@@ -1665,6 +1670,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0)
@@ -1693,6 +1699,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0)
@@ -1716,6 +1723,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0)
@@ -1728,6 +1736,7 @@ mod tests {
             let mut printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0)
@@ -1760,6 +1769,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0),
@@ -1778,6 +1788,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0),
@@ -1794,6 +1805,7 @@ mod tests {
             let mut printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0)
@@ -1825,6 +1837,7 @@ mod tests {
             let mut printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0)
@@ -1847,6 +1860,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 pstr_loc_as_cell!(0)
@@ -1874,6 +1888,7 @@ mod tests {
             let printer = HCPrinter::new(
                 &mut wam.machine_st.heap,
                 &mut wam.machine_st.atom_tbl,
+                &wam.machine_st.stack,
                 &wam.op_dir,
                 PrinterOutputter::new(),
                 heap_loc_as_cell!(0),

src/machine/machine_state.rs

@@ -766,6 +766,7 @@ impl MachineState {
                 let mut printer = HCPrinter::new(
                     &mut self.heap,
                     &mut self.atom_tbl,
+                    &mut self.stack,
                     op_dir,
                     PrinterOutputter::new(),
                     term_to_be_printed,

src/machine/mock_wam.rs

@@ -62,6 +62,7 @@ impl MockWAM {
         let mut printer = HCPrinter::new(
             &mut self.machine_st.heap,
             &mut self.machine_st.atom_tbl,
+            &mut self.machine_st.stack,
             &self.op_dir,
             PrinterOutputter::new(),
             heap_loc_as_cell!(term_write_result.heap_loc),