chore(ACIR): prefer displaying ASSERT return_value = ...

acvm-repo/acir/src/circuit/mod.rs

@@ -5,6 +5,7 @@ pub mod brillig;
 pub mod opcodes;
 
 use crate::{
+    circuit::opcodes::display_opcode,
     native_types::{Expression, Witness},
     serialization::{deserialize_any_format, serialize_with_format_from_env},
 };
@@ -355,7 +356,8 @@ impl<F: AcirField> std::fmt::Display for Circuit<F> {
         write_witness_indices(f, &self.return_values.indices())?;
 
         for opcode in &self.opcodes {
-            writeln!(f, "{opcode}")?;
+            display_opcode(opcode, Some(&self.return_values), f)?;
+            writeln!(f)?;
         }
         Ok(())
     }

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

@@ -6,7 +6,10 @@ use super::brillig::{BrilligFunctionId, BrilligInputs, BrilligOutputs};
 pub mod function_id;
 pub use function_id::AcirFunctionId;
 
-use crate::native_types::{Expression, Witness, display_expression};
+use crate::{
+    circuit::PublicInputs,
+    native_types::{Expression, Witness, display_expression},
+};
 use acir_field::AcirField;
 use serde::{Deserialize, Serialize};
 
@@ -145,75 +148,80 @@ pub enum Opcode<F: AcirField> {
 
 impl<F: AcirField> std::fmt::Display for Opcode<F> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Opcode::AssertZero(expr) => {
-                write!(f, "ASSERT ")?;
-                display_expression(expr, true, f)
+        display_opcode(self, None, f)
+    }
+}
+
+impl<F: AcirField> std::fmt::Debug for Opcode<F> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        std::fmt::Display::fmt(self, f)
+    }
+}
+
+/// Displays an opcode, optionally using the provided return values to prefer displaying
+/// `ASSERT return_value = ...` when possible.
+pub(super) fn display_opcode<F: AcirField>(
+    opcode: &Opcode<F>,
+    return_values: Option<&PublicInputs>,
+    f: &mut std::fmt::Formatter<'_>,
+) -> std::fmt::Result {
+    match opcode {
+        Opcode::AssertZero(expr) => {
+            write!(f, "ASSERT ")?;
+            display_expression(expr, true, return_values, f)
+        }
+        Opcode::BlackBoxFuncCall(g) => std::fmt::Display::fmt(&g, f),
+        Opcode::MemoryOp { block_id, op } => {
+            let is_read = op.operation.is_zero();
+            if is_read {
+                write!(f, "READ {} = b{}[{}]", op.value, block_id.0, op.index)
+            } else {
+                write!(f, "WRITE b{}[{}] = {}", block_id.0, op.index, op.value)
             }
-            Opcode::BlackBoxFuncCall(g) => g.fmt(f),
-            Opcode::MemoryOp { block_id, op } => {
-                let is_read = op.operation.is_zero();
-                if is_read {
-                    write!(f, "READ {} = b{}[{}]", op.value, block_id.0, op.index)
-                } else {
-                    write!(f, "WRITE b{}[{}] = {}", block_id.0, op.index, op.value)
-                }
+        }
+        Opcode::MemoryInit { block_id, init, block_type: databus } => {
+            match databus {
+                BlockType::Memory => write!(f, "INIT ")?,
+                BlockType::CallData(id) => write!(f, "INIT CALLDATA {id} ")?,
+                BlockType::ReturnData => write!(f, "INIT RETURNDATA ")?,
             }
-            Opcode::MemoryInit { block_id, init, block_type: databus } => {
-                match databus {
-                    BlockType::Memory => write!(f, "INIT ")?,
-                    BlockType::CallData(id) => write!(f, "INIT CALLDATA {id} ")?,
-                    BlockType::ReturnData => write!(f, "INIT RETURNDATA ")?,
-                }
-                let witnesses =
-                    init.iter().map(|w| format!("{w}")).collect::<Vec<String>>().join(", ");
-                write!(f, "b{} = [{witnesses}]", block_id.0)
+            let witnesses = init.iter().map(|w| format!("{w}")).collect::<Vec<String>>().join(", ");
+            write!(f, "b{} = [{witnesses}]", block_id.0)
+        }
+        // We keep the display for a BrilligCall and circuit Call separate as they
+        // are distinct in their functionality and we should maintain this separation for debugging.
+        Opcode::BrilligCall { id, inputs, outputs, predicate } => {
+            write!(f, "BRILLIG CALL func: {id}, ")?;
+            if let Some(pred) = predicate {
+                write!(f, "predicate: {pred}, ")?;
             }
-            // We keep the display for a BrilligCall and circuit Call separate as they
-            // are distinct in their functionality and we should maintain this separation for debugging.
-            Opcode::BrilligCall { id, inputs, outputs, predicate } => {
-                write!(f, "BRILLIG CALL func: {id}, ")?;
-                if let Some(pred) = predicate {
-                    write!(f, "predicate: {pred}, ")?;
-                }
 
-                let inputs = inputs
-                    .iter()
-                    .map(|input| format!("{input}"))
-                    .collect::<Vec<String>>()
-                    .join(", ");
-                let outputs = outputs
-                    .iter()
-                    .map(|output| format!("{output}"))
-                    .collect::<Vec<String>>()
-                    .join(", ");
+            let inputs =
+                inputs.iter().map(|input| format!("{input}")).collect::<Vec<String>>().join(", ");
+            let outputs = outputs
+                .iter()
+                .map(|output| format!("{output}"))
+                .collect::<Vec<String>>()
+                .join(", ");
 
-                write!(f, "inputs: [{inputs}], ")?;
-                write!(f, "outputs: [{outputs}]")
+            write!(f, "inputs: [{inputs}], ")?;
+            write!(f, "outputs: [{outputs}]")
+        }
+        Opcode::Call { id, inputs, outputs, predicate } => {
+            write!(f, "CALL func: {id}, ")?;
+            if let Some(pred) = predicate {
+                write!(f, "predicate: {pred}, ")?;
             }
-            Opcode::Call { id, inputs, outputs, predicate } => {
-                write!(f, "CALL func: {id}, ")?;
-                if let Some(pred) = predicate {
-                    write!(f, "predicate: {pred}, ")?;
-                }
-                let inputs =
-                    inputs.iter().map(|w| format!("{w}")).collect::<Vec<String>>().join(", ");
-                let outputs =
-                    outputs.iter().map(|w| format!("{w}")).collect::<Vec<String>>().join(", ");
+            let inputs = inputs.iter().map(|w| format!("{w}")).collect::<Vec<String>>().join(", ");
+            let outputs =
+                outputs.iter().map(|w| format!("{w}")).collect::<Vec<String>>().join(", ");
 
-                write!(f, "inputs: [{inputs}], ")?;
-                write!(f, "outputs: [{outputs}]")
-            }
+            write!(f, "inputs: [{inputs}], ")?;
+            write!(f, "outputs: [{outputs}]")
         }
     }
 }
 
-impl<F: AcirField> std::fmt::Debug for Opcode<F> {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        std::fmt::Display::fmt(self, f)
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use acir_field::FieldElement;

acvm-repo/acir/src/native_types/expression/mod.rs

@@ -1,4 +1,4 @@
-use crate::native_types::Witness;
+use crate::{circuit::PublicInputs, native_types::Witness};
 use acir_field::AcirField;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
@@ -51,7 +51,7 @@ impl<F: AcirField> Default for Expression<F> {
 
 impl<F: AcirField> std::fmt::Display for Expression<F> {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        display_expression(self, false, f)
+        display_expression(self, false, None, f)
     }
 }
 
@@ -371,9 +371,12 @@ impl<F: AcirField> From<Witness> for Expression<F> {
 /// Displays an expression as a quadratic polynomial.
 /// If `as_equal_to_zero` is true, the expression is displayed as equaling zero,
 /// where it's tried to shown as a polynomial equal to the largest witness, if possible.
+/// If the optional `return_values` is provided, the expression is displayed preferring to show
+/// `ASSERT return_value = ...` when possible.
 pub(crate) fn display_expression<F: AcirField>(
     expr: &Expression<F>,
     as_equal_to_zero: bool,
+    return_values: Option<&PublicInputs>,
     f: &mut std::fmt::Formatter<'_>,
 ) -> std::fmt::Result {
     // This is set to an index if we show this expression "as a witness assignment", meaning
@@ -388,11 +391,21 @@ pub(crate) fn display_expression<F: AcirField>(
     // Find a linear combination with a coefficient of 1 or -1 and, if there are many,
     // keep the one with the largest witness.
     let linear_witness_one = if as_equal_to_zero {
-        expr.linear_combinations
-            .iter()
-            .enumerate()
-            .filter(|(_, (coefficient, _))| coefficient.is_one() || (-*coefficient).is_one())
-            .max_by_key(|(_, (_, witness))| witness)
+        // Prefer equating to a return value if possible
+        let linear_witness_one = return_values.and_then(|return_values| {
+            expr.linear_combinations.iter().enumerate().find(|(_, (coefficient, witness))| {
+                (coefficient.is_one() || (-*coefficient).is_one())
+                    && return_values.0.contains(witness)
+            })
+        });
+        linear_witness_one.or_else(|| {
+            // Otherwise just pick the largest witness
+            expr.linear_combinations
+                .iter()
+                .enumerate()
+                .filter(|(_, (coefficient, _))| coefficient.is_one() || (-*coefficient).is_one())
+                .max_by_key(|(_, (_, witness))| witness)
+        })
     } else {
         None
     };

compiler/noirc_evaluator/src/acir/tests/arrays.rs

@@ -30,9 +30,9 @@ fn array_set_not_mutable() {
     READ w12 = b1[w11]
     ASSERT w13 = 2
     READ w14 = b1[w13]
-    ASSERT w10 = w5
-    ASSERT w12 = w6
-    ASSERT w14 = w7
+    ASSERT w5 = w10
+    ASSERT w6 = w12
+    ASSERT w7 = w14
     ");
 }
 
@@ -63,9 +63,9 @@ fn array_set_mutable() {
     READ w12 = b0[w11]
     ASSERT w13 = 2
     READ w14 = b0[w13]
-    ASSERT w10 = w5
-    ASSERT w12 = w6
-    ASSERT w14 = w7
+    ASSERT w5 = w10
+    ASSERT w6 = w12
+    ASSERT w7 = w14
     ");
 }
 
@@ -192,9 +192,9 @@ fn generates_memory_op_for_dynamic_write() {
     READ w11 = b1[w10]
     ASSERT w12 = 2
     READ w13 = b1[w12]
-    ASSERT w9 = w4
-    ASSERT w11 = w5
-    ASSERT w13 = w6
+    ASSERT w4 = w9
+    ASSERT w5 = w11
+    ASSERT w6 = w13
     ");
 }
 
@@ -272,9 +272,9 @@ fn generates_predicated_index_and_dummy_value_for_dynamic_write() {
     READ w14 = b1[w13]
     ASSERT w15 = 2
     READ w16 = b1[w15]
-    ASSERT w12 = w5
-    ASSERT w14 = w6
-    ASSERT w16 = w7
+    ASSERT w5 = w12
+    ASSERT w6 = w14
+    ASSERT w7 = w16
     ");
 }
 

compiler/noirc_evaluator/src/acir/tests/call.rs

@@ -151,7 +151,7 @@ fn basic_nested_call(inline_type: InlineType) {
     return values: [w2]
     ASSERT w3 = w0 + 2
     CALL func: 2, predicate: 1, inputs: [w3, w1], outputs: [w4]
-    ASSERT w4 = w2
+    ASSERT w2 = w4
 
     func 2
     private parameters: [w0, w1]
@@ -187,7 +187,7 @@ fn call_with_predicate() {
     return values: [w2]
     BLACKBOX::RANGE input: w1, bits: 1
     CALL func: 1, predicate: w1, inputs: [w0], outputs: [w3]
-    ASSERT w3 = w2
+    ASSERT w2 = w3
 
     func 1
     private parameters: [w0]

compiler/noirc_evaluator/src/acir/tests/instructions.rs

@@ -124,7 +124,7 @@ fn truncate_u16_to_6_bits() {
     BLACKBOX::RANGE input: w2, bits: 10
     BLACKBOX::RANGE input: w3, bits: 6
     ASSERT w3 = w0 - 64*w2
-    ASSERT w3 = w1
+    ASSERT w1 = w3
 
     unconstrained func 0: directive_integer_quotient
     0: @10 = const u32 2
@@ -179,7 +179,7 @@ fn truncate_field_to_6_bits() {
     ASSERT 0 = -w2*w6 + 342003794872488675347600089769644923258568193756500536620284440415247007744*w6
     ASSERT w7 = w3*w6
     ASSERT w7 = 0
-    ASSERT w3 = w1
+    ASSERT w1 = w3
 
     unconstrained func 0: directive_integer_quotient
     0: @10 = const u32 2