From f939a9e1290f86f5fbdf30105aab5671af6ebdfe Mon Sep 17 00:00:00 2001
From: Mark Thom <markjordanthom@gmail.com>
Date: Sat, 27 Jan 2018 01:27:17 -0700
Subject: [PATCH] correct marker bug, edit README.

---
 README.md                      |  2 +-
 src/prolog/allocator.rs        |  1 +
 src/prolog/codegen.rs          | 10 +++++-----
 src/prolog/debray_allocator.rs | 28 ++++++++++++++++++----------
 src/tests.rs                   |  4 +++-
 5 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 91194b12..fc709ac8 100644
--- a/README.md
+++ b/README.md
@@ -20,7 +20,7 @@ Extend rusty-wam to include the following, among other features:
 * Built-in and user-defined operators of all fixities, with custom
   associativity and precedence (_done_). 
 * Bignum, rational number and floating point arithmetic (_done_).
-* Built-in control operators (`,`, `;`, `->`, etc.) (_in progress_).
+* Built-in control operators (`,`, `;`, `->`, etc.) (_done_).
 * Built-in predicates for list processing and top-level declarative
   control (`setup_call_control/3`, `call_with_inference_limit/3`,
   etc.)
diff --git a/src/prolog/allocator.rs b/src/prolog/allocator.rs
index cd9f523b..7aa0eb7c 100644
--- a/src/prolog/allocator.rs
+++ b/src/prolog/allocator.rs
@@ -18,6 +18,7 @@ pub trait Allocator<'a>
     fn reset(&mut self);
     fn reset_contents(&mut self) {}
     fn reset_arg(&mut self, usize);
+    fn reset_arg_at_head(&mut self, &Term);
     
     fn advance_arg(&mut self);
 
diff --git a/src/prolog/codegen.rs b/src/prolog/codegen.rs
index ec620bc1..811ddbec 100644
--- a/src/prolog/codegen.rs
+++ b/src/prolog/codegen.rs
@@ -483,12 +483,12 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<'a, TermMarker>
         let conjunct_info = self.collect_var_data(iter);
 
         let &Rule { head: (ref p0, ref p1), ref clauses } = rule;
-        let mut code = Vec::new();
-
-        self.marker.reset_arg(p0.arity());
-        self.compile_seq_prelude(&conjunct_info, &mut code);
+        let mut code = Vec::new();        
 
         if let &QueryTerm::Term(ref term) = p0 {
+            self.marker.reset_arg_at_head(term);
+            self.compile_seq_prelude(&conjunct_info, &mut code);
+            
             if let &Term::Clause(..) = term {
                 let iter = FactInstruction::iter(term);
                 let fact = self.compile_target(iter, GenContext::Head, false);
@@ -557,7 +557,7 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<'a, TermMarker>
         vs.populate_restricting_sets();
 
         self.marker.drain_var_data(vs);
-        self.marker.reset_arg(term.arity());
+        self.marker.reset_arg_at_head(term);
 
         let mut code = Vec::new();
 
diff --git a/src/prolog/debray_allocator.rs b/src/prolog/debray_allocator.rs
index a7f4a0f2..e922d60a 100644
--- a/src/prolog/debray_allocator.rs
+++ b/src/prolog/debray_allocator.rs
@@ -138,10 +138,7 @@ impl<'a> DebrayAllocator<'a> {
         };
     }
 
-    fn alloc_reg_to_var<Target>(&mut self,
-                                var: &'a Var,
-                                lvl: Level,
-                                term_loc: GenContext,
+    fn alloc_reg_to_var<Target>(&mut self, var: &'a Var, lvl: Level, term_loc: GenContext,
                                 target: &mut Vec<Target>)
                                 -> usize
         where Target: CompilationTarget<'a>
@@ -250,12 +247,8 @@ impl<'a> Allocator<'a> for DebrayAllocator<'a>
         }
     }
 
-    fn mark_var<Target>(&mut self,
-                        var: &'a Var,
-                        lvl: Level,
-                        cell: &'a Cell<VarReg>,
-                        term_loc: GenContext,
-                        target: &mut Vec<Target>)
+    fn mark_var<Target>(&mut self, var: &'a Var, lvl: Level, cell: &'a Cell<VarReg>,
+                        term_loc: GenContext, target: &mut Vec<Target>)
         where Target: CompilationTarget<'a>
     {
         let (r, is_new_var) = match self.get(var) {
@@ -352,4 +345,19 @@ impl<'a> Allocator<'a> for DebrayAllocator<'a>
         self.arg_c   = 1;
         self.temp_lb = arity + 1;
     }
+
+    fn reset_arg_at_head(&mut self, term: &Term) {
+        self.arg_c = 1;
+        self.temp_lb = term.arity() + 1;
+
+        match term {
+            &Term::Clause(_, _, ref subterms, _) =>
+                for (idx, tr) in subterms.iter().enumerate() {
+                    if let &Term::Var(_, _) = tr.as_ref() {
+                        self.in_use.insert(idx + 1);
+                    }
+                },
+            _ => {}
+        };
+    }
 }
diff --git a/src/tests.rs b/src/tests.rs
index 737e2361..afd5414b 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -1172,7 +1172,9 @@ fn test_queries_on_conditionals()
     assert_prolog_success!(&mut wam, "?- catch(test(a, [a]), type_error(E), true).",
                            [["E = _6"], ["E = _6"]]);
 
-    //TODO: write tests for calling ;, ->, confirm behavior is correct.
+    //TODO: write tests for calling ;, ->, to confirm behavior is correct.
+    assert_prolog_success!(&mut wam, "?- f(X), call(->, atomic(X), true).",
+                           [["X = a"], ["X = b"]]);
 }
 
 #[test]
-- 
2.54.0