From 3356366da0a3fa755d753d37db69a7e11662101a Mon Sep 17 00:00:00 2001
From: Tom French <tom@tomfren.ch>
Date: Wed, 29 Mar 2023 12:04:09 +0100
Subject: [PATCH] chore: handle public parameters and return values separately

---
 crates/noirc_evaluator/src/lib.rs             | 46 ++++++++++---------
 .../src/ssa/acir_gen/operations/return.rs     | 12 +----
 2 files changed, 25 insertions(+), 33 deletions(-)

diff --git a/crates/noirc_evaluator/src/lib.rs b/crates/noirc_evaluator/src/lib.rs
index 9f39dc3baa4..7063afe9bc1 100644
--- a/crates/noirc_evaluator/src/lib.rs
+++ b/crates/noirc_evaluator/src/lib.rs
@@ -13,7 +13,7 @@ use acvm::{
 };
 use errors::{RuntimeError, RuntimeErrorKind};
 use iter_extended::btree_map;
-use noirc_abi::{Abi, AbiType, AbiVisibility, MAIN_RETURN_NAME};
+use noirc_abi::{Abi, AbiType, AbiVisibility};
 use noirc_frontend::monomorphization::ast::*;
 use ssa::{node, ssa_gen::IrGenerator};
 use std::collections::{BTreeMap, BTreeSet};
@@ -38,10 +38,12 @@ pub struct Evaluator {
     // creating the private/public inputs of the ABI.
     num_witnesses_abi_len: usize,
     param_witnesses: BTreeMap<String, Vec<Witness>>,
-    // This is the list of witness indices which are linked to public inputs.
+    // This is the list of witness indices which are linked to public parameters.
     // Witnesses below `num_witnesses_abi_len` and not included in this set
-    // correspond to private inputs and must not be made public.
-    public_inputs: BTreeSet<Witness>,
+    // correspond to private parameters and must not be made public.
+    public_parameters: BTreeSet<Witness>,
+    return_values: BTreeSet<Witness>,
+
     opcodes: Vec<AcirOpcode>,
 }
 
@@ -62,30 +64,30 @@ pub fn create_circuit(
     // First evaluate the main function
     evaluator.evaluate_main_alt(program.clone(), enable_logging, show_output)?;
 
-    let witness_index = evaluator.current_witness_index();
-
-    let (parameters, return_type) = program.main_function_signature;
-
-    // TODO: remove return value from `param_witnesses` once we track public outputs
-    // see https://github.com/noir-lang/acvm/pull/56
-    let mut param_witnesses = evaluator.param_witnesses;
-    let return_witnesses = param_witnesses.remove(MAIN_RETURN_NAME).unwrap_or_default();
-    let return_witnesses_set: BTreeSet<Witness> = return_witnesses.iter().copied().collect();
-
-    let abi = Abi { parameters, param_witnesses, return_type, return_witnesses };
-
+    let Evaluator {
+        current_witness_index,
+        param_witnesses,
+        public_parameters,
+        return_values,
+        opcodes,
+        ..
+    } = evaluator;
     let optimized_circuit = acvm::compiler::compile(
         Circuit {
-            current_witness_index: witness_index,
-            opcodes: evaluator.opcodes,
-            public_parameters: PublicInputs(evaluator.public_inputs),
-            return_values: PublicInputs(return_witnesses_set),
+            current_witness_index,
+            opcodes,
+            public_parameters: PublicInputs(public_parameters),
+            return_values: PublicInputs(return_values.clone()),
         },
         np_language,
         is_opcode_supported,
     )
     .map_err(|_| RuntimeErrorKind::Spanless(String::from("produced an acvm compile error")))?;
 
+    let (parameters, return_type) = program.main_function_signature;
+    let return_witnesses: Vec<Witness> = return_values.into_iter().collect();
+    let abi = Abi { parameters, param_witnesses, return_type, return_witnesses };
+
     Ok((optimized_circuit, abi))
 }
 
@@ -101,7 +103,7 @@ impl Evaluator {
         // an intermediate variable.
         let is_intermediate_variable = witness_index.as_usize() > self.num_witnesses_abi_len;
 
-        let is_public_input = self.public_inputs.contains(&witness_index);
+        let is_public_input = self.public_parameters.contains(&witness_index);
 
         !is_intermediate_variable && !is_public_input
     }
@@ -213,7 +215,7 @@ impl Evaluator {
         };
 
         if param_visibility == &AbiVisibility::Public {
-            self.public_inputs.extend(witnesses.clone());
+            self.public_parameters.extend(witnesses.clone());
         }
         self.param_witnesses.insert(name.to_owned(), witnesses);
 
diff --git a/crates/noirc_evaluator/src/ssa/acir_gen/operations/return.rs b/crates/noirc_evaluator/src/ssa/acir_gen/operations/return.rs
index 7b99b98c45e..1e1183f13ab 100644
--- a/crates/noirc_evaluator/src/ssa/acir_gen/operations/return.rs
+++ b/crates/noirc_evaluator/src/ssa/acir_gen/operations/return.rs
@@ -1,6 +1,3 @@
-use acvm::acir::native_types::Witness;
-use noirc_abi::MAIN_RETURN_NAME;
-
 use crate::{
     errors::RuntimeErrorKind,
     ssa::{
@@ -40,7 +37,6 @@ pub(crate) fn evaluate(
             }
         };
 
-        let mut witnesses: Vec<Witness> = Vec::new();
         for object in objects {
             let witness = var_cache.get_or_compute_witness_unwrap(object, evaluator, ctx);
             // Before pushing to the public inputs, we need to check that
@@ -50,14 +46,8 @@ pub(crate) fn evaluate(
                     "we do not allow private ABI inputs to be returned as public outputs",
                 )));
             }
-            witnesses.push(witness);
+            evaluator.return_values.insert(witness);
         }
-        evaluator.public_inputs.extend(witnesses.clone());
-        evaluator
-            .param_witnesses
-            .entry(MAIN_RETURN_NAME.to_owned())
-            .or_default()
-            .append(&mut witnesses);
     }
 
     Ok(None)