constraints modify array and struct values in place

simple.program

@@ -1,6 +1,10 @@
-struct Foo {
-    bool x
+struct Point {
+    field x
+    field y
 }
-def main() -> (bool):
-    Foo f = Foo {x: true}
-    return f.x
+
+def main() -> ():
+    Point p = Point { x: 1, y: 1 }
+    p.x = 2
+    p.y = 2
+    return p

src/aleo_program/constraints.rs

@@ -1,5 +1,5 @@
 use crate::aleo_program::{
-    BooleanExpression, BooleanSpreadOrExpression, Expression, FieldExpression,
+    Assignee, BooleanExpression, BooleanSpreadOrExpression, Expression, FieldExpression,
     FieldRangeOrExpression, FieldSpreadOrExpression, Function, Program, Statement, Struct,
     StructMember, Type, Variable,
 };
@@ -75,8 +75,7 @@ impl fmt::Display for ResolvedValue {
             }
             ResolvedValue::Function(ref _function) => {
                 unimplemented!("cannot return function definition in program")
-            }
-            // _ => unimplemented!("display not impl for value"),
+            } // _ => unimplemented!("display not impl for value"),
         }
     }
 }
@@ -695,20 +694,146 @@ impl ResolvedProgram {
         }
     }
 
+    fn resolve_assignee(&mut self, scope: String, assignee: Assignee) -> String {
+        match assignee {
+            Assignee::Variable(name) => new_scope_from_variable(scope, &name),
+            Assignee::Array(array, _index) => self.resolve_assignee(scope, *array),
+            Assignee::StructMember(struct_variable, _member) => {
+                self.resolve_assignee(scope, *struct_variable)
+            }
+        }
+    }
+
+    // fn modify_array<F: Field + PrimeField, CS: ConstraintSystem<F>>(
+    //     &mut self,
+    //     cs: &mut CS,
+    //     scope: String,
+    //     assignee: Assignee,
+    //     expression: Expression,
+    // )
+
     fn enforce_definition_statement<F: Field + PrimeField, CS: ConstraintSystem<F>>(
         &mut self,
         cs: &mut CS,
         scope: String,
-        name: Variable,
+        assignee: Assignee,
         expression: Expression,
     ) {
-        // Evaluate the expression in the current function scope
-        let result = self.enforce_expression(cs, scope.clone(), expression);
+        // Create or modify the lhs variable in the current function scope
+        match assignee {
+            Assignee::Variable(name) => {
+                // Store the variable in the current scope
+                let definition_name = new_scope_from_variable(scope.clone(), &name);
+                let result = self.enforce_expression(cs, scope, expression);
 
-        // Store the definition in the current scope
-        let definition_name = new_scope_from_variable(scope, &name);
+                self.store(definition_name, result);
+            }
+            Assignee::Array(array, index_expression) => {
+                // Evaluate the rhs expression in the current function scope
+                let result = &mut self.enforce_expression(cs, scope.clone(), expression);
 
-        self.store(definition_name, result);
+                // Check that array exists
+                let expected_array_name = self.resolve_assignee(scope.clone(), *array);
+
+                // Resolve index so we know if we are assigning to a single value or a range of values
+                match index_expression {
+                    FieldRangeOrExpression::FieldExpression(index) => {
+                        let index = self.enforce_index(cs, scope.clone(), index);
+
+                        // Modify the single value of the array in place
+                        match self.get_mut(&expected_array_name) {
+                            Some(value) => match (value, result) {
+                                (
+                                    ResolvedValue::FieldElementArray(old),
+                                    ResolvedValue::FieldElement(new),
+                                ) => {
+                                    old[index] = new.to_owned();
+                                }
+                                (ResolvedValue::BooleanArray(old), ResolvedValue::Boolean(new)) => {
+                                    old[index] = new.to_owned();
+                                }
+                                _ => {
+                                    unimplemented!("Cannot assign single index to array of values ")
+                                }
+                            },
+                            None => unimplemented!(
+                                "tried to assign to unknown array {}",
+                                expected_array_name
+                            ),
+                        }
+                    }
+                    FieldRangeOrExpression::Range(from, to) => {
+                        let from_index = match from {
+                            Some(expression) => self.enforce_index(cs, scope.clone(), expression),
+                            None => 0usize,
+                        };
+                        let to_index_option = match to {
+                            Some(expression) => {
+                                Some(self.enforce_index(cs, scope.clone(), expression))
+                            }
+                            None => None,
+                        };
+
+                        // Modify the range of values of the array in place
+                        match self.get_mut(&expected_array_name) {
+                            Some(value) => match (value, result) {
+                                (
+                                    ResolvedValue::FieldElementArray(old),
+                                    ResolvedValue::FieldElementArray(new),
+                                ) => {
+                                    let to_index = to_index_option.unwrap_or(old.len());
+                                    old.splice(from_index..to_index, new.iter().cloned());
+                                }
+                                (
+                                    ResolvedValue::BooleanArray(old),
+                                    ResolvedValue::BooleanArray(new),
+                                ) => {
+                                    let to_index = to_index_option.unwrap_or(old.len());
+                                    old.splice(from_index..to_index, new.iter().cloned());
+                                }
+                                _ => unimplemented!(
+                                    "Cannot assign a range of array values to single value"
+                                ),
+                            },
+                            None => unimplemented!(
+                                "tried to assign to unknown array {}",
+                                expected_array_name
+                            ),
+                        }
+                    }
+                }
+            }
+            Assignee::StructMember(struct_variable, struct_member) => {
+                // Check that struct exists
+                let expected_struct_name = self.resolve_assignee(scope.clone(), *struct_variable);
+
+                match self.get_mut(&expected_struct_name) {
+                    Some(value) => match value {
+                        ResolvedValue::StructExpression(_variable, members) => {
+                            // Modify the struct member in place
+                            let matched_member = members
+                                .into_iter()
+                                .find(|member| member.variable == struct_member);
+                            match matched_member {
+                                Some(mut member) => member.expression = expression,
+                                None => unimplemented!(
+                                    "struct member {} does not exist in {}",
+                                    struct_member,
+                                    expected_struct_name
+                                ),
+                            }
+                        }
+                        _ => unimplemented!(
+                            "tried to assign to unknown struct {}",
+                            expected_struct_name
+                        ),
+                    },
+                    None => {
+                        unimplemented!("tried to assign to unknown struct {}", expected_struct_name)
+                    }
+                }
+            }
+        };
     }
 
     fn enforce_return_statement<F: Field + PrimeField, CS: ConstraintSystem<F>>(

src/aleo_program/types.rs

@@ -98,11 +98,18 @@ pub enum Expression {
     FunctionCall(Box<Expression>, Vec<Expression>),
 }
 
+#[derive(Debug, Clone)]
+pub enum Assignee {
+    Variable(Variable),
+    Array(Box<Assignee>, FieldRangeOrExpression),
+    StructMember(Box<Assignee>, Variable),
+}
+
 /// Program statement that defines some action (or expression) to be carried out.
 #[derive(Clone)]
 pub enum Statement {
     // Declaration(Variable),
-    Definition(Variable, Expression),
+    Definition(Assignee, Expression),
     For(Variable, FieldExpression, FieldExpression, Vec<Statement>),
     Return(Vec<Expression>),
 }

src/aleo_program/types_display.rs

@@ -5,9 +5,9 @@
 //! @date 2020
 
 use crate::aleo_program::{
-    BooleanExpression, BooleanSpread, BooleanSpreadOrExpression, Expression, FieldExpression,
-    FieldRangeOrExpression, FieldSpread, FieldSpreadOrExpression, Function, FunctionName,
-    Parameter, Statement, Struct, StructField, Type, Variable,
+    Assignee, BooleanExpression, BooleanSpread, BooleanSpreadOrExpression, Expression,
+    FieldExpression, FieldRangeOrExpression, FieldSpread, FieldSpreadOrExpression, Function,
+    FunctionName, Parameter, Statement, Struct, StructField, Type, Variable,
 };
 
 use std::fmt;
@@ -163,6 +163,19 @@ impl<'ast> fmt::Display for Expression {
         }
     }
 }
