chore: use witnesses from the generated acir in the ABI (#2095)

* Use witnesses from the generated acir in the ABI

* Code review

crates/noirc_evaluator/src/ssa_refactor.rs

@@ -77,10 +77,16 @@ pub fn create_circuit(
     show_output: bool,
 ) -> Result<(Circuit, DebugInfo, Abi), RuntimeError> {
     let func_sig = program.main_function_signature.clone();
-    let GeneratedAcir { current_witness_index, opcodes, return_witnesses, locations, .. } =
-        optimize_into_acir(program, show_output, enable_ssa_logging, enable_brillig_logging)?;
+    let GeneratedAcir {
+        current_witness_index,
+        opcodes,
+        return_witnesses,
+        locations,
+        input_witnesses,
+        ..
+    } = optimize_into_acir(program, show_output, enable_ssa_logging, enable_brillig_logging)?;
 
-    let abi = gen_abi(func_sig, return_witnesses.clone());
+    let abi = gen_abi(func_sig, &input_witnesses, return_witnesses.clone());
     let public_abi = abi.clone().public_abi();
 
     let public_parameters =

crates/noirc_evaluator/src/ssa_refactor/abi_gen/mod.rs

@@ -1,29 +1,36 @@
 use std::collections::BTreeMap;
 
 use acvm::acir::native_types::Witness;
-use iter_extended::{btree_map, vecmap};
+use iter_extended::btree_map;
 use noirc_abi::{Abi, AbiParameter, FunctionSignature};
 
 /// Arranges a function signature and a generated circuit's return witnesses into a
 /// `noirc_abi::Abi`.
-pub(crate) fn gen_abi(func_sig: FunctionSignature, return_witnesses: Vec<Witness>) -> Abi {
+pub(crate) fn gen_abi(
+    func_sig: FunctionSignature,
+    input_witnesses: &[Witness],
+    return_witnesses: Vec<Witness>,
+) -> Abi {
     let (parameters, return_type) = func_sig;
-    let param_witnesses = param_witnesses_from_abi_param(&parameters);
+    let param_witnesses = param_witnesses_from_abi_param(&parameters, input_witnesses);
     Abi { parameters, return_type, param_witnesses, return_witnesses }
 }
 
 // Takes each abi parameter and shallowly maps to the expected witness range in which the
 // parameter's constituent values live.
 fn param_witnesses_from_abi_param(
     abi_params: &Vec<AbiParameter>,
+    input_witnesses: &[Witness],
 ) -> BTreeMap<String, Vec<Witness>> {
-    let mut offset = 1;
+    let mut idx = 0_usize;
+
     btree_map(abi_params, |param| {
         let num_field_elements_needed = param.typ.field_count();
-        let idx_start = offset;
-        let idx_end = idx_start + num_field_elements_needed;
-        let witnesses = vecmap(idx_start..idx_end, Witness);
-        offset += num_field_elements_needed;
-        (param.name.clone(), witnesses)
+        let mut wit = Vec::new();
+        for _ in 0..num_field_elements_needed {
+            wit.push(input_witnesses[idx]);
+            idx += 1;
+        }
+        (param.name.clone(), wit)
     })
 }

crates/noirc_evaluator/src/ssa_refactor/acir_gen/acir_ir/acir_variable.rs

@@ -110,6 +110,26 @@ pub(crate) struct AcirContext {
 }
 
 impl AcirContext {
+    pub(crate) fn current_witness_index(&self) -> Witness {
+        self.acir_ir.current_witness_index()
+    }
+
+    pub(crate) fn extract_witness(&self, inputs: &[AcirValue]) -> Vec<u32> {
+        inputs
+            .iter()
+            .flat_map(|value| value.clone().flatten())
+            .map(|value| {
+                self.vars
+                    .get(&value.0)
+                    .expect("ICE: undeclared AcirVar")
+                    .to_expression()
+                    .to_witness()
+                    .expect("ICE - cannot extract a witness")
+                    .0
+            })
+            .collect()
+    }
+
     /// Adds a constant to the context and assigns a Variable to represent it
     pub(crate) fn add_constant(&mut self, constant: FieldElement) -> AcirVar {
         let constant_data = AcirVarData::Const(constant);
@@ -808,7 +828,8 @@ impl AcirContext {
     }
 
     /// Terminates the context and takes the resulting `GeneratedAcir`
-    pub(crate) fn finish(self) -> GeneratedAcir {
+    pub(crate) fn finish(mut self, inputs: Vec<u32>) -> GeneratedAcir {
+        self.acir_ir.input_witnesses = vecmap(inputs, Witness);
         self.acir_ir
     }
 

crates/noirc_evaluator/src/ssa_refactor/acir_gen/acir_ir/generated_acir.rs

@@ -40,6 +40,9 @@ pub(crate) struct GeneratedAcir {
     /// abi's return type.
     pub(crate) return_witnesses: Vec<Witness>,
 
+    /// All witness indices which are inputs to the main function
+    pub(crate) input_witnesses: Vec<Witness>,
+
     /// Correspondance between an opcode index (in opcodes) and the source code location which generated it
     pub(crate) locations: HashMap<usize, Location>,
 

crates/noirc_evaluator/src/ssa_refactor/acir_gen/mod.rs

@@ -2,6 +2,7 @@
 
 use std::collections::{HashMap, HashSet};
 use std::fmt::Debug;
+use std::ops::RangeInclusive;
 
 use crate::brillig::brillig_ir::BrilligContext;
 use crate::{
@@ -172,16 +173,15 @@ impl Context {
     ) -> Result<GeneratedAcir, AcirGenError> {
         let dfg = &main_func.dfg;
         let entry_block = &dfg[main_func.entry_block()];
-
-        self.convert_ssa_block_params(entry_block.parameters(), dfg)?;
+        let input_witness = self.convert_ssa_block_params(entry_block.parameters(), dfg)?;
 
         for instruction_id in entry_block.instructions() {
             self.convert_ssa_instruction(*instruction_id, dfg, ssa, &brillig, allow_log_ops)?;
         }
 
         self.convert_ssa_return(entry_block.unwrap_terminator(), dfg);
 
-        Ok(self.acir_context.finish())
+        Ok(self.acir_context.finish(input_witness.collect()))
     }
 
     fn convert_brillig_main(
@@ -195,6 +195,7 @@ impl Context {
             let typ = dfg.type_of_value(*param_id);
             self.create_value_from_type(&typ, &mut |this, _| Ok(this.acir_context.add_variable()))
         })?;
+        let witness_inputs = self.acir_context.extract_witness(&inputs);
 
         let outputs: Vec<AcirType> =
             vecmap(main_func.returns(), |result_id| dfg.type_of_value(*result_id).into());
@@ -213,15 +214,17 @@ impl Context {
             self.acir_context.return_var(acir_var);
         }
 
-        Ok(self.acir_context.finish())
+        Ok(self.acir_context.finish(witness_inputs))
     }
 
     /// Adds and binds `AcirVar`s for each numeric block parameter or block parameter array element.
     fn convert_ssa_block_params(
         &mut self,
         params: &[ValueId],
         dfg: &DataFlowGraph,
-    ) -> Result<(), AcirGenError> {
+    ) -> Result<RangeInclusive<u32>, AcirGenError> {
+        // The first witness (if any) is the next one
+        let start_witness = self.acir_context.current_witness_index().0 + 1;
         for param_id in params {
             let typ = dfg.type_of_value(*param_id);
             let value = self.convert_ssa_block_param(&typ)?;
@@ -238,7 +241,8 @@ impl Context {
             }
             self.ssa_values.insert(*param_id, value);
         }
-        Ok(())
+        let end_witness = self.acir_context.current_witness_index().0;
+        Ok(start_witness..=end_witness)
     }
 
     fn convert_ssa_block_param(&mut self, param_type: &Type) -> Result<AcirValue, AcirGenError> {