follow marked variables to end in eager_stackful_iter (#2100, #2101)

src/heap_iter.rs

@@ -56,41 +56,58 @@ impl<'a> EagerStackfulPreOrderHeapIter<'a> {
         }
     }
 
-    fn follow(&mut self) -> Option<HeapCellValue> {
-        while let Some(value) = self.iter_stack.pop() {
-            if value.get_mark_bit() == self.mark_phase {
+    #[inline]
+    fn is_self_ref_var(&self, value: HeapCellValue) -> bool {
         if value.is_var() {
             let h = value.get_value() as usize;
 
             if self.heap[h].is_var() && self.heap[h].get_value() as usize == h {
-                        return Some(unmark_cell_bits!(value));
+                return true;
             }
         }
 
+        false
+    }
+
+    fn follow(&mut self) -> Option<HeapCellValue> {
+        while let Some(value) = self.iter_stack.pop() {
+            if value.get_mark_bit() == self.mark_phase {
+                // follow marked variables to their end. only marked
+                // non-variables are ignored.
+                if self.is_self_ref_var(value) {
+                    return Some(unmark_cell_bits!(value));
+                } else if !value.is_var() {
                     continue;
                 }
+            }
 
             read_heap_cell!(value,
                 (HeapCellValueTag::Str, s) => {
                     let arity = cell_as_atom_cell!(self.heap[s]).get_arity();
 
                     for idx in (s + 1 .. s + arity + 1).rev() {
+			if self.heap[idx].get_mark_bit() != self.mark_phase {
                             self.iter_stack.push(self.heap[idx]);
                             self.heap[idx].set_mark_bit(self.mark_phase);
 			}
                     }
+                }
                 (HeapCellValueTag::Lis, l) => {
+		    if self.heap[l+1].get_mark_bit() != self.mark_phase {
 			self.iter_stack.push(self.heap[l+1]);
-                    self.iter_stack.push(self.heap[l]);
-
-                    self.heap[l].set_mark_bit(self.mark_phase);
 			self.heap[l+1].set_mark_bit(self.mark_phase);
 		    }
+
+		    if self.heap[l].get_mark_bit() != self.mark_phase {
+			self.iter_stack.push(self.heap[l]);
+			self.heap[l].set_mark_bit(self.mark_phase);
+		    }
+                }
                 (HeapCellValueTag::AttrVar | HeapCellValueTag::Var, h) => {
                     let var_value = self.heap[h];
                     self.heap[h].set_mark_bit(self.mark_phase);
 
-                    if !(var_value.is_var() && var_value.get_value() as usize == h) {
+                    if !(self.heap[h].is_var() && self.heap[h].get_value() as usize == h) {
                         self.iter_stack.push(var_value);
                         continue;
                     }

src/tests/term_variables.pl

@@ -43,12 +43,34 @@ test("term_variables#2097", (
     T = [[[A|B]|A]|A], Vs == [A,B]
 )).
 
+test("term_variables#2100", (
+    termt2(T), term_variables(T,Vs),
+    T = [[T|A]|B], Vs == [A,B]
+)).
+
+test("term_variables#2101", (
+    termt3(T), term_variables(T,Vs),
+    T = [[[[A|B]|A]|A]|A], Vs == [A, B]
+)).
+
 termt(T) :-
    T = [T1|T2],
    T1 = [T3|A],
    T3 = [A|_],
    T2 = A.
 
+termt2(T) :-
+   T = [T1|_B],
+   T1 = [T|_A].
+
+termt3(T) :-
+   T = [T1|T0],
+   T1 = [T2|T3],
+   T2 = [T4|A],
+   T4 = [A|_],
+   T3 = A,
+   T0 = A.
+
 main :-
     findall(test(Name, Goal), test(Name, Goal), Tests),
     run_tests(Tests, Failed),