Support array access expressions in AST and visitors

compiler/ast/src/access/array_access.rs

@@ -0,0 +1,42 @@
+// Copyright (C) 2019-2023 Aleo Systems Inc.
+// This file is part of the Leo library.
+
+// The Leo library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// The Leo library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
+
+use crate::{Expression, Node, NodeID};
+use leo_span::Span;
+
+use serde::{Deserialize, Serialize};
+use std::fmt;
+
+/// An array access expression, e.g., `foo[index]`.
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
+pub struct ArrayAccess {
+    /// An expression evaluating to some array type, e.g., `[false, true]`.
+    pub array: Box<Expression>,
+    /// The index to access in the array expression. E.g., `0` for `[false, true]` would yield `false`.
+    pub index: Box<Expression>,
+    /// The span for the entire expression `foo[index]`.
+    pub span: Span,
+    /// The ID of the node.
+    pub id: NodeID,
+}
+
+impl fmt::Display for ArrayAccess {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}.{}", self.array, self.index)
+    }
+}
+
+crate::simple_node_impl!(ArrayAccess);

compiler/ast/src/access/mod.rs

@@ -14,6 +14,9 @@
 // You should have received a copy of the GNU General Public License
 // along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
 
+mod array_access;
+pub use array_access::*;
+
 mod associated_constant_access;
 pub use associated_constant_access::*;
 

compiler/ast/src/expressions/access.rs

@@ -23,8 +23,8 @@ use std::fmt;
 /// An access expressions, extracting a smaller part out of a whole.
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub enum AccessExpression {
-    // /// An `array[index]` expression.
-    // Array(ArrayAccess),
+    /// An `array[index]` expression.
+    Array(ArrayAccess),
     // /// An expression accessing a range of an array.
     // ArrayRange(ArrayRangeAccess),
     /// Access to an associated variable of a struct e.g `u8::MAX`.
@@ -40,6 +40,7 @@ pub enum AccessExpression {
 impl Node for AccessExpression {
     fn span(&self) -> Span {
         match self {
+            AccessExpression::Array(n) => n.span(),
             AccessExpression::AssociatedConstant(n) => n.span(),
             AccessExpression::AssociatedFunction(n) => n.span(),
             AccessExpression::Member(n) => n.span(),
@@ -49,6 +50,7 @@ impl Node for AccessExpression {
 
     fn set_span(&mut self, span: Span) {
         match self {
+            AccessExpression::Array(n) => n.set_span(span),
             AccessExpression::AssociatedConstant(n) => n.set_span(span),
             AccessExpression::AssociatedFunction(n) => n.set_span(span),
             AccessExpression::Member(n) => n.set_span(span),
@@ -58,6 +60,7 @@ impl Node for AccessExpression {
 
     fn id(&self) -> NodeID {
         match self {
+            AccessExpression::Array(n) => n.id(),
             AccessExpression::AssociatedConstant(n) => n.id(),
             AccessExpression::AssociatedFunction(n) => n.id(),
             AccessExpression::Member(n) => n.id(),
@@ -67,6 +70,7 @@ impl Node for AccessExpression {
 
     fn set_id(&mut self, id: NodeID) {
         match self {
+            AccessExpression::Array(n) => n.set_id(id),
             AccessExpression::AssociatedConstant(n) => n.set_id(id),
             AccessExpression::AssociatedFunction(n) => n.set_id(id),
             AccessExpression::Member(n) => n.set_id(id),
@@ -79,6 +83,7 @@ impl fmt::Display for AccessExpression {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         use AccessExpression::*;
         match self {
+            Array(access) => access.fmt(f),
             AssociatedConstant(access) => access.fmt(f),
             AssociatedFunction(access) => access.fmt(f),
             Member(access) => access.fmt(f),

compiler/ast/src/passes/reconstructor.rs

@@ -45,6 +45,12 @@ pub trait ExpressionReconstructor {
     fn reconstruct_access(&mut self, input: AccessExpression) -> (Expression, Self::AdditionalOutput) {
         (
             Expression::Access(match input {
+                AccessExpression::Array(array) => AccessExpression::Array(ArrayAccess {
+                    array: Box::new(self.reconstruct_expression(*array.array).0),
+                    index: Box::new(self.reconstruct_expression(*array.index).0),
+                    span: array.span,
+                    id: array.id,
+                }),
                 AccessExpression::AssociatedFunction(function) => {
                     AccessExpression::AssociatedFunction(AssociatedFunction {
                         ty: function.ty,

compiler/ast/src/passes/visitor.rs

@@ -45,6 +45,10 @@ pub trait ExpressionVisitor<'a> {
 
     fn visit_access(&mut self, input: &'a AccessExpression, additional: &Self::AdditionalInput) -> Self::Output {
         match input {
+            AccessExpression::Array(array) => {
+                self.visit_expression(&array.array, additional);
+                self.visit_expression(&array.index, additional);
+            }
             AccessExpression::AssociatedFunction(function) => {
                 function.arguments.iter().for_each(|arg| {
                     self.visit_expression(arg, &Default::default());

tests/tests/compiler/array/access_array_with_loop_counter.leo

@@ -6,7 +6,7 @@ expectation: Pass
 
 program test.aleo {
     transition foo(a: [bool; 4]) {
-        for i in 0u32..4u32 {
+        for i: i32 in 0u32..4u32 {
             assert(a[i]);
         }
     }

tests/tests/compiler/array/array_with_units_fail.leo

@@ -5,6 +5,6 @@ expectation: Pass
 
 program test.aleo {
     transition foo() -> bool {
-        let bar: [(), 2] = [(), ()];
+        let bar: [(); 2] = [(), ()];
     }
 }

tests/tests/execution/array_sum.leo

@@ -16,7 +16,7 @@ program test.aleo {
 
     transition sum_with_loop(a: [u64; 4]) -> u64 {
         let sum: u64 = 0u64;
-        for i in 0u8..4u8 {
+        for i: u8 in 0u8..4u8 {
             sum += a[i];
         }
         return sum;