feat: backpropagate constants in ACIR during optimization (#3926)

# Description

## Problem\*

Resolves <!-- Link to GitHub Issue -->

## Summary\*

This is a mildy bruteforce-y optimisation method where we just literally
attempt to execute the circuit backwards. Any witnesses which we can
determine from this can just be written into the circuit directly.

A lot of the complexity here comes from the fact that memory opcodes,
etc. require witnesses to be unassigned at the point at which the opcode
is encountered so we need to "forget" certain witnesses so that we don't
optimise them away.

Draft as I'm just pushing this up to track effects.

## Additional Context



## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** Documentation to be submitted in a separate
PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: kevaundray <kevtheappdev@gmail.com>
Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>

Author: 3 people

acvm-repo/acir/src/circuit/opcodes/black_box_function_call.rs

@@ -217,8 +217,10 @@ impl BlackBoxFuncCall {
             | BlackBoxFuncCall::PedersenCommitment { inputs, .. }
             | BlackBoxFuncCall::PedersenHash { inputs, .. }
             | BlackBoxFuncCall::BigIntFromLeBytes { inputs, .. }
-            | BlackBoxFuncCall::Poseidon2Permutation { inputs, .. }
-            | BlackBoxFuncCall::Sha256Compression { inputs, .. } => inputs.to_vec(),
+            | BlackBoxFuncCall::Poseidon2Permutation { inputs, .. } => inputs.to_vec(),
+            BlackBoxFuncCall::Sha256Compression { inputs, hash_values, .. } => {
+                inputs.iter().chain(hash_values).copied().collect()
+            }
             BlackBoxFuncCall::AND { lhs, rhs, .. } | BlackBoxFuncCall::XOR { lhs, rhs, .. } => {
                 vec![*lhs, *rhs]
             }

acvm-repo/acvm/src/compiler/optimizers/constant_backpropagation.rs

@@ -0,0 +1,331 @@
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+
+use crate::{
+    compiler::optimizers::GeneralOptimizer,
+    pwg::{
+        arithmetic::ExpressionSolver, blackbox::solve_range_opcode, directives::solve_directives,
+        BrilligSolver, BrilligSolverStatus,
+    },
+};
+use acir::{
+    circuit::{
+        brillig::{Brillig, BrilligInputs, BrilligOutputs},
+        directives::Directive,
+        opcodes::BlackBoxFuncCall,
+        Circuit, Opcode,
+    },
+    native_types::{Expression, Witness, WitnessMap},
+};
+use acvm_blackbox_solver::StubbedBlackBoxSolver;
+
+/// `ConstantBackpropagationOptimizer` will attempt to determine any constant witnesses within the program.
+/// It does this by attempting to solve the program without any inputs (i.e. using an empty witness map),
+/// any values which it can determine are then enforced to be constant values.
+///
+/// The optimizer will then replace any witnesses wherever they appear within the circuit with these constant values.
+/// This is repeated until the circuit stabilizes.
+pub(crate) struct ConstantBackpropagationOptimizer {
+    circuit: Circuit,
+}
+
+impl ConstantBackpropagationOptimizer {
+    /// Creates a new `ConstantBackpropagationOptimizer`
+    pub(crate) fn new(circuit: Circuit) -> Self {
+        Self { circuit }
+    }
+
+    fn gather_known_witnesses(&self) -> (WitnessMap, BTreeSet<Witness>) {
+        // We do not want to affect the circuit's interface so avoid optimizing away these witnesses.
+        let mut required_witnesses: BTreeSet<Witness> = self
+            .circuit
+            .private_parameters
+            .union(&self.circuit.public_parameters.0)
+            .chain(&self.circuit.return_values.0)
+            .copied()
+            .collect();
+
+        for opcode in &self.circuit.opcodes {
+            match &opcode {
+                Opcode::BlackBoxFuncCall(func_call) => {
+                    required_witnesses.extend(
+                        func_call.get_inputs_vec().into_iter().map(|func_input| func_input.witness),
+                    );
+                    required_witnesses.extend(func_call.get_outputs_vec());
+                }
+
+                Opcode::MemoryInit { init, .. } => {
+                    required_witnesses.extend(init);
+                }
+
+                Opcode::MemoryOp { op, .. } => {
+                    required_witnesses.insert(op.index.to_witness().unwrap());
+                    required_witnesses.insert(op.value.to_witness().unwrap());
+                }
+
+                _ => (),
+            };
+        }
+
+        let mut known_witnesses = WitnessMap::new();
+        for opcode in self.circuit.opcodes.iter().rev() {
+            if let Opcode::AssertZero(expr) = opcode {
+                let solve_result = ExpressionSolver::solve(&mut known_witnesses, expr);
+                // It doesn't matter what the result is. We expect most opcodes to not be solved successfully so we discard errors.
+                // At the same time, if the expression can be solved then we track this by the updates to `known_witnesses`
+                drop(solve_result);
+            }
+        }
+
+        // We want to retain any references to required witnesses so we "forget" these assignments.
+        let known_witnesses: BTreeMap<_, _> = known_witnesses
+            .into_iter()
+            .filter(|(witness, _)| !required_witnesses.contains(witness))
+            .collect();
+
+        (known_witnesses.into(), required_witnesses)
+    }
+
+    /// Returns a `Circuit` where with any constant witnesses replaced with the constant they resolve to.
+    #[tracing::instrument(level = "trace", skip_all)]
+    pub(crate) fn backpropagate_constants(
+        circuit: Circuit,
+        order_list: Vec<usize>,
+    ) -> (Circuit, Vec<usize>) {
+        let old_circuit_size = circuit.opcodes.len();
+
+        let optimizer = Self::new(circuit);
+        let (circuit, order_list) = optimizer.backpropagate_constants_iteration(order_list);
+
+        let new_circuit_size = circuit.opcodes.len();
+        if new_circuit_size < old_circuit_size {
+            Self::backpropagate_constants(circuit, order_list)
+        } else {
+            (circuit, order_list)
+        }
+    }
+
+    /// Applies a single round of constant backpropagation to a `Circuit`.
+    pub(crate) fn backpropagate_constants_iteration(
+        mut self,
+        order_list: Vec<usize>,
+    ) -> (Circuit, Vec<usize>) {
+        let (mut known_witnesses, required_witnesses) = self.gather_known_witnesses();
+
+        let opcodes = std::mem::take(&mut self.circuit.opcodes);
+
+        fn remap_expression(known_witnesses: &WitnessMap, expression: Expression) -> Expression {
+            GeneralOptimizer::optimize(ExpressionSolver::evaluate(&expression, known_witnesses))
+        }
+
+        let mut new_order_list = Vec::with_capacity(order_list.len());
+        let mut new_opcodes = Vec::with_capacity(opcodes.len());
+        for (idx, opcode) in opcodes.into_iter().enumerate() {
+            let new_opcode = match opcode {
+                Opcode::AssertZero(expression) => {
+                    let new_expr = remap_expression(&known_witnesses, expression);
+                    if new_expr.is_zero() {
+                        continue;
+                    }
+
+                    // Attempt to solve the opcode to see if we can determine the value of any witnesses in the expression.
+                    // We only do this _after_ we apply any simplifications to create the new opcode as we want to
+                    // keep the constraint on the witness which we are solving for here.
+                    let solve_result = ExpressionSolver::solve(&mut known_witnesses, &new_expr);
+                    // It doesn't matter what the result is. We expect most opcodes to not be solved successfully so we discard errors.
+                    // At the same time, if the expression can be solved then we track this by the updates to `known_witnesses`
+                    drop(solve_result);
+
+                    Opcode::AssertZero(new_expr)
+                }
+                Opcode::Brillig(brillig) => {
+                    let remapped_inputs = brillig
+                        .inputs
+                        .into_iter()
+                        .map(|input| match input {
+                            BrilligInputs::Single(expr) => {
+                                BrilligInputs::Single(remap_expression(&known_witnesses, expr))
+                            }
+                            BrilligInputs::Array(expr_array) => {
+                                let new_input: Vec<_> = expr_array
+                                    .into_iter()
+                                    .map(|expr| remap_expression(&known_witnesses, expr))
+                                    .collect();
+
+                                BrilligInputs::Array(new_input)
+                            }
+                            input @ BrilligInputs::MemoryArray(_) => input,
+                        })
+                        .collect();
+
+                    let remapped_predicate = brillig
+                        .predicate
+                        .map(|predicate| remap_expression(&known_witnesses, predicate));
+
+                    let new_brillig = Brillig {
+                        inputs: remapped_inputs,
+                        predicate: remapped_predicate,
+                        ..brillig
+                    };
+
+                    let brillig_output_is_required_witness =
+                        new_brillig.outputs.iter().any(|output| match output {
+                            BrilligOutputs::Simple(witness) => required_witnesses.contains(witness),
+                            BrilligOutputs::Array(witness_array) => witness_array
+                                .iter()
+                                .any(|witness| required_witnesses.contains(witness)),
+                        });
+
+                    if brillig_output_is_required_witness {
+                        // If one of the brillig opcode's outputs is a required witness then we can't remove the opcode. In this case we can't replace
+                        // all of the uses of this witness with the calculated constant so we'll be attempting to use an uninitialized witness.
+                        //
+                        // We then do not attempt execution of this opcode and just simplify the inputs.
+                        Opcode::Brillig(new_brillig)
+                    } else if let Ok(mut solver) = BrilligSolver::new(
+                        &known_witnesses,
+                        &HashMap::new(),
+                        &new_brillig,
+                        &StubbedBlackBoxSolver,
+                        idx,
+                    ) {
+                        match solver.solve() {
+                            Ok(BrilligSolverStatus::Finished) => {
+                                // Write execution outputs
+                                match solver.finalize(&mut known_witnesses, &new_brillig) {
+                                    Ok(()) => {
+                                        // If we've managed to execute the brillig opcode at compile time, we can now just write in the
+                                        // results as constants for the rest of the circuit.
+                                        continue;
+                                    }
+                                    _ => Opcode::Brillig(new_brillig),
+                                }
+                            }
+                            Ok(BrilligSolverStatus::InProgress) => unreachable!(
+                                "Solver should either finish, block on foreign call, or error."
+                            ),
+                            Ok(BrilligSolverStatus::ForeignCallWait(_)) | Err(_) => {
+                                Opcode::Brillig(new_brillig)
+                            }
+                        }
+                    } else {
+                        Opcode::Brillig(new_brillig)
+                    }
+                }
+
+                Opcode::Directive(Directive::ToLeRadix { a, b, radix }) => {
+                    if b.iter().all(|output| known_witnesses.contains_key(output)) {
+                        continue;
+                    } else if b.iter().any(|witness| required_witnesses.contains(witness)) {
+                        // If one of the brillig opcode's outputs is a required witness then we can't remove the opcode. In this case we can't replace
+                        // all of the uses of this witness with the calculated constant so we'll be attempting to use an uninitialized witness.
+                        //
+                        // We then do not attempt execution of this opcode and just simplify the inputs.
+                        Opcode::Directive(Directive::ToLeRadix {
+                            a: remap_expression(&known_witnesses, a),
+                            b,
+                            radix,
+                        })
+                    } else {
+                        let directive = Directive::ToLeRadix {
+                            a: remap_expression(&known_witnesses, a),
+                            b,
+                            radix,
+                        };
+                        let result = solve_directives(&mut known_witnesses, &directive);
+
+                        match result {
+                            Ok(()) => continue,
+                            Err(_) => Opcode::Directive(directive),
+                        }
+                    }
+                }
+
+                Opcode::BlackBoxFuncCall(BlackBoxFuncCall::RANGE { input }) => {
+                    if solve_range_opcode(&known_witnesses, &input).is_ok() {
+                        continue;
+                    } else {
+                        opcode
+                    }
+                }
+
+                Opcode::BlackBoxFuncCall(_)
+                | Opcode::MemoryOp { .. }
+                | Opcode::MemoryInit { .. } => opcode,
+            };
+
+            new_opcodes.push(new_opcode);
+            new_order_list.push(order_list[idx]);
+        }
+
+        self.circuit.opcodes = new_opcodes;
+
+        (self.circuit, new_order_list)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeSet;
+
+    use crate::compiler::optimizers::constant_backpropagation::ConstantBackpropagationOptimizer;
+    use acir::{
+        brillig::MemoryAddress,
+        circuit::{
+            brillig::{Brillig, BrilligOutputs},
+            opcodes::{BlackBoxFuncCall, FunctionInput},
+            Circuit, ExpressionWidth, Opcode, PublicInputs,
+        },
+        native_types::Witness,
+    };
+    use brillig_vm::brillig::Opcode as BrilligOpcode;
+
+    fn test_circuit(opcodes: Vec<Opcode>) -> Circuit {
+        Circuit {
+            current_witness_index: 1,
+            expression_width: ExpressionWidth::Bounded { width: 3 },
+            opcodes,
+            private_parameters: BTreeSet::new(),
+            public_parameters: PublicInputs::default(),
+            return_values: PublicInputs::default(),
+            assert_messages: Default::default(),
+            recursive: false,
+        }
+    }
+
+    #[test]
+    fn retain_brillig_with_required_witness_outputs() {
+        let brillig_opcode = Opcode::Brillig(Brillig {
+            inputs: Vec::new(),
+            outputs: vec![BrilligOutputs::Simple(Witness(1))],
+            bytecode: vec![
+                BrilligOpcode::Const {
+                    destination: MemoryAddress(0),
+                    bit_size: 32,
+                    value: 1u128.into(),
+                },
+                BrilligOpcode::Stop { return_data_offset: 0, return_data_size: 1 },
+            ],
+            predicate: None,
+        });
+        let blackbox_opcode = Opcode::BlackBoxFuncCall(BlackBoxFuncCall::AND {
+            lhs: FunctionInput { witness: Witness(1), num_bits: 64 },
+            rhs: FunctionInput { witness: Witness(2), num_bits: 64 },
+            output: Witness(3),
+        });
+
+        let opcodes = vec![brillig_opcode, blackbox_opcode];
+        // The optimizer should keep the lowest bit size range constraint
+        let circuit = test_circuit(opcodes);
+        let acir_opcode_positions = circuit.opcodes.iter().enumerate().map(|(i, _)| i).collect();
+        let optimizer = ConstantBackpropagationOptimizer::new(circuit);
+
+        let (optimized_circuit, _) =
+            optimizer.backpropagate_constants_iteration(acir_opcode_positions);
+
+        assert_eq!(
+            optimized_circuit.opcodes.len(),
+            2,
+            "The brillig opcode should not be removed as the output is needed as a witness"
+        );
+    }
+}

acvm-repo/acvm/src/compiler/optimizers/general.rs

@@ -13,7 +13,8 @@ impl GeneralOptimizer {
     pub(crate) fn optimize(opcode: Expression) -> Expression {
         // XXX: Perhaps this optimization can be done on the fly
         let opcode = remove_zero_coefficients(opcode);
-        simplify_mul_terms(opcode)
+        let opcode = simplify_mul_terms(opcode);
+        simplify_linear_terms(opcode)
     }
 }
 
@@ -42,3 +43,20 @@ fn simplify_mul_terms(mut gate: Expression) -> Expression {
     gate.mul_terms = hash_map.into_iter().map(|((w_l, w_r), scale)| (scale, w_l, w_r)).collect();
     gate
 }
+
+// Simplifies all linear terms with the same variables
+fn simplify_linear_terms(mut gate: Expression) -> Expression {
+    let mut hash_map: IndexMap<Witness, FieldElement> = IndexMap::new();
+
+    // Canonicalize the ordering of the terms, lets just order by variable name
+    for (scale, witness) in gate.linear_combinations.into_iter() {
+        *hash_map.entry(witness).or_insert_with(FieldElement::zero) += scale;
+    }
+
+    gate.linear_combinations = hash_map
+        .into_iter()
+        .filter(|(_, scale)| scale != &FieldElement::zero())
+        .map(|(witness, scale)| (scale, witness))
+        .collect();
+    gate
+}