[VAN-540] Generalize mapped offset computation, add tests

circom/tests/subcmps/mapped3.circom

@@ -0,0 +1,50 @@
+pragma circom 2.0.0;
+// REQUIRES: circom
+// RUN: rm -rf %t && mkdir %t && %circom --llvm -o %t %s | sed -n 's/.*Written successfully:.* \(.*\)/\1/p' | xargs cat | FileCheck %s
+
+template ArrayOp(q) {
+    signal input inp[15];
+    signal output outp[15];
+
+    for (var i = 0; i < 15; i++) {
+        outp[i] <== inp[i] + q;
+    }
+}
+
+//CHECK-LABEL: define void @ArrayOp_{{[0-9]+}}_build
+//CHECK-SAME: ({ [0 x i256]*, i32 }* %{{.*}})
+//CHECK: alloca [30 x i256]
+//CHECK: %[[DIM_REG:.*]] = getelementptr { [0 x i256]*, i32 }, { [0 x i256]*, i32 }* %0, i32 0, i32 1
+//CHECK: store i32 15, i32* %{{.*}}[[DIM_REG]]
+
+template Wrapper() {
+    signal input inp[15];
+    signal output outp;
+
+    component m[4];
+
+    for (var q = 0; q < 4; q++) {
+        // This test exhibits the behavior because the array of different subcomponents
+        // (differentiated by the template parameter changing)
+        m[q] = ArrayOp(q);
+        for (var i = 0; i < 15; i++) {
+            m[q].inp[i] <== inp[i];
+        }
+    }
+
+    outp <== m[2].outp[3];
+}
+
+component main = Wrapper();
+
+//CHECK-LABEL: define void @Wrapper_{{[0-9]+}}_run
+//CHECK-SAME: ([0 x i256]* %{{.*}})
+//CHECK: %lvars = alloca [2 x i256]
+//COM: offset = (1 * (3 * 7)) + (2 * (7)) + (3) + 1 (since 0 is output) = 21 + 14 + 3 + 1 = 39
+//CHECK: store{{.*}}:{{.*}}; preds = %unrolled_loop{{.*}}
+//CHECK: %[[SUB_PTR:.*]] = getelementptr [4 x { [0 x i256]*, i32 }], [4 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0
+//CHECK: %[[SUBCMP:.*]] = load [0 x i256]*, [0 x i256]** %[[SUB_PTR]]
+//CHECK: %[[VAL_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %[[SUBCMP]], i32 0, i32 3
+//CHECK: %[[VAL:.*]] = load i256, i256* %[[VAL_PTR]]
+//CHECK: %[[OUTP_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i32 0
+//CHECK: store i256 %[[VAL]], i256* %[[OUTP_PTR]]

circom/tests/subcmps/mapped4.circom