+
+impl fmt::Display for Assignee {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            Assignee::Variable(ref variable) => write!(f, "{}", variable),
+            Assignee::Array(ref array, ref index) => write!(f, "{}[{}]", array, index),
+            Assignee::StructMember(ref struct_variable, ref member) => {
+                write!(f, "{}.{}", struct_variable, member)
+            }
+        }
+    }
+}
+
 impl fmt::Display for Statement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {

src/aleo_program/types_from.rs

@@ -336,7 +336,7 @@ impl<'ast> From<ast::PostfixExpression<'ast>> for types::Expression {
                     Box::new(acc),
                     types::Variable::from(struct_member.variable),
                 ),
-                ast::Access::Select(array) => types::Expression::ArrayAccess(
+                ast::Access::Array(array) => types::Expression::ArrayAccess(
                     Box::new(acc),
                     types::FieldRangeOrExpression::from(array.expression),
                 ),
@@ -364,10 +364,37 @@ impl<'ast> From<ast::Expression<'ast>> for types::Expression {
     }
 }
 
+impl<'ast> From<ast::Variable<'ast>> for types::Assignee {
+    fn from(variable: ast::Variable<'ast>) -> Self {
+        types::Assignee::Variable(types::Variable::from(variable))
+    }
+}
+
+impl<'ast> From<ast::Assignee<'ast>> for types::Assignee {
+    fn from(assignee: ast::Assignee<'ast>) -> Self {
+        let variable = types::Assignee::from(assignee.variable);
+
+        // we start with the id, and we fold the array of accesses by wrapping the current value
+        assignee
+            .accesses
+            .into_iter()
+            .fold(variable, |acc, access| match access {
+                ast::AssigneeAccess::Array(array) => types::Assignee::Array(
+                    Box::new(acc),
+                    types::FieldRangeOrExpression::from(array.expression),
+                ),
+                ast::AssigneeAccess::Member(struct_member) => types::Assignee::StructMember(
+                    Box::new(acc),
+                    types::Variable::from(struct_member.variable),
+                ),
+            })
+    }
+}
+
 impl<'ast> From<ast::AssignStatement<'ast>> for types::Statement {
     fn from(statement: ast::AssignStatement<'ast>) -> Self {
         types::Statement::Definition(
-            types::Variable::from(statement.variable),
+            types::Assignee::from(statement.assignee),
             types::Expression::from(statement.expression),
         )
     }
@@ -543,7 +570,7 @@ impl<'ast> types::Expression {
 impl<'ast> From<ast::DefinitionStatement<'ast>> for types::Statement {
     fn from(statement: ast::DefinitionStatement<'ast>) -> Self {
         types::Statement::Definition(
-            types::Variable::from(statement.variable),
+            types::Assignee::from(statement.variable),
             types::Expression::from_type(statement.ty, statement.expression),
         )
     }

src/ast.rs

@@ -13,7 +13,6 @@ use pest::{
 };
 use pest_ast::FromPest;
 use std::fmt;
-use std::fmt::Formatter;
 
 #[derive(Parser)]
 #[grammar = "language.pest"]
@@ -219,7 +218,7 @@ pub enum Type<'ast> {
 }
 
 impl<'ast> fmt::Display for Type<'ast> {
-    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match *self {
             Type::Basic(ref _ty) => write!(f, "basic"),
             Type::Array(ref _ty) => write!(f, "array"),
@@ -387,6 +386,28 @@ pub enum RangeOrExpression<'ast> {
     Expression(Expression<'ast>),
 }
 
+impl<'ast> fmt::Display for RangeOrExpression<'ast> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            RangeOrExpression::Expression(ref expression) => write!(f, "{}", expression),
+            RangeOrExpression::Range(ref range) => write!(
+                f,
+                "{}..{}",
+                range
+                    .from
+                    .as_ref()
+                    .map(|e| e.0.to_string())
+                    .unwrap_or("".to_string()),
+                range
+                    .to
+                    .as_ref()
+                    .map(|e| e.0.to_string())
+                    .unwrap_or("".to_string())
+            ),
+        }
+    }
+}
+
 #[derive(Clone, Debug, FromPest, PartialEq)]
 #[pest_ast(rule(Rule::access_call))]
 pub struct CallAccess<'ast> {
@@ -414,8 +435,8 @@ pub struct MemberAccess<'ast> {
 #[derive(Clone, Debug, FromPest, PartialEq)]
 #[pest_ast(rule(Rule::access))]
 pub enum Access<'ast> {
+    Array(ArrayAccess<'ast>),
     Call(CallAccess<'ast>),
-    Select(ArrayAccess<'ast>),
     Member(MemberAccess<'ast>),
 }
 
@@ -428,6 +449,44 @@ pub struct PostfixExpression<'ast> {
     pub span: Span<'ast>,
 }
 
+#[derive(Clone, Debug, FromPest, PartialEq)]
+#[pest_ast(rule(Rule::assignee_access))]
+pub enum AssigneeAccess<'ast> {
+    Array(ArrayAccess<'ast>),
+    Member(MemberAccess<'ast>),
+}
+
+impl<'ast> fmt::Display for AssigneeAccess<'ast> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match *self {
+            AssigneeAccess::Array(ref array) => write!(f, "[{}]", array.expression),
+            AssigneeAccess::Member(ref member) => write!(f, ".{}", member.variable),
+        }
+    }
+}
+
+#[derive(Clone, Debug, FromPest, PartialEq)]
+#[pest_ast(rule(Rule::assignee))]
+pub struct Assignee<'ast> {
+    pub variable: Variable<'ast>,
+    pub accesses: Vec<AssigneeAccess<'ast>>,
+    #[pest_ast(outer())]
+    pub span: Span<'ast>,
+}
+
+impl<'ast> fmt::Display for Assignee<'ast> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.variable)?;
+        for (i, access) in self.accesses.iter().enumerate() {
+            write!(f, "{}", access)?;
+            if i < self.accesses.len() - 1 {
+                write!(f, ", ")?;
+            }
+        }
+        write!(f, "")
+    }
+}
+
 #[derive(Clone, Debug, FromPest, PartialEq)]
 #[pest_ast(rule(Rule::spread))]
 pub struct Spread<'ast> {
