From 01f73b11d8f0070580dafdb4ac4f9ab53e1704ff Mon Sep 17 00:00:00 2001
From: Mark Thom <markjordanthom@gmail.com>
Date: Sat, 26 Dec 2020 22:35:28 -0700
Subject: [PATCH] fix bug in code walker causing entire clauses to be skipped
 (#736)

---
 src/machine/code_walker.rs | 59 +++++++++++++++++++++++++++-----------
 1 file changed, 43 insertions(+), 16 deletions(-)

diff --git a/src/machine/code_walker.rs b/src/machine/code_walker.rs
index 7cc7be4f..152395c5 100644
--- a/src/machine/code_walker.rs
+++ b/src/machine/code_walker.rs
@@ -2,30 +2,57 @@ use crate::instructions::*;
 
 use std::collections::VecDeque;
 
-fn scan_for_trust_me(code: &Code, jmp_offsets: &mut VecDeque<usize>, after_idx: &mut usize) {
-    for (idx, instr) in code[*after_idx..].iter().enumerate() {
-        match instr {
-            &Line::Choice(ChoiceInstruction::TrustMe)
-          | &Line::IndexedChoice(IndexedChoiceInstruction::Trust(..)) => {
-                *after_idx += idx;
-                return;
+fn scan_for_trust_me(
+    code: &Code,
+    jmp_offsets: &mut VecDeque<usize>,
+    before_idx: usize,
+    after_idx: &mut usize,
+) {
+    // record the location of the line after the TrustMe capping the
+    // choice instruction sequence to after_idx.
+    loop {
+        match &code[*after_idx] {
+            &Line::Choice(ChoiceInstruction::DefaultRetryMeElse(offset)) |
+            &Line::Choice(ChoiceInstruction::RetryMeElse(offset)) |
+            &Line::IndexedChoice(IndexedChoiceInstruction::Retry(offset)) => {
+                *after_idx += offset;
+            }
+            &Line::Choice(ChoiceInstruction::DefaultTrustMe) |
+            &Line::Choice(ChoiceInstruction::TrustMe) |
+            &Line::IndexedChoice(IndexedChoiceInstruction::Trust(..)) => {
+                break;
             }
+            _ => {
+                *after_idx += 1;
+            }
+        }
+    }
+
+    // search the code in the range for JmpBy instructions and record their
+    // offsets for future scanning.
+    for (idx, instr) in code[before_idx .. *after_idx].iter().enumerate() {
+        match instr {
             &Line::Control(ControlInstruction::JmpBy(_, offset, ..)) => {
-                jmp_offsets.push_back(*after_idx + idx + offset)
+                jmp_offsets.push_back(before_idx + idx + offset)
+            }
+            _ => {
             }
-            _ => {}
         }
     }
+
+    *after_idx += 1;
 }
 
 fn capture_next_range(code: &Code, queue: &mut VecDeque<usize>, last_idx: &mut usize) {
     loop {
         match &code[*last_idx] {
-            &Line::Choice(ChoiceInstruction::TryMeElse(..))
-          | &Line::IndexedChoice(IndexedChoiceInstruction::Try(..)) => {
-                    *last_idx += 1;
-                    scan_for_trust_me(code, queue, last_idx);
-                }
+            &Line::Choice(ChoiceInstruction::TryMeElse(offset)) |
+            &Line::IndexedChoice(IndexedChoiceInstruction::Try(offset)) => {
+                let before_idx = *last_idx;
+                *last_idx += offset;
+
+                scan_for_trust_me(code, queue, before_idx, last_idx);
+            }
             &Line::Control(ControlInstruction::JmpBy(_, offset, _, false)) => {
                 queue.push_back(*last_idx + offset);
                 *last_idx += 1;
@@ -34,8 +61,8 @@ fn capture_next_range(code: &Code, queue: &mut VecDeque<usize>, last_idx: &mut u
                 queue.push_back(*last_idx + offset);
                 break;
             }
-            &Line::Control(ControlInstruction::Proceed)
-          | &Line::Control(ControlInstruction::CallClause(_, _, _, true, _)) =>
+            &Line::Control(ControlInstruction::Proceed) |
+            &Line::Control(ControlInstruction::CallClause(_, _, _, true, _)) =>
                 break,
             _ =>
                 *last_idx += 1,
-- 
2.54.0