From afe9c7a38bb9d4245205d3aa46d4ce23d70a5671 Mon Sep 17 00:00:00 2001
From: Tom French <15848336+TomAFrench@users.noreply.github.com>
Date: Thu, 4 Jan 2024 12:45:23 +0000
Subject: [PATCH] feat: remove range constraints from witnesses which are
 constrained to be constants (#3928)

# Description

## Problem\*

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

## Summary\*

A range constraint is a weaker constraint than directly asserting the
witness to be equal to a value (as long as that value fits within the
set number of bits). We can then remove range constraints from these
witnesses.

I've made sure to keep any range constraints which would conflict with
the constant witness assignment.

## 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.
---
 .../compiler/optimizers/redundant_range.rs    | 79 +++++++++++++------
 1 file changed, 56 insertions(+), 23 deletions(-)

diff --git a/acvm-repo/acvm/src/compiler/optimizers/redundant_range.rs b/acvm-repo/acvm/src/compiler/optimizers/redundant_range.rs
index 7b40c35960a..5d19f9629ba 100644
--- a/acvm-repo/acvm/src/compiler/optimizers/redundant_range.rs
+++ b/acvm-repo/acvm/src/compiler/optimizers/redundant_range.rs
@@ -48,25 +48,51 @@ impl RangeOptimizer {
     /// only store the fact that we have constrained it to
     /// be 16 bits.
     fn collect_ranges(circuit: &Circuit) -> BTreeMap<Witness, u32> {
-        let mut witness_to_bit_sizes = BTreeMap::new();
+        let mut witness_to_bit_sizes: BTreeMap<Witness, u32> = BTreeMap::new();
 
         for opcode in &circuit.opcodes {
-            // Extract the witness index and number of bits,
-            // if it is a range constraint
-            let (witness, num_bits) = match extract_range_opcode(opcode) {
-                Some(func_inputs) => func_inputs,
-                None => continue,
+            let Some((witness, num_bits)) = (match opcode {
+                Opcode::AssertZero(expr) => {
+                    // If the opcode is constraining a witness to be equal to a value then it can be considered
+                    // as a range opcode for the number of bits required to hold that value.
+                    if expr.is_degree_one_univariate() {
+                        let (k, witness) = expr.linear_combinations[0];
+                        let constant = expr.q_c;
+                        let witness_value = -constant / k;
+
+                        if witness_value.is_zero() {
+                            Some((witness, 0))
+                        } else {
+                            // We subtract off 1 bit from the implied witness value to give the weakest range constraint
+                            // which would be stricter than the constraint imposed by this opcode.
+                            let implied_range_constraint_bits = witness_value.num_bits() - 1;
+                            Some((witness, implied_range_constraint_bits))
+                        }
+                    } else {
+                        None
+                    }
+                }
+
+
+                Opcode::BlackBoxFuncCall(BlackBoxFuncCall::RANGE {
+                    input: FunctionInput { witness, num_bits },
+                }) => {
+                    Some((*witness, *num_bits))
+                }
+
+                _ => None,
+            }) else {
+                continue;
             };
 
             // Check if the witness has already been recorded and if the witness
             // size is more than the current one, we replace it
-            let should_replace = match witness_to_bit_sizes.get(&witness).copied() {
-                Some(old_range_bits) => old_range_bits > num_bits,
-                None => true,
-            };
-            if should_replace {
-                witness_to_bit_sizes.insert(witness, num_bits);
-            }
+            witness_to_bit_sizes
+                .entry(witness)
+                .and_modify(|old_range_bits| {
+                    *old_range_bits = std::cmp::min(*old_range_bits, num_bits);
+                })
+                .or_insert(num_bits);
         }
         witness_to_bit_sizes
     }
@@ -116,16 +142,10 @@ impl RangeOptimizer {
 /// Extract the range opcode from the `Opcode` enum
 /// Returns None, if `Opcode` is not the range opcode.
 fn extract_range_opcode(opcode: &Opcode) -> Option<(Witness, u32)> {
-    // Range constraints are blackbox function calls
-    // so we first extract the function call
-    let func_call = match opcode {
-        acir::circuit::Opcode::BlackBoxFuncCall(func_call) => func_call,
-        _ => return None,
-    };
-
-    // Skip if it is not a range constraint
-    match func_call {
-        BlackBoxFuncCall::RANGE { input } => Some((input.witness, input.num_bits)),
+    match opcode {
+        Opcode::BlackBoxFuncCall(BlackBoxFuncCall::RANGE { input }) => {
+            Some((input.witness, input.num_bits))
+        }
         _ => None,
     }
 }
@@ -246,4 +266,17 @@ mod tests {
         let (optimized_circuit, _) = optimizer.replace_redundant_ranges(acir_opcode_positions);
         assert_eq!(optimized_circuit.opcodes.len(), 5);
     }
+
+    #[test]
+    fn constant_implied_ranges() {
+        // The optimizer should use knowledge about constant witness assignments to remove range opcodes.
+        let mut circuit = test_circuit(vec![(Witness(1), 16)]);
+
+        circuit.opcodes.push(Opcode::AssertZero(Witness(1).into()));
+        let acir_opcode_positions = circuit.opcodes.iter().enumerate().map(|(i, _)| i).collect();
+        let optimizer = RangeOptimizer::new(circuit);
+        let (optimized_circuit, _) = optimizer.replace_redundant_ranges(acir_opcode_positions);
+        assert_eq!(optimized_circuit.opcodes.len(), 1);
+        assert_eq!(optimized_circuit.opcodes[0], Opcode::AssertZero(Witness(1).into()));
+    }
 }