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

src/heap_iter.rs

@@ -56,18 +56,29 @@ impl<'a> EagerStackfulPreOrderHeapIter<'a> {
         }
     }
 
+    #[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 true;
+            }
+        }
+
+        false
+    }
+
     fn follow(&mut self) -> Option<HeapCellValue> {
         while let Some(value) = self.iter_stack.pop() {
             if value.get_mark_bit() == self.mark_phase {
-                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));
-                    }
+                // 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;
                 }
-
-                continue;
             }
 
             read_heap_cell!(value,
@@ -90,7 +101,7 @@ impl<'a> EagerStackfulPreOrderHeapIter<'a> {
                     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|_B]|_A], 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),