@@ -0,0 +1,53 @@
+pragma circom 2.0.0;
+// REQUIRES: circom
+// RUN: rm -rf %t && mkdir %t && %circom --llvm -o %t %s | sed -n 's/.*Written successfully:.* \(.*\)/\1/p' | xargs cat | FileCheck %s
+
+template MatrixOp(q) {
+    signal input inp[5][3];
+    signal output outp[5][3];
+
+    for (var i = 0; i < 5; i++) {
+        for (var j = 0; j < 3; j++) {
+            outp[i][j] <== inp[i][j] + q;
+        }
+    }
+}
+
+//CHECK-LABEL: define void @MatrixOp_{{[0-9]+}}_build
+//CHECK-SAME: ({ [0 x i256]*, i32 }* %{{.*}})
+//CHECK: alloca [16 x i256]
+//CHECK: %[[DIM_REG:.*]] = getelementptr { [0 x i256]*, i32 }, { [0 x i256]*, i32 }* %0, i32 0, i32 1
+//CHECK: store i32 15, i32* %{{.*}}[[DIM_REG]]
+
+template Wrapper() {
+    signal input inp[5][3];
+    signal output outp;
+
+    component m[4];
+
+    for (var q = 0; q < 4; q++) {
+        // This test exhibits the behavior because the array of different subcomponents
+        // (differentiated by the template parameter changing)
+        m[q] = MatrixOp(q);
+        for (var i = 0; i < 5; i++) {
+            for (var j = 0; j < 3; j++) {
+                m[q].inp[i][j] <== inp[i][j];
+            }
+        }
+    }
+
+    outp <== m[2].outp[1][2];
+}
+
+component main = Wrapper();
+
+//CHECK-LABEL: define void @Wrapper_{{[0-9]+}}_run
+//CHECK-SAME: ([0 x i256]* %{{.*}})
+//CHECK: %lvars = alloca [3 x i256]
+//CHECK: store{{.*}}:{{.*}}; preds = %unrolled_loop{{.*}}
+//CHECK: %[[SUB_PTR:.*]] = getelementptr [4 x { [0 x i256]*, i32 }], [4 x { [0 x i256]*, i32 }]* %subcmps, i32 0, i32 2, i32 0
+//CHECK: %[[SUBCMP:.*]] = load [0 x i256]*, [0 x i256]** %[[SUB_PTR]]
+//CHECK: %[[VAL_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %[[SUBCMP]], i32 0, i32 5
+//CHECK: %[[VAL:.*]] = load i256, i256* %[[VAL_PTR]]
+//CHECK: %[[OUTP_PTR:.*]] = getelementptr [0 x i256], [0 x i256]* %0, i32 0, i32 0
+//CHECK: store i256 %[[VAL]], i256* %[[OUTP_PTR]]

circuit_passes/src/bucket_interpreter/mod.rs

@@ -12,6 +12,7 @@ use compiler::num_bigint::BigInt;
 use observer::InterpreterObserver;
 use program_structure::constants::UsefulConstants;
 use crate::bucket_interpreter::env::Env;
+use crate::bucket_interpreter::operations::compute_offset;
 use crate::bucket_interpreter::value::{JoinSemiLattice, Value};
 use crate::bucket_interpreter::value::Value::{KnownBigInt, KnownU32, Unknown};
 
@@ -244,19 +245,17 @@ impl<'a> BucketInterpreter<'a> {
                     },
                     LocationRule::Mapped { signal_code, indexes } => {
                         let mut acc_env = env;
-                        let map_access = self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code].offset;
+                        let io_def = &self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code];
+                        let map_access = io_def.offset;
                         if indexes.len() > 0 {
                             let mut indexes_values = vec![];
                             for i in indexes {
                                 let (val, new_env) = self.execute_instruction(i, acc_env, continue_observing);
                                 indexes_values.push(val.expect("Mapped location must produce a value!").get_u32());
                                 acc_env = new_env;
                             }
-                            if indexes.len() == 1 {
-                                (map_access + indexes_values[0], acc_env)
-                            } else {
-                                todo!()
-                            }
+                            let offset = compute_offset(&indexes_values, &io_def.lengths);
+                            (map_access + offset, acc_env)
                         } else {
                             (map_access, acc_env)
                         }
@@ -318,19 +317,17 @@ impl<'a> BucketInterpreter<'a> {
                     LocationRule::Mapped { signal_code, indexes } => {
                         let mut acc_env = env;
                         let name = Some(acc_env.get_subcmp_name(addr).clone());
-                        let map_access = self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code].offset;
+                        let io_def = &self.io_map[&acc_env.get_subcmp_template_id(addr)][*signal_code];
+                        let map_access = io_def.offset;
                         if indexes.len() > 0 {
                             let mut indexes_values = vec![];
                             for i in indexes {
                                 let (val, new_env) = self.execute_instruction(i, acc_env, continue_observing);
                                 indexes_values.push(val.expect("Mapped location must produce a value!").get_u32());
                                 acc_env = new_env;
                             }
-                            if indexes.len() == 1 {
-                                (map_access + indexes_values[0], acc_env, name)
-                            } else {
-                                todo!()
-                            }
+                            let offset = compute_offset(&indexes_values, &io_def.lengths);
+                            (map_access + offset, acc_env, name)
                         } else {
                             (map_access, acc_env, name)
                         }

circuit_passes/src/bucket_interpreter/operations.rs

@@ -40,3 +40,55 @@ pub fn compute_operation(bucket: &ComputeBucket, stack: &Vec<Value>, p: &BigInt)
     });
     computed_value
 }
