Merge pull request #114 from lambdaclass/array-lowering

Array Lowering

crates/concrete_codegen_mlir/src/codegen.rs

@@ -500,6 +500,7 @@ fn compile_rvalue<'c: 'b, 'b>(
                     }
                     PlaceElem::Field(_) => todo!(),
                     PlaceElem::Index(_) => todo!(),
+                    PlaceElem::ConstantIndex(_) => todo!(),
                 }
             }
 
@@ -1012,6 +1013,7 @@ fn compile_store_place<'c: 'b, 'b>(
                 }
             }
             PlaceElem::Index(_) => todo!(),
+            PlaceElem::ConstantIndex(_) => todo!(),
         }
     }
 
@@ -1083,6 +1085,7 @@ fn compile_load_place<'c: 'b, 'b>(
                 }
             }
             PlaceElem::Index(_) => todo!(),
+            PlaceElem::ConstantIndex(_) => todo!(),
         }
     }
 

crates/concrete_ir/src/lib.rs

@@ -201,6 +201,8 @@ pub enum PlaceElem {
     Field(FieldIndex),
     /// array index
     Index(LocalIndex),
+    /// constant array index
+    ConstantIndex(u64),
 }
 
 /// A local, akin to a variable, it can be user defined or compiler-introduced.

crates/concrete_ir/src/lowering.rs

@@ -861,7 +861,27 @@ fn find_expression_type(builder: &mut FnBodyBuilder, info: &Expression) -> Optio
             })
         }
         Expression::Cast(_, _, _) => todo!(),
-        Expression::ArrayInit(_) => todo!(),
+        Expression::ArrayInit(info) => {
+            let first_element = info.values.first()?;
+
+            let first_type = find_expression_type(builder, first_element)?;
+
+            let length = info.values.len() as u64;
+
+            Some(Ty {
+                span: Some(info.span),
+                kind: TyKind::Array(
+                    Box::new(first_type),
+                    Box::new(ConstData {
+                        ty: Ty {
+                            span: None,
+                            kind: TyKind::Uint(UintTy::U64),
+                        },
+                        data: ConstKind::Value(ValueTree::Leaf(ConstValue::U64(length))),
+                    }),
+                ),
+            })
+        }
     }
 }
 
@@ -1036,7 +1056,74 @@ fn lower_expression(
 
             (rvalue, new_ty, *span)
         }
-        Expression::ArrayInit(_) => todo!(),
+        Expression::ArrayInit(info) => {
+            let element_type_hint = type_hint.and_then(|type_hint| match type_hint.kind {
+                TyKind::Array(type_hint, _) => Some(*type_hint),
+                _ => None,
+            });
+
+            let mut values = info.values.iter().enumerate();
+
+            // Extract the first value from the array init. It's type will be used as type hint for the following elements.
+            let (first_idx, first_element) = values.next().expect("array init cannot be empty");
+            let (first_value, element_type, _element_span) =
+                lower_expression(builder, first_element, element_type_hint)?;
+
+            let length = info.values.len() as u64;
+
+            let ty = Ty {
+                span: Some(info.span),
+                kind: TyKind::Array(
+                    Box::new(element_type.clone()),
+                    Box::new(ConstData {
+                        ty: Ty {
+                            span: None,
+                            kind: TyKind::Uint(UintTy::U64),
+                        },
+                        data: ConstKind::Value(ValueTree::Leaf(ConstValue::U64(length))),
+                    }),
+                ),
+            };
+
+            // Create and init local for the array init expression.
+            let array_local = builder.add_local(Local::temp(ty.clone()));
+            let place = Place {
+                local: array_local,
+                projection: Default::default(),
+            };
+            builder.statements.push(Statement {
+                span: None,
+                kind: StatementKind::StorageLive(array_local),
+            });
+
+            // Assign the first value of the expression
+            let mut first_place = place.clone();
+            first_place
+                .projection
+                .push(PlaceElem::ConstantIndex(first_idx as u64));
+            builder.statements.push(Statement {
+                span: Some(info.span),
+                kind: StatementKind::Assign(first_place, first_value),
+            });
+
+            // Loop over the remaining values and assign them
+            for (idx, element) in values {
+                let mut element_place = place.clone();
+                element_place
+                    .projection
+                    .push(PlaceElem::ConstantIndex(idx as u64));
+
+                let (value, _value_ty, _field_span) =
+                    lower_expression(builder, element, Some(element_type.clone()))?;
+
+                builder.statements.push(Statement {
+                    span: Some(info.span),
+                    kind: StatementKind::Assign(element_place, value),
+                });
+            }
+
+            (Rvalue::Use(Operand::Place(place)), ty, info.span)
+        }
     })
 }
 
@@ -1428,7 +1515,35 @@ pub fn lower_path(
                     ty = struct_body.variants[idx].ty.clone();
                 }
             }
-            PathSegment::ArrayIndex(_, _) => todo!(),
+            PathSegment::ArrayIndex(expression, _) => {
+                while let TyKind::Ref(inner, _) = ty.kind {
+                    projection.push(PlaceElem::Deref);
+                    ty = *inner;
+                }
+
+                if let TyKind::Array(element_type, _) = ty.kind {
+                    // Assign the index expression to a temporary local
+                    let (index, index_ty) = lower_value_expr(builder, expression, None)?;
+                    let index_local = builder.add_temp_local(index_ty.kind);
+                    let index_place = Place {
+                        local: index_local,
+                        projection: vec![],
+                    };
+                    builder.statements.push(Statement {
+                        span: None,
+                        kind: StatementKind::StorageLive(index_local),
+                    });
+                    builder.statements.push(Statement {
+                        span: None,
+                        kind: StatementKind::Assign(index_place.clone(), index),
+                    });
+
+                    // Use the local's value as index of the array
+                    projection.push(PlaceElem::Index(index_local));
+
+                    ty = *element_type;
+                }
+            }
         }
     }
 

examples/arrays.con

@@ -0,0 +1,11 @@
+mod Example {
+    fn main() -> i32 {
+        let array: [i32; 4] = [1, 2, 3, 4];
+        let nested_array: [[i32; 2]; 2] = [[1, 2], [3, 4]];
+
+        let a: i32 = array[1];
+        let b: i32 = nested_array[1][0];
+
+        return a + b;
+    }
+}