fix: avoid revisiting locations in CFG traversal (#121)

Author: toolCHAINZ

jingle/examples/unwind.rs

@@ -30,7 +30,7 @@ fn main() {
     z3::set_global_param("trace", "true");
     let loaded = load_with_gimli(bin_path, "/Applications/ghidra").unwrap();
 
-    let mut direct = UnwindingAnalysis::new(1);
+    let mut direct = UnwindingAnalysis::new(10);
     let pcode_graph = direct.run(&loaded, direct.make_initial_state(FUNC_NESTED.into()));
     // let pcode_graph = pcode_graph.basic_blocks();
     let addrs = pcode_graph.nodes().collect::<Vec<_>>();

jingle/src/analysis/cfg/mod.rs

@@ -10,8 +10,10 @@ use petgraph::graph::NodeIndex;
 use petgraph::prelude::DiGraph;
 use petgraph::visit::EdgeRef;
 use std::borrow::Borrow;
-use std::collections::HashMap;
+use std::cell::RefCell;
+use std::collections::{HashMap, HashSet};
 use std::fmt::{Formatter, LowerHex};
+use std::rc::Rc;
 use z3::ast::Bool;
 
 mod model;
@@ -38,17 +40,37 @@ pub struct PcodeCfg<N: CfgState = ConcretePcodeAddress, D = PcodeOperation> {
 pub struct PcodeCfgVisitor<'a, N: CfgState = ConcretePcodeAddress, D = PcodeOperation> {
     cfg: &'a PcodeCfg<N, D>,
     location: N,
+    pub(crate) visited_locations: Rc<RefCell<HashSet<N>>>,
 }
 
 impl<'a, N: CfgState, D: ModelTransition<N::Model>> PcodeCfgVisitor<'a, N, D> {
-    pub(crate) fn successors(&self) -> impl Iterator<Item = Self> {
+    pub(crate) fn successors(&mut self) -> impl Iterator<Item = Self> {
         self.cfg
             .successors(&self.location)
             .into_iter()
             .flatten()
-            .map(|n| Self {
-                cfg: self.cfg,
-                location: n.clone(),
+            .flat_map(|n| {
+                // Use a short-lived borrow so the RefMut is released before we construct the new visitor
+                let is_repeat = {
+                    let mut set = self.visited_locations.borrow_mut();
+                    if set.contains(n) {
+                        println!("Trimming repeat of {n:x?}");
+                        true
+                    } else {
+                        set.insert(n.clone());
+                        false
+                    }
+                };
+
+                if is_repeat {
+                    None
+                } else {
+                    Some(Self {
+                        cfg: self.cfg,
+                        location: n.clone(),
+                        visited_locations: self.visited_locations.clone(),
+                    })
+                }
             })
     }
 
@@ -288,10 +310,11 @@ impl<D: ModelTransition<MachineState>> PcodeCfg<UnwoundLocation, D> {
         location: &UnwoundLocation,
         ctl_model: CtlFormula<UnwoundLocation, D>,
     ) -> Bool {
-        let visitor = PcodeCfgVisitor {
+        let mut visitor = PcodeCfgVisitor {
             location: location.clone(),
             cfg: self,
+            visited_locations: Rc::new(RefCell::new(HashSet::new())),
         };
-        ctl_model.check(&visitor)
+        ctl_model.check(&mut visitor)
     }
 }

jingle/src/analysis/ctl.rs

@@ -309,7 +309,7 @@ impl<N: CfgState, D: ModelTransition<N::Model>> std::fmt::Debug for CtlFormula<N
     }
 }
 impl<N: CfgState, D: ModelTransition<N::Model>> CtlFormula<N, D> {
-    pub fn check(&self, g: &PcodeCfgVisitor<N, D>) -> Bool {
+    pub fn check(&self, g: &mut PcodeCfgVisitor<N, D>) -> Bool {
         match self {
             CtlFormula::Bottom => Bool::from_bool(false),
             CtlFormula::Top => Bool::from_bool(true),
@@ -351,24 +351,26 @@ impl<N: CfgState, D: ModelTransition<N::Model>> CtlFormula<N, D> {
         }
     }
 
-    pub(crate) fn check_next_exists(&self, g: &PcodeCfgVisitor<N, D>) -> Bool {
-        let state = g.state().unwrap();
-        let connect: Vec<_> = g
-            .successors()
+    pub(crate) fn check_next_exists(&self, g: &mut PcodeCfgVisitor<N, D>) -> Bool {
+        let state = g.state().unwrap().clone();
+        let trans = g.transition().cloned();
+        let successors: Vec<_> = g.successors().collect();
+        let connect: Vec<_> = successors
+            .iter()
             .map(|a| {
                 let successor = a.state().unwrap();
-                let after = g.transition().unwrap().transition(state).unwrap();
+                let after = trans.clone().unwrap().transition(&state).unwrap();
 
                 after.location_eq(successor)
             })
             .collect();
         let connect = Bool::or(&connect);
-        let bools: Vec<_> = g
-            .successors()
-            .flat_map(|a| {
+        let bools: Vec<_> = successors
+            .into_iter()
+            .flat_map(|mut a| {
+                let check = self.check(&mut a);
                 let successor = a.state().unwrap();
-                let check = self.check(&a);
-                let after = g.transition().unwrap().transition(state).unwrap();
+                let after = trans.clone().unwrap().transition(&state).unwrap();
                 let imp = after
                     .location_eq(successor)
                     .implies(after.state_eq(successor));
@@ -378,24 +380,26 @@ impl<N: CfgState, D: ModelTransition<N::Model>> CtlFormula<N, D> {
         Bool::or(&bools).bitand(connect)
     }
 
-    pub(crate) fn check_next_universal(&self, g: &PcodeCfgVisitor<N, D>) -> Bool {
-        let state = g.state().unwrap();
-        let connect: Vec<_> = g
-            .successors()
+    pub(crate) fn check_next_universal(&self, g: &mut PcodeCfgVisitor<N, D>) -> Bool {
+        let state = g.state().unwrap().clone();
+        let trans = g.transition().cloned();
+        let successors: Vec<_> = g.successors().collect();
+        let connect: Vec<_> = successors
+            .iter()
             .map(|a| {
                 let successor = a.state().unwrap();
-                let after = g.transition().unwrap().transition(state).unwrap();
+                let after = trans.clone().unwrap().transition(&state).unwrap();
 
                 after.location_eq(successor)
             })
             .collect();
         let connect = Bool::or(&connect);
-        let bools: Vec<_> = g
-            .successors()
-            .flat_map(|a| {
+        let bools: Vec<_> = successors
+            .into_iter()
+            .flat_map(|mut a| {
+                let check = self.check(&mut a);
                 let successor = a.state().unwrap();
-                let check = self.check(&a);
-                let after = g.transition().unwrap().transition(state).unwrap();
+                let after = trans.clone().unwrap().transition(&state).unwrap();
                 let imp = after
                     .location_eq(successor)
                     .implies(after.state_eq(successor));