Enhanced Python Constraint Support (#27)

* Python bindings now use a single-threaded version of synthesis. where ALL work takes place in the main thread and all z3 usage uses the main python z3 context. This is used to enable...

* Constraints on the system preconditions, postconditions and transitions can be expressed as python functions or lambdas. These can be given to the synthesis procedure, which will adapt them into the rust-side's closure-based constraint system. This makes it very easy to specify non-trivial arbitrary constraints from python.

* General refactors around model building to make it easier to reuse

Author: toolCHAINZ

.github/workflows/build.yml

@@ -31,4 +31,4 @@ jobs:
       - uses: dtolnay/rust-toolchain@master
         with:
           toolchain: ${{matrix.rust}}
-      - run: cargo build --all-features
+      - run: cargo build --all-features --manifest-path crackers/Cargo.toml

crackers/src/bin/crackers/main.rs

@@ -129,6 +129,8 @@ fn synthesize(config: PathBuf) -> anyhow::Result<()> {
     match result {
         Ok(res) => match res {
             DecisionResult::AssignmentFound(a) => {
+                let z3 = Context::new(&Config::new());
+                let a = a.build(&z3)?;
                 event!(Level::INFO, "Synthesis successful :)");
                 event!(Level::INFO, "{}", a)
             }

crackers/src/error.rs

@@ -37,6 +37,9 @@ pub enum CrackersError {
     LibraryConfig(#[from] GadgetLibraryConfigBuilderError),
     #[error("Invalid synthesis params")]
     SynthesisParams(#[from] SynthesisParamsBuilderError),
+    #[cfg(feature = "pyo3")]
+    #[error("Python error: {0}")]
+    PythonError(#[from] PyErr),
 }
 
 #[cfg(feature = "pyo3")]

crackers/src/synthesis/assignment_model/builder.rs

@@ -0,0 +1,121 @@
+use crate::error::CrackersError;
+use crate::gadget::Gadget;
+use crate::gadget::library::GadgetLibrary;
+use crate::synthesis::assignment_model::AssignmentModel;
+use crate::synthesis::builder::{StateConstraintGenerator, TransitionConstraintGenerator};
+use crate::synthesis::pcode_theory::pcode_assignment::PcodeAssignment;
+use jingle::JingleContext;
+use jingle::modeling::{ModeledBlock, ModeledInstruction};
+use jingle::sleigh::{ArchInfoProvider, Instruction, SpaceInfo, VarNode};
+use std::fmt::{Debug, Formatter};
+use std::sync::Arc;
+use z3::{Context, Solver};
+
+// todo: this should probably just be a struct in Jingle and we can stop with this trait nonsense
+#[derive(Clone, Debug)]
+pub struct ArchInfo {
+    spaces: Vec<SpaceInfo>,
+    default_code_space_index: usize,
+    registers: Vec<(VarNode, String)>,
+}
+
+impl From<&GadgetLibrary> for ArchInfo {
+    fn from(value: &GadgetLibrary) -> Self {
+        Self {
+            default_code_space_index: value.get_code_space_idx(),
+            registers: value
+                .get_registers()
+                .map(|(a, b)| (a.clone(), b.to_string()))
+                .collect(),
+            spaces: value.get_all_space_info().cloned().collect(),
+        }
+    }
+}
+
+impl ArchInfoProvider for ArchInfo {
+    fn get_space_info(&self, idx: usize) -> Option<&SpaceInfo> {
+        self.spaces.get(idx)
+    }
+
+    fn get_all_space_info(&self) -> impl Iterator<Item = &SpaceInfo> {
+        self.spaces.iter()
+    }
+
+    fn get_code_space_idx(&self) -> usize {
+        self.default_code_space_index
+    }
+
+    fn get_register(&self, name: &str) -> Option<&VarNode> {
+        self.registers.iter().find(|f| f.1 == name).map(|f| &f.0)
+    }
+
+    fn get_register_name(&self, location: &VarNode) -> Option<&str> {
+        self.registers
+            .iter()
+            .find(|f| f.0 == *location)
+            .map(|f| f.1.as_str())
+    }
+
+    fn get_registers(&self) -> impl Iterator<Item = (&VarNode, &str)> {
+        self.registers.iter().map(|(f, v)| (f, v.as_str()))
+    }
+}
+
+#[derive(Clone)]
+pub struct AssignmentModelBuilder {
+    pub templates: Vec<Instruction>,
+    pub gadgets: Vec<Gadget>,
+    pub preconditions: Vec<Arc<StateConstraintGenerator>>,
+    pub postconditions: Vec<Arc<StateConstraintGenerator>>,
+    pub pointer_invariants: Vec<Arc<TransitionConstraintGenerator>>,
+    pub arch_info: ArchInfo,
+}
+
+impl Debug for AssignmentModelBuilder {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("AssignmentModelBuilder")
+            .field("templates", &self.templates)
+            .field("gadgets", &self.gadgets)
+            .field("arch_info", &self.arch_info)
+            .finish()
+    }
+}
+
+impl AssignmentModelBuilder {
+    fn make_pcode_model<'ctx>(
+        &self,
+        jingle: &JingleContext<'ctx>,
+    ) -> Result<PcodeAssignment<'ctx>, CrackersError> {
+        let modeled_spec: Result<Vec<ModeledInstruction<'ctx>>, _> = self
+            .templates
+            .iter()
+            .cloned()
+            .map(|i| ModeledInstruction::new(i, jingle).map_err(CrackersError::from))
+            .collect();
+        let modeled_spec = modeled_spec?;
+        let modeled_gadgets: Result<_, _> = self
+            .gadgets
+            .iter()
+            .cloned()
+            .map(|i| i.model(jingle))
+            .collect();
+        let modeled_gadgets = modeled_gadgets?;
+        Ok(PcodeAssignment::new(
+            modeled_spec,
+            modeled_gadgets,
+            self.preconditions.clone(),
+            self.postconditions.clone(),
+            self.pointer_invariants.clone(),
+        ))
+    }
+    pub fn build<'ctx>(
+        &self,
+        z3: &'ctx Context,
+    ) -> Result<AssignmentModel<'ctx, ModeledBlock<'ctx>>, CrackersError> {
+        let jingle = JingleContext::new(z3, &self.arch_info);
+
+        let pcode_model = self.make_pcode_model(&jingle)?;
+        let s = Solver::new(jingle.z3);
+        pcode_model.check(&jingle, &s)
+    }
+}

…ackers/src/synthesis/assignment_model.rs …rs/src/synthesis/assignment_model/mod.rscrackers/src/synthesis/assignment_model.rs renamed to crackers/src/synthesis/assignment_model/mod.rs crackers/src/synthesis/assignment_model.rs renamed to crackers/src/synthesis/assignment_model/mod.rs

@@ -1,3 +1,5 @@
+pub mod builder;
+
 use std::fmt::{Display, Formatter};
 
 use jingle::modeling::{ModelingContext, State};
@@ -13,7 +15,7 @@ pub struct AssignmentModel<'ctx, T: ModelingContext<'ctx>> {
 }
 
 impl<'ctx, T: ModelingContext<'ctx>> AssignmentModel<'ctx, T> {
-    pub fn generate(model: Model<'ctx>, gadgets: Vec<T>) -> Self {
+    pub fn new(model: Model<'ctx>, gadgets: Vec<T>) -> Self {
         Self { model, gadgets }
     }
 

crackers/src/synthesis/combined.rs

@@ -1,4 +1,3 @@
-use jingle::modeling::ModeledBlock;
 use jingle::sleigh::Instruction;
 use tracing::{Level, event};
 use z3::Context;
@@ -14,7 +13,7 @@ pub struct CombinedAssignmentSynthesis<'a> {
 }
 
 impl<'a> CombinedAssignmentSynthesis<'a> {
-    pub fn decide(&mut self) -> Result<DecisionResult<'a, ModeledBlock<'a>>, CrackersError> {
+    pub fn decide(&mut self) -> Result<DecisionResult, CrackersError> {
         let mut ordering: Vec<Vec<Instruction>> = self
             .base_config
             .instructions
@@ -29,7 +28,7 @@ impl<'a> CombinedAssignmentSynthesis<'a> {
         // todo: gross hack to avoid rewriting the partitioning algorithm to be breadth-first
         ordering.sort_by(|a, b| a.len().partial_cmp(&b.len()).unwrap());
         let iter = ordering.into_iter();
-        let mut last: Option<DecisionResult<'a, ModeledBlock<'a>>> = None;
+        let mut last: Option<_> = None;
         for instructions in iter {
             // todo: filter for instruction combinations that have already been ruled out?
             // if instructions.iter().any(|i| blacklist.contains(i)) {
@@ -66,6 +65,58 @@ impl<'a> CombinedAssignmentSynthesis<'a> {
         last.ok_or(CrackersError::EmptySpecification)
     }
 
+    // gross but I don't feel like rewriting this right now
+    pub fn decide_single_threaded(&mut self) -> Result<DecisionResult, CrackersError> {
+        let mut ordering: Vec<Vec<Instruction>> = self
+            .base_config
+            .instructions
+            .partitions()
+            .map(|part| {
+                part.into_iter()
+                    .map(|instrs| Instruction::try_from(instrs).unwrap())
+                    .collect::<Vec<Instruction>>()
+            })
+            .collect();
+        // let mut blacklist = HashSet::new();
+        // todo: gross hack to avoid rewriting the partitioning algorithm to be breadth-first
+        ordering.sort_by(|a, b| a.len().partial_cmp(&b.len()).unwrap());
+        let iter = ordering.into_iter();
+        let mut last: Option<_> = None;
+        for instructions in iter {
+            // todo: filter for instruction combinations that have already been ruled out?
+            // if instructions.iter().any(|i| blacklist.contains(i)) {
+            //     continue;
+            // }
+            let mut new_config = self.base_config.clone();
+            new_config.instructions = instructions;
+            let synth = AssignmentSynthesis::new(self.z3, &new_config);
+            if let Ok(mut synth) = synth {
+                // this one constructed, let's try it
+                match synth.decide_single_threaded() {
+                    Ok(result) => {
+                        match result {
+                            DecisionResult::AssignmentFound(a) => {
+                                return Ok(DecisionResult::AssignmentFound(a));
+                            }
+                            DecisionResult::Unsat(e) => {
+                                // todo: add in bad combos here
+                                event!(Level::WARN, "{:?}", e);
+                                // e.indexes.iter().for_each(|i|{blacklist.insert(new_config.instructions[*i].clone());});
+                                last = Some(DecisionResult::Unsat(e))
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        event!(Level::ERROR, "{:?}", e)
+                    }
+                }
+            } else {
+                event!(Level::WARN, "Failed to find gadgets for partition")
+            }
+        }
+        // Only an empty specification can possibly result in this being `None`
+        last.ok_or(CrackersError::EmptySpecification)
+    }
     pub fn new(z3: &'a Context, base_config: SynthesisParams) -> Self {
         Self { z3, base_config }
     }

crackers/src/synthesis/mod.rs

@@ -1,28 +1,28 @@
 use jingle::JingleContext;
-use jingle::modeling::{ModeledBlock, ModeledInstruction, ModelingContext};
+use jingle::modeling::ModeledInstruction;
 use jingle::sleigh::Instruction;
 use std::cmp::Ordering;
 use std::sync::Arc;
 use tracing::{Level, event, instrument};
-use z3::{Config, Context, Solver};
+use z3::{Config, Context};
 
 use crate::error::CrackersError;
 use crate::error::CrackersError::EmptySpecification;
 use crate::gadget::candidates::{CandidateBuilder, Candidates};
 use crate::gadget::library::GadgetLibrary;
-use crate::synthesis::assignment_model::AssignmentModel;
+use crate::synthesis::assignment_model::builder::{ArchInfo, AssignmentModelBuilder};
 use crate::synthesis::builder::{
     StateConstraintGenerator, SynthesisParams, SynthesisSelectionStrategy,
     TransitionConstraintGenerator,
 };
 use crate::synthesis::pcode_theory::builder::PcodeTheoryBuilder;
-use crate::synthesis::pcode_theory::pcode_assignment::PcodeAssignment;
 use crate::synthesis::pcode_theory::theory_worker::TheoryWorker;
 use crate::synthesis::selection_strategy::AssignmentResult::{Failure, Success};
 use crate::synthesis::selection_strategy::OuterProblem::{OptimizeProb, SatProb};
 use crate::synthesis::selection_strategy::optimization_problem::OptimizationProblem;
 use crate::synthesis::selection_strategy::sat_problem::SatProblem;
 use crate::synthesis::selection_strategy::{OuterProblem, SelectionFailure, SelectionStrategy};
+use crate::synthesis::slot_assignments::SlotAssignments;
 
 pub mod assignment_model;
 pub mod builder;
@@ -45,8 +45,8 @@ impl PartialOrd for Decision {
 }
 
 #[derive(Debug)]
-pub enum DecisionResult<'ctx, T: ModelingContext<'ctx>> {
-    AssignmentFound(AssignmentModel<'ctx, T>),
+pub enum DecisionResult {
+    AssignmentFound(AssignmentModelBuilder),
     Unsat(SelectionFailure),
 }
 
@@ -107,11 +107,54 @@ impl<'ctx> AssignmentSynthesis<'ctx> {
         })
     }
 
+    fn make_model_builder(&self, slot_assignments: SlotAssignments) -> AssignmentModelBuilder {
+        AssignmentModelBuilder {
+            templates: self.instructions.clone(),
+            gadgets: slot_assignments.interpret_from_library(&self.candidates),
+            preconditions: self.preconditions.clone(),
+            postconditions: self.postconditions.clone(),
+            pointer_invariants: self.pointer_invariants.clone(),
+            arch_info: ArchInfo::from(self.library.as_ref()),
+        }
+    }
+
+    fn make_pcode_theory_builder(&self) -> PcodeTheoryBuilder {
+        PcodeTheoryBuilder::new(self.candidates.clone(), &self.library)
+            .with_pointer_invariants(&self.pointer_invariants)
+            .with_preconditions(&self.preconditions)
+            .with_postconditions(&self.postconditions)
+            .with_max_candidates(self.candidates_per_slot)
+            .with_templates(self.instructions.clone().into_iter())
+    }
+
+    pub fn decide_single_threaded(&mut self) -> Result<DecisionResult, CrackersError> {
+        let theory_builder = self.make_pcode_theory_builder();
+        let theory = theory_builder.build(self.z3)?;
+        loop {
+            let assignment = self.outer_problem.get_assignments()?;
+            match assignment {
+                Success(a) => {
+                    let theory_result = theory.check_assignment(&a)?;
+                    match theory_result {
+                        None => {
+                            // success
+                            return Ok(DecisionResult::AssignmentFound(self.make_model_builder(a)));
+                        }
+                        Some(conflict) => {
+                            self.outer_problem.add_theory_clauses(&conflict);
+                        }
+                    }
+                }
+                Failure(d) => return Ok(DecisionResult::Unsat(d)),
+            }
+        }
+    }
     #[instrument(skip_all)]
-    pub fn decide(&mut self) -> Result<DecisionResult<'ctx, ModeledBlock<'ctx>>, CrackersError> {
+    pub fn decide(&mut self) -> Result<DecisionResult, CrackersError> {
         let mut req_channels = vec![];
         let mut kill_senders = vec![];
-        let theory_builder = PcodeTheoryBuilder::new(self.candidates.clone(), &self.library)
+        let library = self.library.clone();
+        let theory_builder = PcodeTheoryBuilder::new(self.candidates.clone(), &library)
             .with_pointer_invariants(&self.pointer_invariants)
             .with_preconditions(&self.preconditions)
             .with_postconditions(&self.postconditions)
@@ -188,19 +231,9 @@ impl<'ctx> AssignmentSynthesis<'ctx> {
                                     x.send(()).unwrap();
                                 }
                                 kill_senders.clear();
-                                let t = theory_builder.clone();
-                                let jingle = JingleContext::new(self.z3, self.library.as_ref());
-                                let a: PcodeAssignment<'ctx> =
-                                    t.build_assignment(&jingle, response.assignment)?;
-                                event!(
-                                    Level::DEBUG,
-                                    "Workers Terminated; building and checking model"
-                                );
-                                let solver = Solver::new(self.z3);
-                                let jingle = JingleContext::new(self.z3, self.library.as_ref());
-                                let model = a.check(&jingle, &solver)?;
-                                event!(Level::DEBUG, "Model built");
-                                return Ok(DecisionResult::AssignmentFound(model));
+                                return Ok(DecisionResult::AssignmentFound(
+                                    self.make_model_builder(response.assignment),
+                                ));
                             }
                             Some(c) => {
                                 event!(

crackers/src/synthesis/pcode_theory/mod.rs

@@ -73,7 +73,6 @@ impl<'ctx, S: ModelingContext<'ctx>> PcodeTheory<'ctx, S> {
         let final_state = self.j.fresh_state();
         self.solver
             .assert(&assert_concat(self.j.z3, &self.templates)?);
-
         let mut assertions: Vec<ConjunctiveConstraint> = Vec::new();
         for (index, x) in gadgets.windows(2).enumerate() {
             let branch = Bool::fresh_const(self.j.z3, "b");

crackers/src/synthesis/pcode_theory/pcode_assignment.rs

@@ -69,7 +69,7 @@ impl<'ctx> PcodeAssignment<'ctx> {
                 let model = solver
                     .get_model()
                     .ok_or(CrackersError::ModelGenerationError)?;
-                Ok(AssignmentModel::generate(model, self.eval_trace.to_vec()))
+                Ok(AssignmentModel::new(model, self.eval_trace.to_vec()))
             }
         }
     }