feat: initial implementation of slices in brillig (#1932)

* feat: get push/pop working

* feat: get all slice operations working

* chore: remove printing of initial ssa

* chore: remove print

* refactor: use memory op

* feat: implement array_set for slices

* Update crates/nargo_cli/tests/test_data_ssa_refactor/brillig_slices/src/main.nr

* test: added unit testing for slice ops

Author: sirasistant

crates/nargo_cli/tests/test_data_ssa_refactor/brillig_slices/Nargo.toml

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

crates/nargo_cli/tests/test_data_ssa_refactor/brillig_slices/Prover.toml

@@ -0,0 +1,2 @@
+x = "5"
+y = "10"

crates/nargo_cli/tests/test_data_ssa_refactor/brillig_slices/src/main.nr

@@ -0,0 +1,72 @@
+use dep::std::slice;
+use dep::std;
+
+unconstrained fn main(x: Field, y: Field) {
+    // Mark it as mut so the compiler doesn't simplify the following operations
+    // But don't reuse the mut slice variable until this is fixed https://github.com/noir-lang/noir/issues/1931
+    let slice: [Field] = [y, x];
+    assert(slice.len() == 2);
+
+    let mut pushed_back_slice = slice.push_back(7);
+    assert(pushed_back_slice.len() == 3);
+    assert(pushed_back_slice[0] == y);
+    assert(pushed_back_slice[1] == x);
+    assert(pushed_back_slice[2] == 7);
+
+    // Array set on slice target
+    pushed_back_slice[0] = x;
+    pushed_back_slice[1] = y;
+    pushed_back_slice[2] = 1;
+
+    assert(pushed_back_slice[0] == x);
+    assert(pushed_back_slice[1] == y);
+    assert(pushed_back_slice[2] == 1);
+
+    assert(slice.len() == 2);
+
+    let pushed_front_slice = pushed_back_slice.push_front(2);
+    assert(pushed_front_slice.len() == 4);
+    assert(pushed_front_slice[0] == 2);
+    assert(pushed_front_slice[1] == x);
+    assert(pushed_front_slice[2] == y);
+    assert(pushed_front_slice[3] == 1);
+
+    let (item, popped_front_slice) = pushed_front_slice.pop_front();
+    assert(item == 2);
+
+    assert(popped_front_slice.len() == 3);
+    assert(popped_front_slice[0] == x);
+    assert(popped_front_slice[1] == y);
+    assert(popped_front_slice[2] == 1);
+
+    let (popped_back_slice, another_item) = popped_front_slice.pop_back();
+    assert(another_item == 1);
+
+    assert(popped_back_slice.len() == 2);
+    assert(popped_back_slice[0] == x);
+    assert(popped_back_slice[1] == y);
+
+    let inserted_slice = popped_back_slice.insert(1, 2);
+    assert(inserted_slice.len() == 3);
+    assert(inserted_slice[0] == x);
+    assert(inserted_slice[1] == 2);
+    assert(inserted_slice[2] == y);
+
+    let (removed_slice, should_be_2) = inserted_slice.remove(1);
+    assert(should_be_2 == 2);
+
+    assert(removed_slice.len() == 2);
+    assert(removed_slice[0] == x);
+    assert(removed_slice[1] == y);
+
+    let (slice_with_only_x, should_be_y) = removed_slice.remove(1);
+    assert(should_be_y == y);
+
+    assert(slice_with_only_x.len() == 1);
+    assert(removed_slice[0] == x);
+
+    let (empty_slice, should_be_x) = slice_with_only_x.remove(0);
+    assert(should_be_x == x);
+    assert(empty_slice.len() == 0);
+}
+

crates/noirc_evaluator/src/brillig/brillig_gen.rs

@@ -2,6 +2,7 @@ pub(crate) mod brillig_black_box;
 pub(crate) mod brillig_block;
 pub(crate) mod brillig_directive;
 pub(crate) mod brillig_fn;
+pub(crate) mod brillig_slice_ops;
 
 use crate::ssa_refactor::ir::{function::Function, post_order::PostOrder};
 
@@ -18,7 +19,7 @@ pub(crate) fn convert_ssa_function(func: &Function) -> BrilligArtifact {
     reverse_post_order.reverse();
 
     let mut function_context =
-        FunctionContext { function_id: func.id(), ssa_value_to_register: HashMap::new() };
+        FunctionContext { function_id: func.id(), ssa_value_to_brillig_variable: HashMap::new() };
 
     let mut brillig_context = BrilligContext::new(
         FunctionContext::parameters(func),

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

@@ -1,51 +1,137 @@
 use std::collections::HashMap;
 
-use acvm::acir::brillig::RegisterIndex;
+use acvm::brillig_vm::brillig::{HeapArray, HeapVector, RegisterIndex, RegisterOrMemory};
 
 use crate::{
     brillig::brillig_ir::{
         artifact::{BrilligParameter, Label},
         BrilligContext,
     },
     ssa_refactor::ir::{
+        dfg::DataFlowGraph,
         function::{Function, FunctionId},
-        types::Type,
+        types::{CompositeType, Type},
         value::ValueId,
     },
 };
 
-use super::brillig_block::compute_size_of_type;
-
 pub(crate) struct FunctionContext {
     pub(crate) function_id: FunctionId,
-    /// Map from SSA values to Register Indices.
-    pub(crate) ssa_value_to_register: HashMap<ValueId, RegisterIndex>,
+    /// Map from SSA values to register or memory.
+    pub(crate) ssa_value_to_brillig_variable: HashMap<ValueId, RegisterOrMemory>,
 }
 
 impl FunctionContext {
-    /// Gets a `RegisterIndex` for a `ValueId`, if one already exists
-    /// or creates a new `RegisterIndex` using the latest available
-    /// free register.
-    pub(crate) fn get_or_create_register(
+    /// For a given SSA value id, create and cache the a corresponding variable.
+    /// This will allocate the needed registers for the variable.
+    pub(crate) fn create_variable(
+        &mut self,
+        brillig_context: &mut BrilligContext,
+        value: ValueId,
+        dfg: &DataFlowGraph,
+    ) -> RegisterOrMemory {
+        let typ = dfg.type_of_value(value);
+
+        let variable = match typ {
+            Type::Numeric(_) | Type::Reference => {
+                let register = brillig_context.allocate_register();
+                RegisterOrMemory::RegisterIndex(register)
+            }
+            Type::Array(_, _) => {
+                let pointer_register = brillig_context.allocate_register();
+                let size = compute_size_of_type(&typ);
+                RegisterOrMemory::HeapArray(HeapArray { pointer: pointer_register, size })
+            }
+            Type::Slice(_) => {
+                let pointer_register = brillig_context.allocate_register();
+                let size_register = brillig_context.allocate_register();
+                RegisterOrMemory::HeapVector(HeapVector {
+                    pointer: pointer_register,
+                    size: size_register,
+                })
+            }
+            Type::Function => {
+                unreachable!("ICE: Function values should have been removed from the SSA")
+            }
+        };
+
+        // Cache the `ValueId` so that if we call get_variable, it will
+        // return the registers that have just been created.
+        //
+        // WARNING: This assumes that a registers won't be reused for a different value.
+        // If you overwrite the registers, then the cache will be invalid.
+
+        if self.ssa_value_to_brillig_variable.insert(value, variable).is_some() {
+            unreachable!("ICE: ValueId {value:?} was already in cache");
+        }
+
+        variable
+    }
+
+    /// For a given SSA value id, return the corresponding cached variable.
+    pub(crate) fn get_variable(&mut self, value: ValueId) -> RegisterOrMemory {
+        *self
+            .ssa_value_to_brillig_variable
+            .get(&value)
+            .unwrap_or_else(|| panic!("ICE: Value not found in cache {value}"))
+    }
+
+    pub(crate) fn get_or_create_variable(
         &mut self,
         brillig_context: &mut BrilligContext,
         value: ValueId,
+        dfg: &DataFlowGraph,
+    ) -> RegisterOrMemory {
+        if let Some(variable) = self.ssa_value_to_brillig_variable.get(&value) {
+            return *variable;
+        }
+
+        self.create_variable(brillig_context, value, dfg)
+    }
+
+    /// Creates a variable that fits in a single register and returns the register.
+    pub(crate) fn create_register_variable(
+        &mut self,
+        brillig_context: &mut BrilligContext,
+        value: ValueId,
+        dfg: &DataFlowGraph,
     ) -> RegisterIndex {
-        if let Some(register_index) = self.ssa_value_to_register.get(&value) {
-            return *register_index;
+        let variable = self.create_variable(brillig_context, value, dfg);
+        self.extract_register(variable)
+    }
+
+    pub(crate) fn extract_register(&self, variable: RegisterOrMemory) -> RegisterIndex {
+        match variable {
+            RegisterOrMemory::RegisterIndex(register_index) => register_index,
+            _ => unreachable!("ICE: Expected register, got {variable:?}"),
         }
+    }
 
-        let register = brillig_context.allocate_register();
+    pub(crate) fn extract_heap_array(&self, variable: RegisterOrMemory) -> HeapArray {
+        match variable {
+            RegisterOrMemory::HeapArray(array) => array,
+            _ => unreachable!("ICE: Expected array, got {variable:?}"),
+        }
+    }
 
-        // Cache the `ValueId` so that if we call it again, it will
-        // return the register that has just been created.
-        //
-        // WARNING: This assumes that a register has not been
-        // modified. If a MOV instruction has overwritten the value
-        // at a register, then this cache will be invalid.
-        self.ssa_value_to_register.insert(value, register);
+    pub(crate) fn extract_heap_vector(&self, variable: RegisterOrMemory) -> HeapVector {
+        match variable {
+            RegisterOrMemory::HeapVector(vector) => vector,
+            _ => unreachable!("ICE: Expected vector, got {variable:?}"),
+        }
+    }
 
-        register
+    /// Collects the registers that a given variable is stored in.
+    pub(crate) fn extract_registers(&self, variable: RegisterOrMemory) -> Vec<RegisterIndex> {
+        match variable {
+            RegisterOrMemory::RegisterIndex(register_index) => vec![register_index],
+            RegisterOrMemory::HeapArray(array) => {
+                vec![array.pointer]
+            }
+            RegisterOrMemory::HeapVector(vector) => {
+                vec![vector.pointer, vector.size]
+            }
+        }
     }
 
     /// Creates a function label from a given SSA function id.
@@ -62,6 +148,7 @@ impl FunctionContext {
                 match typ {
                     Type::Numeric(_) | Type::Reference => BrilligParameter::Register,
                     Type::Array(..) => BrilligParameter::HeapArray(compute_size_of_type(&typ)),
+                    Type::Slice(_) => BrilligParameter::HeapVector,
                     _ => unimplemented!("Unsupported function parameter type {typ:?}"),
                 }
             })
@@ -77,9 +164,25 @@ impl FunctionContext {
                 match typ {
                     Type::Numeric(_) | Type::Reference => BrilligParameter::Register,
                     Type::Array(..) => BrilligParameter::HeapArray(compute_size_of_type(&typ)),
+                    Type::Slice(_) => BrilligParameter::HeapVector,
                     _ => unimplemented!("Unsupported return value type {typ:?}"),
                 }
             })
             .collect()
     }
 }
+
+/// Computes the size of an SSA composite type
+pub(crate) fn compute_size_of_composite_type(typ: &CompositeType) -> usize {
+    typ.iter().map(compute_size_of_type).sum()
+}
+
+/// Finds out the size of a given SSA type
+/// This is needed to store values in memory
+pub(crate) fn compute_size_of_type(typ: &Type) -> usize {
+    match typ {
+        Type::Numeric(_) => 1,
+        Type::Array(types, item_count) => compute_size_of_composite_type(types) * item_count,
+        _ => todo!("ICE: Type not supported {typ:?}"),
+    }
+}