re: issue #199

src/prolog/allocator.rs

@@ -44,7 +44,11 @@ pub trait Allocator<'a> {
 
     fn take_bindings(self) -> AllocVarDict;
 
-    fn drain_var_data(&mut self, vs: VariableFixtures<'a>) -> VariableFixtures<'a> {
+    fn drain_var_data(
+        &mut self,
+        vs: VariableFixtures<'a>,
+        num_of_chunks: usize
+    ) -> VariableFixtures<'a> {
         let mut perm_vs = VariableFixtures::new();
 
         for (var, (var_status, cells)) in vs.into_iter() {
@@ -52,6 +56,10 @@ pub trait Allocator<'a> {
                 VarStatus::Temp(chunk_num, tvd) => {
                     self.bindings_mut()
                         .insert(var.clone(), VarData::Temp(chunk_num, 0, tvd));
+
+                    if chunk_num + 1 == num_of_chunks {
+                        perm_vs.insert_last_chunk_temp_var(var);
+                    }
                 }
                 VarStatus::Perm(_) => {
                     self.bindings_mut().insert(var.clone(), VarData::Perm(0));

src/prolog/codegen.rs

@@ -51,20 +51,32 @@ impl<'a> ConjunctInfo<'a> {
         self.has_deep_cut as usize
     }
 
-    fn mark_unsafe_vars(&self, mut unsafe_var_marker: UnsafeVarMarker, code: &mut Code) {
+    fn mark_unsafe_vars<Alloc: Allocator<'a>>(
+        &self,
+        mut unsafe_var_marker: UnsafeVarMarker,
+        marker: &Alloc,
+        code: &mut Code
+    ) {
         // target the last goal of the rule for handling unsafe variables.
         // we use this weird logic to find the last goal.
-        let right_index = if let &Line::Control(_) = code.last().unwrap() {
+        let right_index = if let Some(Line::Control(_)) = code.last() {
+            if code.len() >= 2 {
                 code.len() - 2
             } else {
+                return;
+            }
+        } else {
+            if code.len() >= 1 {
                 code.len() - 1
+            } else {
+                return;
+            }
         };
 
         let mut index = right_index;
 
         if let Line::Query(_) = &code[right_index] {
             while let Line::Query(_) = &code[index] {
-                // index >= 0.
                 if index == 0 {
                     break;
                 } else {
@@ -77,10 +89,10 @@ impl<'a> ConjunctInfo<'a> {
                 index += 1;
             }
 
-            unsafe_var_marker.record_unsafe_vars(&self.perm_vs);
+            unsafe_var_marker.record_unsafe_vars(&self.perm_vs, marker);
 
-            for line in code.iter_mut() {
+            for line in code.iter() {
-                if let &mut Line::Query(ref mut query_instr) = line {
+                if let Line::Query(ref query_instr) = line {
                     unsafe_var_marker.mark_safe_vars(query_instr);
                 }
             }
@@ -273,7 +285,8 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<TermMarker> {
         while let Some((chunk_num, lt_arity, chunked_terms)) = iter.next() {
             for (i, chunked_term) in chunked_terms.iter().enumerate() {
                 let term_loc = match chunked_term {
-                    &ChunkedTerm::HeadClause(..) => GenContext::Head,
+                    &ChunkedTerm::HeadClause(..) =>                       
+                        GenContext::Head,
                     &ChunkedTerm::BodyTerm(_) => {
                         if i < chunked_terms.len() - 1 {
                             GenContext::Mid(chunk_num)
@@ -294,8 +307,7 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<TermMarker> {
         vs.populate_restricting_sets();
         vs.set_perm_vals(has_deep_cut);
 
-        let vs = self.marker.drain_var_data(vs);
+        let vs = self.marker.drain_var_data(vs, num_of_chunks);
-
         ConjunctInfo::new(vs, num_of_chunks, has_deep_cut)
     }
 
@@ -703,9 +715,7 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<TermMarker> {
         let iter = ChunkedIterator::from_rule_body(p1, clauses);
         self.compile_seq(iter, &conjunct_info, &mut code, false)?;
 
-        if conjunct_info.allocates() {
+        conjunct_info.mark_unsafe_vars(unsafe_var_marker, &self.marker, &mut code);        
-            conjunct_info.mark_unsafe_vars(unsafe_var_marker, &mut code);
-        }
 
         Self::compile_cleanup(&mut code, &conjunct_info, clauses.last().unwrap_or(p1));
         Ok(code)
@@ -747,7 +757,7 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<TermMarker> {
         vs.mark_vars_in_chunk(post_order_iter(term), term.arity(), GenContext::Head);
 
         vs.populate_restricting_sets();
-        self.marker.drain_var_data(vs);
+        self.marker.drain_var_data(vs, 1);
 
         let mut code = Vec::new();
 
@@ -802,9 +812,7 @@ impl<'a, TermMarker: Allocator<'a>> CodeGenerator<TermMarker> {
         let iter = ChunkedIterator::from_term_sequence(query);
         self.compile_seq(iter, &conjunct_info, &mut code, true)?;
 
-        if conjunct_info.allocates() {
+        conjunct_info.mark_unsafe_vars(UnsafeVarMarker::new(), &self.marker, &mut code);
-            conjunct_info.mark_unsafe_vars(UnsafeVarMarker::new(), &mut code);
-        }
 
         if let Some(query_term) = query.last() {
             Self::compile_cleanup(&mut code, &conjunct_info, query_term);

src/prolog/fixtures.rs

@@ -1,14 +1,14 @@
 use prolog_parser::ast::*;
 
+use crate::prolog::allocator::*;
 use crate::prolog::forms::*;
 use crate::prolog::instructions::*;
 use crate::prolog::iterators::*;
 
-use indexmap::IndexMap;
+use indexmap::{IndexMap, IndexSet};
 
 use std::cell::Cell;
-use std::collections::btree_map::{IntoIter, IterMut, Values};
+use std::collections::BTreeSet;
-use std::collections::{BTreeMap, BTreeSet};
 use std::mem::swap;
 use std::rc::Rc;
 use std::vec::Vec;
@@ -79,15 +79,27 @@ impl TempVarData {
 }
 
 type VariableFixture<'a> = (VarStatus, Vec<&'a Cell<VarReg>>);
-pub struct VariableFixtures<'a>(BTreeMap<Rc<Var>, VariableFixture<'a>>);
+
+pub struct VariableFixtures<'a>{
+    perm_vars: IndexMap<Rc<Var>, VariableFixture<'a>>,
+    last_chunk_temp_vars: IndexSet<Rc<Var>>
+}
 
 impl<'a> VariableFixtures<'a> {
     pub fn new() -> Self {
-        VariableFixtures(BTreeMap::new())
+        VariableFixtures {
+            perm_vars: IndexMap::new(),
+            last_chunk_temp_vars: IndexSet::new()
+        }
+                          
     }
 
     pub fn insert(&mut self, var: Rc<Var>, vs: VariableFixture<'a>) {
-        self.0.insert(var, vs);
+        self.perm_vars.insert(var, vs);
+    }
+
+    pub fn insert_last_chunk_temp_var(&mut self, var: Rc<Var>) {
+        self.last_chunk_temp_vars.insert(var);
     }
 
     // computes no_use and conflict sets for all temp vars.
@@ -140,11 +152,11 @@ impl<'a> VariableFixtures<'a> {
     }
 
     fn get_mut(&mut self, u: Rc<Var>) -> Option<&mut VariableFixture<'a>> {
-        self.0.get_mut(&u)
+        self.perm_vars.get_mut(&u)
     }
 
-    fn iter_mut(&mut self) -> IterMut<Rc<Var>, VariableFixture<'a>> {
+    fn iter_mut(&mut self) -> indexmap::map::IterMut<Rc<Var>, VariableFixture<'a>> {
-        self.0.iter_mut()
+        self.perm_vars.iter_mut()
     }
 
     fn record_temp_info(&mut self, tvd: &mut TempVarData, arg_c: usize, term_loc: GenContext) {
@@ -179,7 +191,7 @@ impl<'a> VariableFixtures<'a> {
 
         for term_ref in iter {
             if let &TermRef::Var(lvl, cell, ref var) = &term_ref {
-                let mut status = self.0.remove(var).unwrap_or((
+                let mut status = self.perm_vars.remove(var).unwrap_or((
                     VarStatus::Temp(chunk_num, TempVarData::new(lt_arity)),
                     Vec::new(),
                 ));
@@ -195,7 +207,7 @@ impl<'a> VariableFixtures<'a> {
                     _ => status.0 = VarStatus::Perm(chunk_num),
                 };
 
-                self.0.insert(var.clone(), status);
+                self.perm_vars.insert(var.clone(), status);
             }
 
             if let Level::Shallow = term_ref.level() {
@@ -204,16 +216,16 @@ impl<'a> VariableFixtures<'a> {
         }
     }
 
-    pub fn into_iter(self) -> IntoIter<Rc<Var>, VariableFixture<'a>> {
+    pub fn into_iter(self) -> indexmap::map::IntoIter<Rc<Var>, VariableFixture<'a>> {
-        self.0.into_iter()
+        self.perm_vars.into_iter()
     }
 
-    fn values(&self) -> Values<Rc<Var>, VariableFixture<'a>> {
+    fn values(&self) -> indexmap::map::Values<Rc<Var>, VariableFixture<'a>> {
-        self.0.values()
+        self.perm_vars.values()
     }
 
     pub fn size(&self) -> usize {
-        self.0.len()
+        self.perm_vars.len()
     }
 
     pub fn set_perm_vals(&self, has_deep_cuts: bool) {
@@ -248,7 +260,11 @@ impl UnsafeVarMarker {
         }
     }
 
-    pub fn record_unsafe_vars(&mut self, fixtures: &VariableFixtures) {
+    pub fn record_unsafe_vars<'a, Alloc: Allocator<'a>>(
+        &mut self,
+        fixtures: &VariableFixtures,
+        marker: &Alloc
+    ) {
         for &(_, ref cb) in fixtures.values() {           
             if let Some(index) = cb.first() {
                 if !self.unsafe_vars.contains_key(&index.get().norm()) {
@@ -256,17 +272,22 @@ impl UnsafeVarMarker {
                 }
             }
         }
+
+        for var in fixtures.last_chunk_temp_vars.iter().cloned() {
+            let r = marker.get(var);
+            self.unsafe_vars.insert(r, false);
+        }
     }
 
-    pub fn mark_safe_vars(&mut self, query_instr: &mut QueryInstruction) {
+    pub fn mark_safe_vars(&mut self, query_instr: &QueryInstruction) {
         match query_instr {
-            &mut QueryInstruction::PutVariable(RegType::Temp(r), _) => {
+            QueryInstruction::PutVariable(RegType::Temp(r), _) => {
-                if let Some(found) = self.unsafe_vars.get_mut(&RegType::Temp(r)) {
+                if let Some(found) = self.unsafe_vars.get_mut(&RegType::Temp(*r)) {
                     *found = true;
                 }
             }
-            &mut QueryInstruction::SetVariable(reg) => {
+            QueryInstruction::SetVariable(reg) => {
-                if let Some(found) = self.unsafe_vars.get_mut(&reg) {
+                if let Some(found) = self.unsafe_vars.get_mut(reg) {
                     *found = true;
                 }
             }