Fix member accesses for nested structs returned from funcs

crates/acvm/src/lib.rs

@@ -43,14 +43,17 @@ pub trait PartialWitnessGenerator {
         initial_witness: &mut BTreeMap<Witness, FieldElement>,
         gates: Vec<Gate>,
     ) -> GateResolution {
+        // dbg!(gates.clone());
         if gates.is_empty() {
+            dbg!("gates are empty");
             return GateResolution::Resolved;
         }
         let mut unsolved_gates: Vec<Gate> = Vec::new();
 
         for gate in gates.into_iter() {
             let unsolved = match &gate {
                 Gate::Arithmetic(arith) => {
+                    // dbg!("arithmetic gate");
                     let result = ArithmeticSolver::solve(initial_witness, arith);
                     match result {
                         GateResolution::Resolved => false,
@@ -197,9 +200,11 @@ pub trait PartialWitnessGenerator {
                 },
             };
             if unsolved {
+                // dbg!(gate.clone());
                 unsolved_gates.push(gate);
             }
         }
+        // dbg!(unsolved_gates.clone());
         self.solve(initial_witness, unsolved_gates)
     }
 

crates/nargo/src/cli/prove_cmd.rs

@@ -176,12 +176,13 @@ pub fn solve_witness<P: AsRef<Path>>(
     }
     // Map initial witnesses with their values
     let abi = compiled_program.abi.as_ref().unwrap();
+
     // Solve the remaining witnesses
     let (mut solved_witness, rv) = process_abi_with_input(abi.clone(), &witness_map)?;
 
     let backend = crate::backends::ConcreteBackend;
     let solver_res = backend.solve(&mut solved_witness, compiled_program.circuit.gates.clone());
-
+    dbg!("got a solver res");
     match solver_res {
         GateResolution::UnsupportedOpcode(opcode) => return Err(CliError::Generic(format!(
                 "backend does not currently support the {} opcode. ACVM does not currently fall back to arithmetic gates.",

crates/nargo/tests/test_data/nested_structs/Nargo.toml

@@ -0,0 +1,5 @@
+[package]
+authors = [""]
+compiler_version = "0.1"
+
+[dependencies]

crates/nargo/tests/test_data/nested_structs/Prover.toml

@@ -0,0 +1,4 @@
+y = "1"
+a = "2"
+b = "3"
+return = ""

crates/nargo/tests/test_data/nested_structs/src/main.nr

@@ -0,0 +1,78 @@
+struct Parent {
+    foo: Field,
+    child: Child,
+    bar: Field,
+}
+
+struct Child {
+    child2: Child2,
+    foo: u32,
+    bar: u32,
+}
+
+struct Child2 {
+    input_a: u30,
+    input_b: u30,
+    flag: bool,
+}
+
+fn main(y : pub Field, a: u32, b: u32) -> pub Field {
+    // TODO: Currently broken for fixed return value problem discussed in the comments below 
+    // let (struct_from_tuple, a_bool) = test_struct_in_tuple(true);
+    // constrain struct_from_tuple.my_bool == true;
+    // constrain a_bool == 10;
+
+    // Always passes as the nested struct is not returned from a function
+    let c2 = Child2 { input_a: 0 as u30, input_b: 1 as u30, flag: true };
+    let c = Child { child2: c2, foo: a, bar: b };
+    let p = Parent { foo: y, child: c, bar: y };
+    constrain y == p.foo;
+
+    // Fails unless we specify the return value. Otherwise we get a stack overflow
+    let p2 = new_parent(c2, a, b);
+    // Passes without a specified return value when I add a constrain stmt on the struct field
+    // Trying to return p2.foo alone though will always fail without a specified return value in Prover.toml
+    let q2 = p2.foo + y;
+    constrain q2 == 3; 
+    // Then a valid return
+    p2.foo
+
+    // Passes only when the field is constrained 
+    // constrain p2.child.child2.input_b == 1 as u30;
+    // p2.child.child2.input_b as Field 
+
+    // Fails without change `let q2 = p2.bar + y`
+    // p2.bar
+
+    // Always passes as parent_wrapper returns a Field
+    // let res = parent_wrapper(a, b);
+    // constrain res == a as Field;
+    // let q2 = res + y;
+    // q2
+}
+
+fn parent_wrapper(a: u32, b: u32) -> Field {
+    let c2 = Child2 { input_a: 0 as u30, input_b: 1 as u30, flag: true };
+
+    let p2 = new_parent(c2, a, b);
+
+    p2.foo
+}
+
+fn new_parent(c2: Child2, a: u32, b: u32) -> Parent {
+    let c = Child { child2: c2, foo: a, bar: b };
+
+    let p = Parent { foo: a as Field, child: c, bar: b as Field };
+
+    p
+}
+
+struct MyStruct {
+    my_bool: bool,
+}
+fn test_struct_in_tuple(a_bool : bool) -> (MyStruct, Field) {
+    let my_struct = MyStruct {
+        my_bool: a_bool,
+    };
+    (my_struct, 10)
+}

crates/noirc_evaluator/src/ssa/code_gen.rs

@@ -30,7 +30,7 @@ pub struct IRGenerator {
     pub program: Program,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Value {
     Single(NodeId),
     Tuple(Vec<Value>),
@@ -82,6 +82,11 @@ impl Value {
     pub fn into_field_member(self, field_index: usize) -> Value {
         match self {
             Value::Single(_) => {
+                // When we return a struct from a function it is flattened.
+                // If a developer returns a one element struct they will panic without this conditional
+                if field_index == 0 {
+                    return self;
+                }
                 unreachable!("Runtime type error, expected struct but found a single value")
             }
             Value::Tuple(mut fields) => fields.remove(field_index),
@@ -194,6 +199,29 @@ impl IRGenerator {
 
     fn codegen_identifier(&mut self, ident: &Ident) -> Value {
         let value = self.variable_values[&ident.id].clone();
+
+        // Determine whether the ident being accessed is a nested struct/tuple that was returned from a function.
+        // When creating an SSA function we flatten the return values. This can lead to a mismatch in the NodeId that is indexed by a member access.
+        // We recreate a struct/tuple with an unflattened structure in order to enable accurate indexing.
+        // If the Value::Tuple has a longer list of NodeIds (meaning it has been flattened) we know to use our recreated value.
+        let value = match value.clone() {
+            Value::Single(_) => value,
+            Value::Tuple(orig_values) => {
+                match ident.typ.clone() {
+                    Type::Tuple(fields) => {
+                        if orig_values.len() > fields.len() {
+                            // TODO: Look into how to avoid creating a new value here everytime we have a tuple member access from a nested tuple that was returned by a func
+                            // We might have to cache the updated SSA values for nested structs that are the result of func returns
+                            self.create_new_value(&ident.typ, &ident.name, Some(ident.id))
+                        } else {
+                            value
+                        }
+                    }
+                    _ => value,
+                }
+            }
+        };
+
         self.get_current_value(&value)
     }
 

crates/noirc_evaluator/src/ssa/function.rs

@@ -175,6 +175,7 @@ impl IRGenerator {
                 func.result_types.push(ObjectType::NotAnObject);
             }
         }
+
         self.context.new_instruction(
             node::Operation::Return(returned_values),
             node::ObjectType::NotAnObject,