feat!: Remove solve from PWG trait & introduce separate solvers for…

… each blackbox (#257)

* feat!: Remove solve from PWG trait & introduce separate solvers for each blackbox

* clippy

* code review

acvm/src/pwg.rs

@@ -1,19 +1,23 @@
 // Re-usable methods that backends can use to implement their PWG
 
-use crate::{OpcodeNotSolvable, OpcodeResolutionError};
+use crate::{OpcodeNotSolvable, OpcodeResolutionError, PartialWitnessGenerator};
 use acir::{
+    circuit::opcodes::{Opcode, OracleData},
     native_types::{Expression, Witness},
     FieldElement,
 };
 use std::collections::BTreeMap;
 
-use self::arithmetic::ArithmeticSolver;
+use self::{
+    arithmetic::ArithmeticSolver, block::Blocks, directives::solve_directives, oracle::OracleSolver,
+};
 
 // arithmetic
 pub mod arithmetic;
 // Directives
 pub mod directives;
 // black box functions
+mod blackbox;
 pub mod block;
 pub mod hash;
 pub mod logic;
@@ -22,6 +26,109 @@ pub mod range;
 pub mod signature;
 pub mod sorting;
 
+#[derive(Debug, PartialEq)]
+pub enum PartialWitnessGeneratorStatus {
+    /// All opcodes have been solved.
+    Solved,
+
+    /// The `PartialWitnessGenerator` has encountered a request for [oracle data][Opcode::Oracle].
+    ///
+    /// The caller must resolve these opcodes externally and insert the results into the intermediate witness.
+    /// Once this is done, the `PartialWitnessGenerator` can be restarted to solve the remaining opcodes.
+    RequiresOracleData { required_oracle_data: Vec<OracleData>, unsolved_opcodes: Vec<Opcode> },
+}
+
+#[derive(Debug, PartialEq)]
+pub enum OpcodeResolution {
+    /// The opcode is resolved
+    Solved,
+    /// The opcode is not solvable
+    Stalled(OpcodeNotSolvable),
+    /// The opcode is not solvable but could resolved some witness
+    InProgress,
+}
+
+pub fn solve(
+    backend: &impl PartialWitnessGenerator,
+    initial_witness: &mut BTreeMap<Witness, FieldElement>,
+    blocks: &mut Blocks,
+    mut opcode_to_solve: Vec<Opcode>,
+) -> Result<PartialWitnessGeneratorStatus, OpcodeResolutionError> {
+    let mut unresolved_opcodes: Vec<Opcode> = Vec::new();
+    let mut unresolved_oracles: Vec<OracleData> = Vec::new();
+    while !opcode_to_solve.is_empty() || !unresolved_oracles.is_empty() {
+        unresolved_opcodes.clear();
+        let mut stalled = true;
+        let mut opcode_not_solvable = None;
+        for opcode in &opcode_to_solve {
+            let mut solved_oracle_data = None;
+            let resolution = match opcode {
+                Opcode::Arithmetic(expr) => ArithmeticSolver::solve(initial_witness, expr),
+                Opcode::BlackBoxFuncCall(bb_func) => {
+                    blackbox::solve(backend, initial_witness, bb_func)
+                }
+                Opcode::Directive(directive) => solve_directives(initial_witness, directive),
+                Opcode::Block(block) | Opcode::ROM(block) | Opcode::RAM(block) => {
+                    blocks.solve(block.id, &block.trace, initial_witness)
+                }
+                Opcode::Oracle(data) => {
+                    let mut data_clone = data.clone();
+                    let result = OracleSolver::solve(initial_witness, &mut data_clone)?;
+                    solved_oracle_data = Some(data_clone);
+                    Ok(result)
+                }
+            };
+            match resolution {
+                Ok(OpcodeResolution::Solved) => {
+                    stalled = false;
+                }
+                Ok(OpcodeResolution::InProgress) => {
+                    stalled = false;
+                    // InProgress Oracles must be externally re-solved
+                    if let Some(oracle) = solved_oracle_data {
+                        unresolved_oracles.push(oracle);
+                    } else {
+                        unresolved_opcodes.push(opcode.clone());
+                    }
+                }
+                Ok(OpcodeResolution::Stalled(not_solvable)) => {
+                    if opcode_not_solvable.is_none() {
+                        // we keep track of the first unsolvable opcode
+                        opcode_not_solvable = Some(not_solvable);
+                    }
+                    // We push those opcodes not solvable to the back as
+                    // it could be because the opcodes are out of order, i.e. this assignment
+                    // relies on a later opcodes' results
+                    unresolved_opcodes.push(match solved_oracle_data {
+                        Some(oracle_data) => Opcode::Oracle(oracle_data),
+                        None => opcode.clone(),
+                    });
+                }
+                Err(OpcodeResolutionError::OpcodeNotSolvable(_)) => {
+                    unreachable!("ICE - Result should have been converted to GateResolution")
+                }
+                Err(err) => return Err(err),
+            }
+        }
+        // We have oracles that must be externally resolved
+        if !unresolved_oracles.is_empty() {
+            return Ok(PartialWitnessGeneratorStatus::RequiresOracleData {
+                required_oracle_data: unresolved_oracles,
+                unsolved_opcodes: unresolved_opcodes,
+            });
+        }
+        // We are stalled because of an opcode being bad
+        if stalled && !unresolved_opcodes.is_empty() {
+            return Err(OpcodeResolutionError::OpcodeNotSolvable(
+                opcode_not_solvable
+                    .expect("infallible: cannot be stalled and None at the same time"),
+            ));
+        }
+        std::mem::swap(&mut opcode_to_solve, &mut unresolved_opcodes);
+    }
+    Ok(PartialWitnessGeneratorStatus::Solved)
+}
+
 // Returns the concrete value for a particular witness
 // If the witness has no assignment, then
 // an error is returned

acvm/src/pwg/arithmetic.rs

@@ -4,7 +4,7 @@ use acir::{
 };
 use std::collections::BTreeMap;
 
-use crate::{OpcodeNotSolvable, OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::OpcodeResolution, OpcodeNotSolvable, OpcodeResolutionError};
 
 /// An Arithmetic solver will take a Circuit's arithmetic gates with witness assignments
 /// and create the other witness variables

acvm/src/pwg/blackbox.rs

@@ -0,0 +1,94 @@
+use std::collections::BTreeMap;
+
+use acir::{
+    circuit::opcodes::{BlackBoxFuncCall, FunctionInput},
+    native_types::Witness,
+    BlackBoxFunc, FieldElement,
+};
+
+use crate::{OpcodeNotSolvable, OpcodeResolutionError, PartialWitnessGenerator};
+
+use super::OpcodeResolution;
+
+/// Check if all of the inputs to the function have assignments
+///
+/// Returns the first missing assignment if any are missing
+fn first_missing_assignment(
+    witness_assignments: &BTreeMap<Witness, FieldElement>,
+    inputs: &[FunctionInput],
+) -> Option<Witness> {
+    inputs.iter().find_map(|input| {
+        if witness_assignments.contains_key(&input.witness) {
+            None
+        } else {
+            Some(input.witness)
+        }
+    })
+}
+
+/// Check if all of the inputs to the function have assignments
+fn contains_all_inputs(
+    witness_assignments: &BTreeMap<Witness, FieldElement>,
+    inputs: &[FunctionInput],
+) -> bool {
+    inputs.iter().all(|input| witness_assignments.contains_key(&input.witness))
+}
+
+pub(crate) fn solve(
+    backend: &impl PartialWitnessGenerator,
+    initial_witness: &mut BTreeMap<Witness, FieldElement>,
+    bb_func: &BlackBoxFuncCall,
+) -> Result<OpcodeResolution, OpcodeResolutionError> {
+    match bb_func {
+        BlackBoxFuncCall { inputs, .. } if !contains_all_inputs(initial_witness, inputs) => {
+            if let Some(unassigned_witness) = first_missing_assignment(initial_witness, inputs) {
+                Ok(OpcodeResolution::Stalled(OpcodeNotSolvable::MissingAssignment(
+                    unassigned_witness.0,
+                )))
+            } else {
+                // This only exists because Rust won't let us bind in a pattern guard.
+                // See https://github.com/rust-lang/rust/issues/51114
+                unreachable!("Only reachable if the blackbox is stalled")
+            }
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::AES, inputs, outputs } => {
+            backend.aes(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::AND, inputs, outputs } => {
+            backend.and(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::XOR, inputs, outputs } => {
+            backend.xor(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::RANGE, inputs, outputs } => {
+            backend.range(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::SHA256, inputs, outputs } => {
+            backend.sha256(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::Blake2s, inputs, outputs } => {
+            backend.blake2s(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::ComputeMerkleRoot, inputs, outputs } => {
+            backend.compute_merkle_root(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::SchnorrVerify, inputs, outputs } => {
+            backend.schnorr_verify(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::Pedersen, inputs, outputs } => {
+            backend.pedersen(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::HashToField128Security, inputs, outputs } => {
+            backend.hash_to_field128_security(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::EcdsaSecp256k1, inputs, outputs } => {
+            backend.ecdsa_secp256k1(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::FixedBaseScalarMul, inputs, outputs } => {
+            backend.fixed_base_scalar_mul(initial_witness, inputs, outputs)
+        }
+        BlackBoxFuncCall { name: BlackBoxFunc::Keccak256, inputs, outputs } => {
+            backend.keccak256(initial_witness, inputs, outputs)
+        }
+    }
+}

acvm/src/pwg/block.rs

@@ -6,7 +6,7 @@ use acir::{
     FieldElement,
 };
 
-use crate::{OpcodeNotSolvable, OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::OpcodeResolution, OpcodeNotSolvable, OpcodeResolutionError};
 
 use super::{
     arithmetic::{ArithmeticSolver, GateStatus},

acvm/src/pwg/directives.rs

@@ -8,7 +8,7 @@ use acir::{
 use num_bigint::BigUint;
 use num_traits::Zero;
 
-use crate::{OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::OpcodeResolution, OpcodeResolutionError};
 
 use super::{get_value, insert_value, sorting::route, witness_to_value};
 

acvm/src/pwg/hash.rs

@@ -4,7 +4,7 @@ use sha2::Sha256;
 use sha3::Keccak256;
 use std::collections::BTreeMap;
 
-use crate::{OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::OpcodeResolution, OpcodeResolutionError};
 
 use super::{insert_value, witness_to_value};
 

acvm/src/pwg/logic.rs

@@ -1,5 +1,5 @@
 use super::{insert_value, witness_to_value};
-use crate::{OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::OpcodeResolution, OpcodeResolutionError};
 use acir::{circuit::opcodes::BlackBoxFuncCall, native_types::Witness, BlackBoxFunc, FieldElement};
 use std::collections::BTreeMap;
 

acvm/src/pwg/oracle.rs

@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
 
 use acir::{circuit::opcodes::OracleData, native_types::Witness, FieldElement};
 
-use crate::{OpcodeNotSolvable, OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::OpcodeResolution, OpcodeNotSolvable, OpcodeResolutionError};
 
 use super::{arithmetic::ArithmeticSolver, insert_value};
 

acvm/src/pwg/range.rs

@@ -1,4 +1,4 @@
-use crate::{pwg::witness_to_value, OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::witness_to_value, pwg::OpcodeResolution, OpcodeResolutionError};
 use acir::{circuit::opcodes::BlackBoxFuncCall, native_types::Witness, BlackBoxFunc, FieldElement};
 use std::collections::BTreeMap;
 

acvm/src/pwg/signature/ecdsa.rs

@@ -1,7 +1,7 @@
 use acir::{circuit::opcodes::BlackBoxFuncCall, native_types::Witness, FieldElement};
 use std::collections::BTreeMap;
 
-use crate::{pwg::witness_to_value, OpcodeResolution, OpcodeResolutionError};
+use crate::{pwg::witness_to_value, pwg::OpcodeResolution, OpcodeResolutionError};
 
 pub fn secp256k1_prehashed(
     initial_witness: &mut BTreeMap<Witness, FieldElement>,