@@ -684,7 +743,7 @@ impl<'ast> FromPest<'ast> for Expression<'ast> {
 #[derive(Clone, Debug, FromPest, PartialEq)]
 #[pest_ast(rule(Rule::statement_assign))]
 pub struct AssignStatement<'ast> {
-    pub variable: Variable<'ast>,
+    pub assignee: Assignee<'ast>,
     pub expression: Expression<'ast>,
     #[pest_ast(outer())]
     pub span: Span<'ast>,
@@ -730,7 +789,7 @@ pub enum Statement<'ast> {
 
 impl<'ast> fmt::Display for AssignStatement<'ast> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{} = {}", self.variable, self.expression)
+        write!(f, "{} = {}", self.assignee, self.expression)
     }
 }
 

src/language.pest

@@ -88,6 +88,9 @@ access_member = { "." ~ variable }
 access = { access_array | access_call | access_member }
 expression_postfix = { variable ~ access+ }
 
+assignee = { variable ~ assignee_access* }
+assignee_access = { access_array | access_member }
+
 spread = { "..." ~ expression }
 spread_or_expression = { spread | expression }
 
@@ -132,7 +135,7 @@ expression_tuple = _{ (expression ~ ("," ~ expression)*)? }
 
 /// Statements
 
-statement_assign = { variable ~ "=" ~ expression }
+statement_assign = { assignee ~ "=" ~ expression }
 statement_definition = { ty ~ variable ~ "=" ~ expression }
 statement_return = { "return" ~ expression_tuple }
 statement_for = { "for" ~ variable ~ "in" ~ expression ~ ".." ~ expression ~ "do" ~ NEWLINE* ~ statement* ~ "endfor"}