feat(abi)!: add explicit return type field to ABI. (#865)

* feat: add an explicit `return_type` field to ABI

* chore: inline `check_for_unexpected_inputs`

* chore: inline `self.return_witnesses` and remove borrow

* chore: add necessary errors

* chore: replace `matches` with `match` statement!

* chore: clippy

* chore: remove filtering of return dummy parameter

* chore: use `MAIN_RETURN_NAME` instead of "return"

* chore: add proper error variants for return value issues

* chore: improve comment

crates/nargo/src/cli/execute_cmd.rs

@@ -4,7 +4,7 @@ use acvm::acir::native_types::Witness;
 use acvm::PartialWitnessGenerator;
 use clap::Args;
 use noirc_abi::input_parser::{Format, InputValue};
-use noirc_abi::{InputMap, WitnessMap, MAIN_RETURN_NAME};
+use noirc_abi::{InputMap, WitnessMap};
 use noirc_driver::CompiledProgram;
 
 use super::NargoConfig;
@@ -67,8 +67,7 @@ fn execute_with_path<P: AsRef<Path>>(
     let solved_witness = execute_program(&compiled_program, &inputs_map)?;
 
     let public_abi = compiled_program.abi.public_abi();
-    let public_inputs = public_abi.decode(&solved_witness)?;
-    let return_value = public_inputs.get(MAIN_RETURN_NAME).cloned();
+    let (_, return_value) = public_abi.decode(&solved_witness)?;
 
     Ok((return_value, solved_witness))
 }
@@ -77,7 +76,7 @@ pub(crate) fn execute_program(
     compiled_program: &CompiledProgram,
     inputs_map: &InputMap,
 ) -> Result<WitnessMap, CliError> {
-    let mut solved_witness = compiled_program.abi.encode(inputs_map, true)?;
+    let mut solved_witness = compiled_program.abi.encode(inputs_map, None)?;
 
     let backend = crate::backends::ConcreteBackend;
     backend.solve(&mut solved_witness, compiled_program.circuit.opcodes.clone())?;

crates/nargo/src/cli/prove_cmd.rs

@@ -2,7 +2,7 @@ use std::path::{Path, PathBuf};
 
 use acvm::ProofSystemCompiler;
 use clap::Args;
-use noirc_abi::input_parser::Format;
+use noirc_abi::{input_parser::Format, MAIN_RETURN_NAME};
 
 use super::{
     create_named_dir, fetch_pk_and_vk, read_inputs_from_file, write_inputs_to_file, write_to_file,
@@ -92,16 +92,30 @@ pub fn prove_with_path<P: AsRef<Path>>(
 
     // Write public inputs into Verifier.toml
     let public_abi = compiled_program.abi.clone().public_abi();
-    let public_inputs = public_abi.decode(&solved_witness)?;
-    write_inputs_to_file(&public_inputs, &program_dir, VERIFIER_INPUT_FILE, Format::Toml)?;
+    let (public_inputs, return_value) = public_abi.decode(&solved_witness)?;
+
+    if let Some(return_value) = return_value.clone() {
+        // Insert return value into public inputs so it's written to file.
+        let mut public_inputs_with_return = public_inputs.clone();
+        public_inputs_with_return.insert(MAIN_RETURN_NAME.to_owned(), return_value);
+        write_inputs_to_file(
+            &public_inputs_with_return,
+            &program_dir,
+            VERIFIER_INPUT_FILE,
+            Format::Toml,
+        )?;
+    } else {
+        write_inputs_to_file(&public_inputs, &program_dir, VERIFIER_INPUT_FILE, Format::Toml)?;
+    }
 
     let backend = crate::backends::ConcreteBackend;
     let proof =
         backend.prove_with_pk(compiled_program.circuit.clone(), solved_witness, proving_key);
 
     println!("Proof successfully created");
     if check_proof {
-        let valid_proof = verify_proof(compiled_program, public_inputs, &proof, verification_key)?;
+        let valid_proof =
+            verify_proof(compiled_program, public_inputs, return_value, &proof, verification_key)?;
         println!("Proof verified : {valid_proof}");
         if !valid_proof {
             return Err(CliError::Generic("Could not verify generated proof".to_owned()));

crates/nargo/src/cli/verify_cmd.rs

@@ -8,7 +8,8 @@ use crate::{
 };
 use acvm::{FieldElement, ProofSystemCompiler};
 use clap::Args;
-use noirc_abi::input_parser::Format;
+use noirc_abi::input_parser::{Format, InputValue};
+use noirc_abi::MAIN_RETURN_NAME;
 use noirc_driver::CompiledProgram;
 use std::{collections::BTreeMap, path::Path};
 
@@ -64,20 +65,22 @@ pub fn verify_with_path<P: AsRef<Path>>(
     let (_, verification_key) =
         fetch_pk_and_vk(compiled_program.circuit.clone(), circuit_build_path, false, true)?;
 
-    let mut public_inputs_map: InputMap = BTreeMap::new();
-
     // Load public inputs (if any) from `VERIFIER_INPUT_FILE`.
     let public_abi = compiled_program.abi.clone().public_abi();
-    let num_pub_params = public_abi.num_parameters();
-    if num_pub_params != 0 {
+    let (public_inputs_map, return_value) = if public_abi.has_public_inputs() {
         let current_dir = program_dir;
-        public_inputs_map =
+        let mut public_inputs_map =
             read_inputs_from_file(current_dir, VERIFIER_INPUT_FILE, Format::Toml, &public_abi)?;
-    }
+        let return_value = public_inputs_map.remove(MAIN_RETURN_NAME);
+        (public_inputs_map, return_value)
+    } else {
+        (BTreeMap::new(), None)
+    };
 
     let valid_proof = verify_proof(
         compiled_program,
         public_inputs_map,
+        return_value,
         &load_hex_data(proof_path)?,
         verification_key,
     )?;
@@ -88,11 +91,12 @@ pub fn verify_with_path<P: AsRef<Path>>(
 pub(crate) fn verify_proof(
     compiled_program: CompiledProgram,
     public_inputs_map: InputMap,
+    return_value: Option<InputValue>,
     proof: &[u8],
     verification_key: Vec<u8>,
 ) -> Result<bool, CliError> {
     let public_abi = compiled_program.abi.public_abi();
-    let public_inputs = public_abi.encode(&public_inputs_map, false)?;
+    let public_inputs = public_abi.encode(&public_inputs_map, return_value)?;
 
     let public_inputs_vec: Vec<FieldElement> = public_inputs.values().copied().collect();
 

crates/noirc_abi/src/errors.rs

@@ -42,4 +42,10 @@ pub enum AbiError {
         "Could not read witness value at index {witness_index:?} (required for parameter \"{name}\")"
     )]
     MissingParamWitnessValue { name: String, witness_index: Witness },
+    #[error("Attempted to write to witness index {0:?} but it is already initialized to a different value")]
+    InconsistentWitnessAssignment(Witness),
+    #[error("The return value is expected to be a {return_type:?} but found incompatible value {value:?}")]
+    ReturnTypeMismatch { return_type: AbiType, value: InputValue },
+    #[error("No return value is expected but received {0:?}")]
+    UnexpectedReturnValue(InputValue),
 }

crates/noirc_abi/src/lib.rs

@@ -130,6 +130,8 @@ pub struct Abi {
     /// A map from the ABI's parameters to the indices they are written to in the [`WitnessMap`].
     /// This defines how to convert between the [`InputMap`] and [`WitnessMap`].
     pub param_witnesses: BTreeMap<String, Vec<Witness>>,
+    pub return_type: Option<AbiType>,
+    pub return_witnesses: Vec<Witness>,
 }
 
 impl Abi {
@@ -146,6 +148,11 @@ impl Abi {
         self.parameters.iter().map(|param| param.typ.field_count()).sum()
     }
 
+    /// Returns whether any values are needed to be made public for verification.
+    pub fn has_public_inputs(&self) -> bool {
+        self.return_type.is_some() || self.parameters.iter().any(|param| param.is_public())
+    }
+
     pub fn to_btree_map(&self) -> BTreeMap<String, AbiType> {
         let mut map = BTreeMap::new();
         for param in self.parameters.iter() {
@@ -164,18 +171,32 @@ impl Abi {
             .into_iter()
             .filter(|(param_name, _)| parameters.iter().any(|param| &param.name == param_name))
             .collect();
-        Abi { parameters, param_witnesses }
+        Abi {
+            parameters,
+            param_witnesses,
+            return_type: self.return_type,
+            return_witnesses: self.return_witnesses,
+        }
     }
 
     /// Encode a set of inputs as described in the ABI into a `WitnessMap`.
-    pub fn encode(&self, input_map: &InputMap, skip_output: bool) -> Result<WitnessMap, AbiError> {
-        self.check_for_unexpected_inputs(input_map)?;
+    pub fn encode(
+        &self,
+        input_map: &InputMap,
+        return_value: Option<InputValue>,
+    ) -> Result<WitnessMap, AbiError> {
+        // Check that no extra witness values have been provided.
+        let param_names = self.parameter_names();
+        if param_names.len() < input_map.len() {
+            let unexpected_params: Vec<String> =
+                input_map.keys().filter(|param| !param_names.contains(param)).cloned().collect();
+            return Err(AbiError::UnexpectedParams(unexpected_params));
+        }
 
         // First encode each input separately, performing any input validation.
         let encoded_input_map: BTreeMap<String, Vec<FieldElement>> = self
             .to_btree_map()
             .into_iter()
-            .filter(|(param_name, _)| !skip_output || param_name != MAIN_RETURN_NAME)
             .map(|(param_name, expected_type)| {
                 let value = input_map
                     .get(&param_name)
@@ -197,7 +218,7 @@ impl Abi {
             .collect::<Result<_, _>>()?;
 
         // Write input field elements into witness indices specified in `self.param_witnesses`.
-        let witness_map = encoded_input_map
+        let mut witness_map: WitnessMap = encoded_input_map
             .iter()
             .flat_map(|(param_name, encoded_param_fields)| {
                 let param_witness_indices = &self.param_witnesses[param_name];
@@ -208,19 +229,39 @@ impl Abi {
             })
             .collect();
 
-        Ok(witness_map)
-    }
-
-    /// Checks that no extra witness values have been provided.
-    fn check_for_unexpected_inputs(&self, inputs: &InputMap) -> Result<(), AbiError> {
-        let param_names = self.parameter_names();
-        if param_names.len() < inputs.len() {
-            let unexpected_params: Vec<String> =
-                inputs.keys().filter(|param| !param_names.contains(param)).cloned().collect();
-            return Err(AbiError::UnexpectedParams(unexpected_params));
+        // When encoding public inputs to be passed to the verifier, the user can must provide a return value
+        // to be inserted into the witness map. This is not needed when generating a witness when proving the circuit.
+        match (&self.return_type, return_value) {
+            (Some(return_type), Some(return_value)) => {
+                if !return_value.matches_abi(return_type) {
+                    return Err(AbiError::ReturnTypeMismatch {
+                        return_type: return_type.clone(),
+                        value: return_value,
+                    });
+                }
+                let encoded_return_fields = Self::encode_value(return_value)?;
+
+                // We need to be more careful when writing the return value's witness values.
+                // This is as it may share witness indices with other public inputs so we must check that when
+                // this occurs the witness values are consistent with each other.
+                self.return_witnesses.iter().zip(encoded_return_fields.iter()).try_for_each(
+                    |(&witness, &field_element)| match witness_map.insert(witness, field_element) {
+                        Some(existing_value) if existing_value != field_element => {
+                            Err(AbiError::InconsistentWitnessAssignment(witness))
+                        }
+                        _ => Ok(()),
+                    },
+                )?;
+            }
+            (None, Some(return_value)) => {
+                return Err(AbiError::UnexpectedReturnValue(return_value))
+            }
+            // We allow not passing a return value despite the circuit defining one
+            // in order to generate the initial partial witness.
+            (_, None) => {}
         }
 
-        Ok(())
+        Ok(witness_map)
     }
 
     fn encode_value(value: InputValue) -> Result<Vec<FieldElement>, AbiError> {
@@ -243,7 +284,10 @@ impl Abi {
     }
 
     /// Decode a `WitnessMap` into the types specified in the ABI.
-    pub fn decode(&self, witness_map: &WitnessMap) -> Result<InputMap, AbiError> {
+    pub fn decode(
+        &self,
+        witness_map: &WitnessMap,
+    ) -> Result<(InputMap, Option<InputValue>), AbiError> {
         let public_inputs_map =
             try_btree_map(self.parameters.clone(), |AbiParameter { name, typ, .. }| {
                 let param_witness_values =
@@ -261,7 +305,31 @@ impl Abi {
                     .map(|input_value| (name.clone(), input_value))
             })?;
 
-        Ok(public_inputs_map)
+        // We also attempt to decode the circuit's return value from `witness_map`.
+        let return_value = if let Some(return_type) = &self.return_type {
+            if let Ok(return_witness_values) =
+                try_vecmap(self.return_witnesses.clone(), |witness_index| {
+                    witness_map
+                        .get(&witness_index)
+                        .ok_or_else(|| AbiError::MissingParamWitnessValue {
+                            name: MAIN_RETURN_NAME.to_string(),
+                            witness_index,
+                        })
+                        .copied()
+                })
+            {
+                Some(Self::decode_value(&mut return_witness_values.into_iter(), return_type)?)
+            } else {
+                // Unlike for the circuit inputs, we tolerate not being able to find the witness values for the return value.
+                // This is because the user may be decoding a partial witness map for which is hasn't been calculated yet.
+                // If a return value is expected, this should be checked for by the user.
+                None
+            }
+        } else {
+            None
+        };
+
+        Ok((public_inputs_map, return_value))
     }
 
     fn decode_value(
@@ -323,10 +391,7 @@ mod test {
 
     use acvm::{acir::native_types::Witness, FieldElement};
 
-    use crate::{
-        input_parser::InputValue, Abi, AbiParameter, AbiType, AbiVisibility, InputMap,
-        MAIN_RETURN_NAME,
-    };
+    use crate::{input_parser::InputValue, Abi, AbiParameter, AbiType, AbiVisibility, InputMap};
 
     #[test]
     fn witness_encoding_roundtrip() {
@@ -342,18 +407,14 @@ mod test {
                     typ: AbiType::Field,
                     visibility: AbiVisibility::Public,
                 },
-                AbiParameter {
-                    name: MAIN_RETURN_NAME.to_string(),
-                    typ: AbiType::Field,
-                    visibility: AbiVisibility::Public,
-                },
             ],
             // Note that the return value shares a witness with `thing2`
             param_witnesses: BTreeMap::from([
                 ("thing1".to_string(), vec![Witness(1), Witness(2)]),
                 ("thing2".to_string(), vec![Witness(3)]),
-                (MAIN_RETURN_NAME.to_string(), vec![Witness(3)]),
             ]),
+            return_type: Some(AbiType::Field),
+            return_witnesses: vec![Witness(3)],
         };
 
         // Note we omit return value from inputs
@@ -362,14 +423,14 @@ mod test {
             ("thing2".to_string(), InputValue::Field(FieldElement::zero())),
         ]);
 
-        let witness_map = abi.encode(&inputs, true).unwrap();
-        let reconstructed_inputs = abi.decode(&witness_map).unwrap();
+        let witness_map = abi.encode(&inputs, None).unwrap();
+        let (reconstructed_inputs, return_value) = abi.decode(&witness_map).unwrap();
 
         for (key, expected_value) in inputs {
             assert_eq!(reconstructed_inputs[&key], expected_value);
         }
 
         // We also decode the return value (we can do this immediately as we know it shares a witness with an input).
-        assert_eq!(reconstructed_inputs[MAIN_RETURN_NAME], reconstructed_inputs["thing2"])
+        assert_eq!(return_value.unwrap(), reconstructed_inputs["thing2"])
     }
 }

crates/noirc_evaluator/src/lib.rs

@@ -10,7 +10,7 @@ use acvm::{
 };
 use errors::{RuntimeError, RuntimeErrorKind};
 use iter_extended::btree_map;
-use noirc_abi::{Abi, AbiType, AbiVisibility};
+use noirc_abi::{Abi, AbiType, AbiVisibility, MAIN_RETURN_NAME};
 use noirc_frontend::monomorphization::ast::*;
 use ssa::{node, ssa_gen::IrGenerator};
 use std::collections::{BTreeMap, BTreeSet};
@@ -53,7 +53,13 @@ pub fn create_circuit(
     let witness_index = evaluator.current_witness_index();
 
     let mut abi = program.abi;
-    abi.param_witnesses = evaluator.param_witnesses;
+
+    // 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();
+    abi.param_witnesses = param_witnesses;
+    abi.return_witnesses = return_witnesses;
 
     let public_inputs = evaluator.public_inputs.into_iter().collect();
     let optimized_circuit = acvm::compiler::compile(
@@ -291,15 +297,6 @@ impl Evaluator {
         let main_params = std::mem::take(&mut main.parameters);
         let abi_params = std::mem::take(&mut ir_gen.program.abi.parameters);
 
-        // Remove the return type from the parameters
-        // Since this is not in the main functions parameters.
-        //
-        // TODO(See Issue633) regarding adding a `return_type` field to the ABI struct
-        let abi_params: Vec<_> = abi_params
-            .into_iter()
-            .filter(|param| param.name != noirc_abi::MAIN_RETURN_NAME)
-            .collect();
-
         assert_eq!(main_params.len(), abi_params.len());
 
         for ((param_id, _, param_name, _), abi_param) in main_params.iter().zip(abi_params) {