From c369ce9a7f36723394c2b77fadc8e424abb233ca Mon Sep 17 00:00:00 2001
From: Mark Thom <markjordanthom@gmail.com>
Date: Fri, 4 May 2018 20:51:02 -0600
Subject: [PATCH] add access to HeapVarDict to the printer.

---
 src/prolog/heap_iter.rs                  |  8 +++++++
 src/prolog/heap_print.rs                 | 15 ++++++++----
 src/prolog/machine/machine_state_impl.rs | 30 ++++++++++++++++++++----
 src/prolog/machine/mod.rs                | 10 +++-----
 4 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/src/prolog/heap_iter.rs b/src/prolog/heap_iter.rs
index 33b552aa..8532e55a 100644
--- a/src/prolog/heap_iter.rs
+++ b/src/prolog/heap_iter.rs
@@ -184,6 +184,14 @@ pub struct HCDerefAcyclicIterator<HCIter> {
     seen: HashSet<Addr>
 }
 
+pub type HCDerefAcyclicPreOrderIterator<'a> = HCDerefAcyclicIterator<HCPreOrderIterator<'a>>;
+
+impl<'a> HCPreOrderIterator<'a> {
+    pub fn deref_acyclic_iter(self) -> HCDerefAcyclicPreOrderIterator<'a> {
+        HCDerefAcyclicIterator::new(self)
+    }
+}
+
 impl<HCIter: MutStackHCIterator> HCDerefAcyclicIterator<HCIter>
 {
     pub fn new(iter: HCIter) -> Self {
diff --git a/src/prolog/heap_print.rs b/src/prolog/heap_print.rs
index cb8d7bd6..3923f72b 100644
--- a/src/prolog/heap_print.rs
+++ b/src/prolog/heap_print.rs
@@ -1,5 +1,6 @@
 use prolog::ast::*;
 use prolog::heap_iter::*;
+use prolog::machine::machine_state::MachineState;
 
 use std::borrow::Cow;
 use std::cell::Cell;
@@ -143,7 +144,7 @@ impl HCValueFormatter for TermFormatter {
 pub struct HCPrinter<'a, Formatter, Outputter> {
     formatter:   Formatter,
     outputter:   Outputter,
-    iter:        HCPreOrderIterator<'a>,
+    iter:        HCDerefAcyclicPreOrderIterator<'a>,
     state_stack: Vec<TokenOrRedirect>,
     heap_locs:   Cow<'a, HeapVarDict>
 }
@@ -151,22 +152,26 @@ pub struct HCPrinter<'a, Formatter, Outputter> {
 impl<'a, Formatter: HCValueFormatter, Outputter: HCValueOutputter>
     HCPrinter<'a, Formatter, Outputter>
 {
-    pub fn new(iter: HCPreOrderIterator<'a>, fmt: Formatter, output: Outputter) -> Self
+    pub fn new(machine_st: &'a MachineState, addr: Addr, fmt: Formatter, output: Outputter) -> Self
     {
+        let iter = HCPreOrderIterator::new(&machine_st, addr);
+        
         HCPrinter { formatter: fmt,
                     outputter: output,
-                    iter,
+                    iter: iter.deref_acyclic_iter(),
                     state_stack: vec![],
                     heap_locs: Cow::default() }
     }
 
-    pub fn from_heap_locs(iter: HCPreOrderIterator<'a>, fmt: Formatter,
+    pub fn from_heap_locs(machine_st: &'a MachineState, addr: Addr, fmt: Formatter,
                           output: Outputter, heap_locs: &'a HeapVarDict)
                           -> Self
     {
+        let iter = HCPreOrderIterator::new(&machine_st, addr);
+        
         HCPrinter { formatter: fmt,
                     outputter: output,
-                    iter,
+                    iter: iter.deref_acyclic_iter(),
                     state_stack: vec![],
                     heap_locs: Cow::Borrowed(heap_locs) }
     }
diff --git a/src/prolog/machine/machine_state_impl.rs b/src/prolog/machine/machine_state_impl.rs
index e425d2b0..d0cf402e 100644
--- a/src/prolog/machine/machine_state_impl.rs
+++ b/src/prolog/machine/machine_state_impl.rs
@@ -93,12 +93,34 @@ impl MachineState {
         self.trail(r1);
     }
 
-    pub(super) fn print_term<Fmt, Outputter>(&self, a: Addr, fmt: Fmt, output: Outputter) -> Outputter
-      where Fmt: HCValueFormatter, Outputter: HCValueOutputter
+    pub(super)
+    fn print_var_eq<Fmt, Outputter>(&self, var: Rc<Var>, addr: Addr, var_dir: &HeapVarDict,
+                                    fmt: Fmt, mut output: Outputter)
+                                    -> Outputter
+        where Fmt: HCValueFormatter, Outputter: HCValueOutputter
     {
-        let iter    = HCPreOrderIterator::new(&self, a);
-        let printer = HCPrinter::new(iter, fmt, output);
+        let orig_len = output.len();
+        
+        output.begin_new_var();
+
+        output.append(var.as_str());
+        output.append(" = ");
+        
+        let printer    = HCPrinter::from_heap_locs(&self, addr, fmt, output, var_dir);
+        let mut output = printer.print();
+
+        if output.ends_with(var.as_str()) {
+            output.truncate(orig_len);
+        }
 
+        output
+    }
+    
+    pub(super)
+    fn print_term<Fmt, Outputter>(&self, addr: Addr, fmt: Fmt, output: Outputter) -> Outputter
+      where Fmt: HCValueFormatter, Outputter: HCValueOutputter
+    {
+        let printer = HCPrinter::new(&self, addr, fmt, output);
         printer.print()
     }
 
diff --git a/src/prolog/machine/mod.rs b/src/prolog/machine/mod.rs
index fa19f6e5..b8294658 100644
--- a/src/prolog/machine/mod.rs
+++ b/src/prolog/machine/mod.rs
@@ -4,7 +4,7 @@ use prolog::heap_print::*;
 use prolog::tabled_rc::*;
 
 mod machine_errors;
-pub(crate) mod machine_state;
+pub(super) mod machine_state;
 #[macro_use]
 mod machine_state_impl;
 
@@ -427,12 +427,8 @@ impl Machine {
        where Outputter: HCValueOutputter
     {
         for (var, addr) in var_dir {
-            output.begin_new_var();
-
-            output.append(var.as_str());
-            output.append(" = ");
-
-            output = self.ms.print_term(addr.clone(), TermFormatter {}, output);
+            let fmt = TermFormatter {};
+            output = self.ms.print_var_eq(var.clone(), addr.clone(), var_dir, fmt, output);
         }
 
         output
-- 
2.54.0