feat: add utils::collapse

noir-projects/aztec-nr/aztec/src/lib.nr

@@ -12,5 +12,6 @@ mod prelude;
 mod public_storage;
 mod encrypted_logs;
 use dep::protocol_types;
+mod utils;
 
 mod test;

noir-projects/aztec-nr/aztec/src/utils.nr

@@ -0,0 +1,84 @@
+use dep::protocol_types::traits::Eq;
+
+mod test;
+
+// Collapses an array of Options with sparse Some values into a BoundedVec, essentially unwrapping the Options and
+// removing the None values. For example, given:
+//   input: [some(3), none(), some(1)]
+// this returns 
+//   collapsed: [3, 1]
+pub fn collapse<T, N>(input: [Option<T>; N]) -> BoundedVec<T, N> where T: Eq {
+    let (collapsed, collapsed_to_input_index_mapping) = get_collapse_hints(input);
+    verify_collapse_hints(input, collapsed, collapsed_to_input_index_mapping);
+    collapsed
+}
+
+fn verify_collapse_hints<T, N>(
+    input: [Option<T>; N],
+    collapsed: BoundedVec<T, N>,
+    collapsed_to_input_index_mapping: BoundedVec<u32, N>
+) where T: Eq {
+    // collapsed should be a BoundedVec with all the non-none elements in input, in the same order. We need to lay down
+    // multiple constraints to guarantee this.
+
+    // First we check that the number of elements is correct
+    let mut count = 0;
+    for i in 0..N {
+        if input[i].is_some() {
+            count += 1;
+        }
+    }
+    assert_eq(count, collapsed.len(), "Wrong collapsed vec length");
+
+    // Then we check that all elements exist in the original array, and are in the same order. To do this we use the
+    // auxiliary collapsed_to_input_index_mapping array, which at index n contains the index in the input array that 
+    // corresponds to the collapsed entry at index n.
+    // Example:
+    //  - input: [some(3), none(), some(1)]
+    //  - collapsed: [3, 1]
+    //  - collapsed_to_input_index_mapping: [0, 2]
+    // These two arrays should therefore have the same length.
+    assert_eq(collapsed.len(), collapsed_to_input_index_mapping.len(), "Collapse hint vec length mismatch");
+
+    // We now look at each collapsed entry and check that there is a valid equal entry in the input array.
+    let mut last_index = Option::none();
+    for i in 0..N {
+        if i < collapsed.len() {
+            let input_index = collapsed_to_input_index_mapping.get_unchecked(i);
+            assert(input_index < N, "Out of bounds index hint");
+
+            assert_eq(collapsed.get_unchecked(i), input[input_index].unwrap(), "Wrong collapsed vec content");
+
+            // By requiring increasing input indices, we both guarantee that we're not looking at the same input
+            // element more than once, and that we're going over them in the original order.
+            if last_index.is_some() {
+                assert(input_index > last_index.unwrap_unchecked(), "Wrong collapsed vec order");
+            }
+            last_index = Option::some(input_index);
+        } else {
+            // We don't technically need to check past the length of the BoundedVec since those slots should not be
+            // accessible, but its fairly cheap to prevent dirty data from being stored there.
+            assert_eq(collapsed.get_unchecked(i), dep::std::unsafe::zeroed(), "Dirty collapsed vec");
+        }
+    }
+    // We now know that:
+    //  - all values in the collapsed array exist in the input array
+    //  - the order of the collapsed values is the same as in the input array
+    //  - no input value is present more than once in the collapsed array
+    //  - the number of elements in the collapsed array is the same as in the input array.
+    // Therefore, the collapsed array is correct.
+}
+
+unconstrained fn get_collapse_hints<T, N>(input: [Option<T>; N]) -> (BoundedVec<T, N>, BoundedVec<u32, N>) {
+    let mut collapsed: BoundedVec<T, N> = BoundedVec::new();
+    let mut collapsed_to_input_index_mapping: BoundedVec<u32, N> = BoundedVec::new();
+
+    for i in 0..N {
+        if input[i].is_some() {
+            collapsed.push(input[i].unwrap_unchecked());
+            collapsed_to_input_index_mapping.push(i);
+        }
+    }
+
+    (collapsed, collapsed_to_input_index_mapping)
+}

noir-projects/aztec-nr/aztec/src/utils/test.nr

@@ -0,0 +1,124 @@
+use crate::utils::{collapse, verify_collapse_hints};
+
+#[test]
+fn collapse_empty_array() {
+    let original: [Option<Field>; 2] = [Option::none(), Option::none()];
+    let collapsed = collapse(original);
+
+    assert_eq(collapsed.len(), 0);
+}
+
+#[test]
+fn collapse_non_sparse_array() {
+    let original = [Option::some(7), Option::some(3), Option::none()];
+    let collapsed = collapse(original);
+
+    assert_eq(collapsed.len(), 2);
+    assert_eq(collapsed.get(0), 7);
+    assert_eq(collapsed.get(1), 3);
+}
+
+#[test]
+fn collapse_sparse_array() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = collapse(original);
+
+    assert_eq(collapsed.len(), 2);
+    assert_eq(collapsed.get(0), 7);
+    assert_eq(collapsed.get(1), 3);
+}
+
+#[test]
+fn collapse_array_front_padding() {
+    let original = [Option::none(), Option::none(), Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = collapse(original);
+
+    assert_eq(collapsed.len(), 2);
+    assert_eq(collapsed.get(0), 7);
+    assert_eq(collapsed.get(1), 3);
+}
+
+#[test]
+fn collapse_array_back_padding() {
+    let original = [Option::some(7), Option::none(), Option::some(3), Option::none(), Option::none()];
+    let collapsed = collapse(original);
+
+    assert_eq(collapsed.len(), 2);
+    assert_eq(collapsed.get(0), 7);
+    assert_eq(collapsed.get(1), 3);
+}
+
+#[test]
+fn verify_collapse_hints_good_hints() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([7, 3]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0, 2]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail_with="Wrong collapsed vec length")]
+fn verify_collapse_hints_wrong_length() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([7]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail_with="Collapse hint vec length mismatch")]
+fn verify_collapse_hints_hint_length_mismatch() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([7, 3]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail_with="Out of bounds index hint")]
+fn verify_collapse_hints_out_of_bounds_index_hint() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([7, 3]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0, 5]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail)]
+fn verify_collapse_hints_hint_to_none() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([7, 0]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0, 1]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail_with="Wrong collapsed vec content")]
+fn verify_collapse_hints_wrong_vec_content() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([7, 42]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0, 2]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail_with="Wrong collapsed vec order")]
+fn verify_collapse_hints_wrong_vec_order() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+    let collapsed = BoundedVec::from_array([3, 7]);
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([2, 0]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}
+
+#[test(should_fail_with="Dirty collapsed vec")]
+fn verify_collapse_hints_dirty_storage() {
+    let original = [Option::some(7), Option::none(), Option::some(3)];
+
+    let mut collapsed: BoundedVec<u32, 3> = BoundedVec::from_array([7, 3]);
+    collapsed.storage[2] = 1;
+
+    let collapsed_to_input_index_mapping = BoundedVec::from_array([0, 2]);
+
+    verify_collapse_hints(original, collapsed, collapsed_to_input_index_mapping);
+}