fix: add support for nested arrays returned by oracles

acvm-repo/brillig_vm/src/lib.rs

@@ -462,124 +462,211 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> VM<'a, F, B> {
         }
     }
 
+    /// Writes flatten values to memory, using the provided type
+    /// Function calls itself recursively in order to work with recursive types (nested arrays)
+    /// values_idx is the current index in the values vector and is incremented every time
+    /// a value is written to memory
+    /// The function returns the address of the next value to be written
+    fn write_slice_of_values_to_memory(
+        &mut self,
+        destination: MemoryAddress,
+        values: &Vec<F>,
+        values_idx: &mut usize,
+        value_type: &HeapValueType,
+    ) -> Result<MemoryAddress, String> {
+        let mut current_pointer = destination;
+        match value_type {
+            HeapValueType::Simple(bit_size) => {
+                self.write_value_to_memory(destination, values, values_idx, *bit_size)?;
+                Ok(MemoryAddress(destination.to_usize() + 1))
+            }
+            HeapValueType::Array { value_types, size } => {
+                for _ in 0..*size {
+                    for typ in value_types {
+                        match typ {
+                            HeapValueType::Simple(len) => {
+                                self.write_value_to_memory(
+                                    current_pointer,
+                                    values,
+                                    values_idx,
+                                    *len,
+                                )?;
+                                current_pointer = MemoryAddress(current_pointer.to_usize() + 1);
+                            }
+                            HeapValueType::Array { .. } => {
+                                let destination = self.memory.read_ref(current_pointer);
+                                let destination = self.memory.read_ref(destination);
+                                self.write_slice_of_values_to_memory(
+                                    destination,
+                                    values,
+                                    values_idx,
+                                    typ,
+                                )?;
+                                current_pointer = MemoryAddress(current_pointer.to_usize() + 1);
+                            }
+                            HeapValueType::Vector { .. } => {
+                                return Err(format!(
+                                    "Unsupported returned type in foreign calls {:?}",
+                                    typ
+                                ));
+                            }
+                        }
+                    }
+                }
+                Ok(current_pointer)
+            }
+            HeapValueType::Vector { .. } => {
+                Err(format!("Unsupported returned type in foreign calls {:?}", value_type))
+            }
+        }
+    }
+
+    fn write_value_to_memory(
+        &mut self,
+        destination: MemoryAddress,
+        values: &Vec<F>,
+        values_idx: &mut usize,
+        value_bit_size: u32,
+    ) -> Result<(), String> {
+        if *values_idx >= values.len() {
+            return Err("Missing returned values, you should make sure the foreign call returns a list of flattened values".to_string());
+        }
+        let memory_value = MemoryValue::new_checked(values[*values_idx], value_bit_size);
+
+        if let Some(memory_value) = memory_value {
+            self.memory.write(destination, memory_value);
+        } else {
+            return Err(format!(
+                "Foreign call result value {} does not fit in bit size {}",
+                values[*values_idx], value_bit_size
+            ));
+        }
+        *values_idx += 1;
+        Ok(())
+    }
+
     fn write_foreign_call_result(
         &mut self,
         destinations: &[ValueOrArray],
         destination_value_types: &[HeapValueType],
         foreign_call_index: usize,
     ) -> Result<(), String> {
         let values = &self.foreign_call_results[foreign_call_index].values;
-
         if destinations.len() != values.len() {
             return Err(format!(
                 "{} output values were provided as a foreign call result for {} destination slots",
                 values.len(),
                 destinations.len()
             ));
         }
+        let flatten_values = if values.is_empty() { Vec::new() } else { values[0].fields() };
+        let mut flatten_values_idx = 0; //index of values read from flatten_values
+        let values = values.clone();
 
         for ((destination, value_type), output) in
-            destinations.iter().zip(destination_value_types).zip(values)
+            destinations.iter().zip(destination_value_types).zip(&values)
         {
             match (destination, value_type) {
-                (ValueOrArray::MemoryAddress(value_index), HeapValueType::Simple(bit_size)) => {
+            (ValueOrArray::MemoryAddress(value_index), HeapValueType::Simple(bit_size)) => {
+                match output {
+                    ForeignCallParam::Single(value) => {
+                        let memory_value = MemoryValue::new_checked(*value, *bit_size);
+                        if let Some(memory_value) = memory_value {
+                            self.memory.write(*value_index, memory_value);
+                        } else {
+                            return Err(format!(
+                                "Foreign call result value {} does not fit in bit size {}",
+                                value,
+                                bit_size
+                            ));
+                        }
+                    }
+                    _ => return Err(format!(
+                        "Function result size does not match brillig bytecode. Expected 1 result but got {output:?}")
+                    ),
+                }
+            }
+            (
+                ValueOrArray::HeapArray(HeapArray { pointer: pointer_index, size }),
+                HeapValueType::Array { value_types, size: type_size },
+            ) if size == type_size => {
+                if HeapValueType::all_simple(value_types) {
+                    let bit_sizes_iterator = value_types.iter().map(|typ| match typ {
+                        HeapValueType::Simple(bit_size) => *bit_size,
+                        _ => unreachable!("Expected simple value type"),
+                    }).cycle();
                     match output {
-                        ForeignCallParam::Single(value) => {
-                            let memory_value = MemoryValue::new_checked(*value, *bit_size);
-                            if let Some(memory_value) = memory_value {
-                                self.memory.write(*value_index, memory_value);
+                        ForeignCallParam::Array(values) => {
+                            if values.len() != *size {
+                                return Err("Foreign call result array doesn't match expected size".to_string());
+                            }
+                            // Convert the destination pointer to a usize
+                            let destination = self.memory.read_ref(*pointer_index);
+                            // Write to our destination memory
+                            let memory_values: Option<Vec<_>> = values.iter().zip(bit_sizes_iterator).map(
+                                |(value, bit_size)| MemoryValue::new_checked(*value, bit_size)).collect();
+                            if let Some(memory_values) = memory_values {
+                                self.memory.write_slice(destination, &memory_values);
                             } else {
                                 return Err(format!(
-                                    "Foreign call result value {} does not fit in bit size {}",
-                                    value,
-                                    bit_size
+                                    "Foreign call result values {:?} do not match expected bit sizes",
+                                    values,
                                 ));
                             }
                         }
-                        _ => return Err(format!(
-                            "Function result size does not match brillig bytecode. Expected 1 result but got {output:?}")
-                        ),
-                    }
-                }
-                (
-                    ValueOrArray::HeapArray(HeapArray { pointer: pointer_index, size }),
-                    HeapValueType::Array { value_types, size: type_size },
-                ) if size == type_size => {
-                    if HeapValueType::all_simple(value_types) {
-                        let bit_sizes_iterator = value_types.iter().map(|typ| match typ {
-                            HeapValueType::Simple(bit_size) => *bit_size,
-                            _ => unreachable!("Expected simple value type"),
-                        }).cycle();
-                        match output {
-                            ForeignCallParam::Array(values) => {
-                                if values.len() != *size {
-                                    return Err("Foreign call result array doesn't match expected size".to_string());
-                                }
-                                // Convert the destination pointer to a usize
-                                let destination = self.memory.read_ref(*pointer_index);
-                                // Write to our destination memory
-                                let memory_values: Option<Vec<_>> = values.iter().zip(bit_sizes_iterator).map(
-                                    |(value, bit_size)| MemoryValue::new_checked(*value, bit_size)).collect();
-                                if let Some(memory_values) = memory_values {
-                                    self.memory.write_slice(destination, &memory_values);
-                                } else {
-                                    return Err(format!(
-                                        "Foreign call result values {:?} do not match expected bit sizes",
-                                        values,
-                                    ));
-                                }
-                            }
-                            _ => {
-                                return Err("Function result size does not match brillig bytecode size".to_string());
-                            }
+                        _ => {
+                            return Err("Function result size does not match brillig bytecode size".to_string());
                         }
-                    } else {
-                        unimplemented!("deflattening heap arrays from foreign calls");
                     }
-                }
-                (
-                    ValueOrArray::HeapVector(HeapVector {pointer: pointer_index, size: size_index }),
-                    HeapValueType::Vector { value_types },
-                ) => {
-                    if HeapValueType::all_simple(value_types) {
-                        let bit_sizes_iterator = value_types.iter().map(|typ| match typ {
-                            HeapValueType::Simple(bit_size) => *bit_size,
-                            _ => unreachable!("Expected simple value type"),
-                        }).cycle();
-                        match output {
-                            ForeignCallParam::Array(values) => {
-                                // Set our size in the size address
-                                self.memory.write(*size_index, values.len().into());
-                                // Convert the destination pointer to a usize
-                                let destination = self.memory.read_ref(*pointer_index);
-                                // Write to our destination memory
-                                let memory_values: Option<Vec<_>> = values.iter().zip(bit_sizes_iterator).map(|(value, bit_size)| MemoryValue::new_checked(*value, bit_size)).collect();
-                                if let Some(memory_values) = memory_values {
-                                    self.memory.write_slice(destination, &memory_values);
-                                }else{
-                                    return Err(format!(
-                                        "Foreign call result values {:?} do not match expected bit sizes",
-                                        values,
-                                    ));
-                                }
-                            }
-                            _ => {
-                                return Err("Function result size does not match brillig bytecode size".to_string());
+                } else {
+                    // foreign call returning flatten values into a nested type, so the sizes do not match
+                    let destination = self.memory.read_ref(*pointer_index);
+                    let return_type = value_type;
+                    self.write_slice_of_values_to_memory(destination, &flatten_values, &mut flatten_values_idx, return_type)?;
+            }
+        }
+            (
+                ValueOrArray::HeapVector(HeapVector {pointer: pointer_index, size: size_index }),
+                HeapValueType::Vector { value_types },
+            ) => {
+                if HeapValueType::all_simple(value_types) {
+                    let bit_sizes_iterator = value_types.iter().map(|typ| match typ {
+                        HeapValueType::Simple(bit_size) => *bit_size,
+                        _ => unreachable!("Expected simple value type"),
+                    }).cycle();
+                    match output {
+                        ForeignCallParam::Array(values) => {
+                            // Set our size in the size address
+                            self.memory.write(*size_index, values.len().into());
+                            // Convert the destination pointer to a usize
+                            let destination = self.memory.read_ref(*pointer_index);
+                            // Write to our destination memory
+                            let memory_values: Option<Vec<_>> = values.iter().zip(bit_sizes_iterator).map(|(value, bit_size)| MemoryValue::new_checked(*value, bit_size)).collect();
+                            if let Some(memory_values) = memory_values {
+                                self.memory.write_slice(destination, &memory_values);
+                            }else{
+                                return Err(format!(
+                                    "Foreign call result values {:?} do not match expected bit sizes",
+                                    values,
+                                ));
                             }
                         }
-                    } else {
-                        unimplemented!("deflattening heap vectors from foreign calls");
+                        _ => {
+                            return Err("Function result size does not match brillig bytecode size".to_string());
+                        }
                     }
-                }
-                _ => {
-                    return Err(format!("Unexpected value type {value_type:?} for destination {destination:?}"));
+                } else {
+                    unimplemented!("deflattening heap vectors from foreign calls");
                 }
             }
+            _ => {
+                return Err(format!("Unexpected value type {value_type:?} for destination {destination:?}"));
+            }
+        }
         }
 
         Ok(())
     }
-
     /// Process a binary operation.
     /// This method will not modify the program counter.
     fn process_binary_field_op(

compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs

@@ -1737,8 +1737,7 @@ impl<'block> BrilligBlock<'block> {
                     dfg,
                 );
                 let array = variable.extract_array();
-                self.brillig_context.codegen_allocate_fixed_length_array(array.pointer, array.size);
-                self.brillig_context.usize_const_instruction(array.rc, 1_usize.into());
+                self.allocate_nested_array(typ, Some(array));
 
                 variable
             }
@@ -1765,6 +1764,39 @@ impl<'block> BrilligBlock<'block> {
         }
     }
 
+    fn allocate_nested_array(
+        &mut self,
+        typ: &Type,
+        array: Option<BrilligArray>,
+    ) -> BrilligVariable {
+        match typ {
+            Type::Array(types, size) => {
+                let array = array.unwrap_or(BrilligArray {
+                    pointer: self.brillig_context.allocate_register(),
+                    size: *size,
+                    rc: self.brillig_context.allocate_register(),
+                });
+                self.brillig_context.codegen_allocate_fixed_length_array(array.pointer, array.size);
+                self.brillig_context.usize_const_instruction(array.rc, 1_usize.into());
+
+                let mut index = 0_usize;
+                for _ in 0..*size {
+                    for element_type in types.iter() {
+                        if matches!(element_type, Type::Array(_, _)) {
+                            let inner_array = self.allocate_nested_array(element_type, None);
+                            let idx =
+                                self.brillig_context.make_usize_constant_instruction(index.into());
+                            self.store_variable_in_array(array.pointer, idx, inner_array);
+                        }
+                        index += 1;
+                    }
+                }
+                BrilligVariable::BrilligArray(array)
+            }
+            _ => unreachable!("ICE: allocate_nested_array() expects an array, got {typ:?}"),
+        }
+    }
+
     /// Gets the "user-facing" length of an array.
     /// An array of structs with two fields would be stored as an 2 * array.len() array/vector.
     /// So we divide the length by the number of subitems in an item to get the user-facing length.

docs/docs/how_to/how-to-oracles.md

@@ -173,6 +173,10 @@ interface ForeignCallResult {
 }
 ```
 
+::: Multidimensional Arrays
+
+If the Oracle function is returning an array containing other arrays, such as `[['1','2],['3','4']]`, you need to provide the values in json as flattened values. In the previous example, it would be `['1', '2', '3', '4']`. In the noir program, the Oracle signature can use a nested type, the flattened values will be automatically converted to the nested type.
+
 :::
 
 ## Step 3 - Usage with Nargo

test_programs/noir_test_success/regression_4561/Nargo.toml

@@ -0,0 +1,6 @@
+[package]
+name = "regression_4561"
+type = "bin"
+authors = [""]
+
+[dependencies]

test_programs/noir_test_success/regression_4561/src/main.nr

@@ -0,0 +1,18 @@
+// Regression test for issue #4561
+use dep::std::test::OracleMock;
+
+type TReturnElem = [Field; 3];
+type TReturn = [TReturnElem; 2];
+
+#[oracle(repro)]
+unconstrained fn repro_oracle() -> TReturn {}
+
+unconstrained fn repro_unconstrained() -> TReturn {
+    repro_oracle()
+}
+
+#[test]
+fn test_mock() {
+    OracleMock::mock("repro").returns([1, 2, 3, 4, 5, 6]);
+    assert_eq(repro_unconstrained(), [[1, 2, 3], [4, 5, 6]]);
+}