From 89068f5a7c37901c7d9f5098604802b442bf517b Mon Sep 17 00:00:00 2001
From: sirasistant <sirasistant@gmail.com>
Date: Tue, 28 Nov 2023 09:32:27 +0000
Subject: [PATCH] fix: make toBits accept arrays and add test

---
 .../src/brillig/brillig_gen/brillig_block.rs  | 21 ++++----
 .../execution_success/brillig_cow/Nargo.toml  |  6 +++
 .../execution_success/brillig_cow/Prover.toml |  7 +++
 .../execution_success/brillig_cow/src/main.nr | 54 +++++++++++++++++++
 4 files changed, 79 insertions(+), 9 deletions(-)
 create mode 100644 tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml
 create mode 100644 tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml
 create mode 100644 tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
index d5e958f032a..0e06a36fd94 100644
--- a/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
+++ b/compiler/noirc_evaluator/src/brillig/brillig_gen/brillig_block.rs
@@ -489,15 +489,18 @@ impl<'block> BrilligBlock<'block> {
                     );
                     let target_len = target_len_variable.extract_register();
 
-                    let target_vector = self
-                        .variables
-                        .define_variable(
-                            self.function_context,
-                            self.brillig_context,
-                            results[1],
-                            dfg,
-                        )
-                        .extract_vector();
+                    let target_vector = match self.variables.define_variable(
+                        self.function_context,
+                        self.brillig_context,
+                        results[1],
+                        dfg,
+                    ) {
+                        BrilligVariable::BrilligArray(array) => {
+                            self.brillig_context.array_to_vector(&array)
+                        }
+                        BrilligVariable::BrilligVector(vector) => vector,
+                        BrilligVariable::Simple(..) => unreachable!("ICE: ToBits on non-array"),
+                    };
 
                     let radix = self.brillig_context.make_constant(2_usize.into());
 
diff --git a/tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml b/tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml
new file mode 100644
index 00000000000..d191eb53ddf
--- /dev/null
+++ b/tooling/nargo_cli/tests/execution_success/brillig_cow/Nargo.toml
@@ -0,0 +1,6 @@
+[package]
+name = "brillig_cow"
+type = "bin"
+authors = [""]
+
+[dependencies]
diff --git a/tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml b/tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml
new file mode 100644
index 00000000000..6533d218b15
--- /dev/null
+++ b/tooling/nargo_cli/tests/execution_success/brillig_cow/Prover.toml
@@ -0,0 +1,7 @@
+original = [0, 1, 2, 3, 4]
+index = 2
+
+[expected_result]
+original = [0, 1, 2, 3, 4]
+modified_once = [0, 1, 27, 3, 4]
+modified_twice = [0, 1, 27, 27, 4]
diff --git a/tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr b/tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr
new file mode 100644
index 00000000000..7d847e085fe
--- /dev/null
+++ b/tooling/nargo_cli/tests/execution_success/brillig_cow/src/main.nr
@@ -0,0 +1,54 @@
+// Tests the copy on write optimization for arrays. We look for cases where we are modifying an array in place when we shouldn't.
+
+global ARRAY_SIZE = 5;
+
+struct ExecutionResult {
+    original: [Field; ARRAY_SIZE],
+    modified_once: [Field; ARRAY_SIZE],
+    modified_twice: [Field; ARRAY_SIZE],
+}
+
+impl ExecutionResult {
+    fn is_equal(self, other: ExecutionResult) -> bool {
+        (self.original == other.original) &
+        (self.modified_once == other.modified_once) &
+        (self.modified_twice == other.modified_twice)
+    }
+}
+
+fn modify_in_inlined_constrained(original: [Field; ARRAY_SIZE], index: u64) -> ExecutionResult {
+    let mut modified = original;
+    
+    modified[index] = 27;
+
+    let modified_once = modified;
+
+    modified[index+1] = 27;
+
+    ExecutionResult {
+        original,
+        modified_once,
+        modified_twice: modified
+    }
+}
+
+unconstrained fn modify_in_unconstrained(original: [Field; ARRAY_SIZE], index: u64) -> ExecutionResult {
+    let mut modified = original;
+    
+    modified[index] = 27;
+
+    let modified_once = modified;
+
+    modified[index+1] = 27;
+
+    ExecutionResult {
+        original,
+        modified_once,
+        modified_twice: modified
+    }
+}
+
+unconstrained fn main(original: [Field; ARRAY_SIZE], index: u64, expected_result: ExecutionResult) {
+    assert(expected_result.is_equal(modify_in_unconstrained(original, index)));
+    assert(expected_result.is_equal(modify_in_inlined_constrained(original, index)));
+}