+
+pub fn compute_offset(indexes: &Vec<usize>, lengths: &Vec<usize>) -> usize {
+    // Lengths are in order, i.e. arr[x][y] => [x, y], same with indices
+    // arr[x][y] is x arrays of length y, laid out sequentially
+    if indexes.len() != lengths.len() {
+        // I did check, both "indexes" and "indices" are valid plurals
+        panic!("must have the same number of indexes and array lengths!");
+    }
+    let mut total_offset = indexes.last().copied().expect("must contain some indexes!");
+    let mut size_multiplier = lengths.last().copied().expect("must contain some array lengths!");
+    for i in (0..lengths.len()-1).rev() {
+        total_offset += indexes[i] * size_multiplier;
+        size_multiplier *= lengths[i];
+    }
+    total_offset
+}
+
+#[cfg(test)]
+mod test {
+    use super::compute_offset;
+
+    #[test]
+    fn test_expected_offset() {
+        let offset = compute_offset(&vec![1, 1], &vec![5, 3]);
+        assert_eq!(4, offset);
+    }
+
+    #[test]
+    fn test_offsets() {
+        let lengths = vec![5, 3, 7];
+        for i in 0..lengths[0] {
+            for j in 0..lengths[1] {
+                for k in 0..lengths[2] {
+                    let offset = compute_offset(&vec![i, j, k], &lengths);
+                    assert_eq!((i * 21) + (j * 7) + k, offset, "i={}, j={}, k={}", i, j, k);
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_increments() {
+        let lengths = vec![5, 7];
+        for i in 0..lengths[0] {
+            for j in 0..lengths[1]-1 {
+                let offset = compute_offset(&vec![i, j], &lengths);
+                let next_offset = compute_offset(&vec![i, j + 1], &lengths);
+                assert_eq!(offset + 1, next_offset);
+            }
+        }
+    }
+}

circuit_passes/src/passes/mapped_to_indexed.rs

@@ -6,6 +6,7 @@ use compiler::intermediate_representation::ir_interface::*;
 use compiler::intermediate_representation::{InstructionPointer, UpdateId};
 use crate::bucket_interpreter::env::Env;
 use crate::bucket_interpreter::observer::InterpreterObserver;
+use crate::bucket_interpreter::operations::compute_offset;
 use crate::bucket_interpreter::value::Value::KnownU32;
 use crate::passes::CircuitTransformationPass;
 use crate::passes::memory::PassMemory;
@@ -35,23 +36,20 @@ impl MappedToIndexedPass {
 
         let mut acc_env = acc_env;
         let name = acc_env.get_subcmp_name(resolved_addr).clone();
-        let map_access = mem.io_map[&acc_env.get_subcmp_template_id(resolved_addr)][signal_code].offset;
+        let io_def = &mem.io_map[&acc_env.get_subcmp_template_id(resolved_addr)][signal_code];
+        let map_access = io_def.offset;
         if indexes.len() > 0 {
             let mut indexes_values = vec![];
             for i in indexes {
                 let (val, new_env) = interpreter.execute_instruction(i, acc_env, false);
                 indexes_values.push(val.expect("Mapped location must produce a value!").get_u32());
                 acc_env = new_env;
             }
-            if indexes.len() == 1 {
-                let value = map_access + indexes_values[0];
-                let mut unused = vec![];
-                LocationRule::Indexed {
-                    location: KnownU32(value).to_value_bucket(&mut unused).allocate(),
-                    template_header: Some(name),
-                }
-            } else {
-                todo!()
+            let offset = compute_offset(&indexes_values, &io_def.lengths);
+            let mut unused = vec![];
+            LocationRule::Indexed {
+                location: KnownU32(map_access + offset).to_value_bucket(&mut unused).allocate(),
+                template_header: Some(name),
             }
         } else {
             let mut unused = vec![];