Merge pull request #1007 from AleoHQ/dyn-array-index-assignment

Dynamic array index assignment

asg/src/statement/assign.rs

@@ -213,7 +213,7 @@ impl<'a> FromAst<'a, leo_ast::AssignStatement> for &'a Statement<'a> {
         let statement = scope.context.alloc_statement(Statement::Assign(AssignStatement {
             parent: Cell::new(None),
             span: Some(statement.span.clone()),
-            operation: statement.operation.clone(),
+            operation: statement.operation,
             target_variable: Cell::new(variable),
             target_accesses,
             value: Cell::new(value),
@@ -231,7 +231,7 @@ impl<'a> FromAst<'a, leo_ast::AssignStatement> for &'a Statement<'a> {
 impl<'a> Into<leo_ast::AssignStatement> for &AssignStatement<'a> {
     fn into(self) -> leo_ast::AssignStatement {
         leo_ast::AssignStatement {
-            operation: self.operation.clone(),
+            operation: self.operation,
             assignee: leo_ast::Assignee {
                 identifier: self.target_variable.get().borrow().name.clone(),
                 accesses: self

ast/src/reducer/canonicalization.rs

@@ -347,7 +347,7 @@ impl Canonicalizer {
                 Statement::Assign(AssignStatement {
                     assignee,
                     value,
-                    operation: assign.operation.clone(),
+                    operation: assign.operation,
                     span: assign.span.clone(),
                 })
             }

ast/src/reducer/reconstructing_reducer.rs

@@ -325,7 +325,7 @@ pub trait ReconstructingReducer {
         value: Expression,
     ) -> Result<AssignStatement, ReducerError> {
         Ok(AssignStatement {
-            operation: assign.operation.clone(),
+            operation: assign.operation,
             assignee,
             value,
             span: assign.span.clone(),

ast/src/statements/assign/mod.rs

@@ -22,7 +22,7 @@ use std::fmt;
 mod assignee;
 pub use assignee::*;
 
-#[derive(Clone, PartialEq, Eq, Serialize, Deserialize, Debug)]
+#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Debug)]
 pub enum AssignOperation {
     Assign,
     Add,

compiler/src/expression/array/access.rs

@@ -37,6 +37,33 @@ use snarkvm_gadgets::utilities::{
 use snarkvm_r1cs::ConstraintSystem;
 
 impl<'a, F: PrimeField, G: GroupType<F>> ConstrainedProgram<'a, F, G> {
+    pub fn array_bounds_check<CS: ConstraintSystem<F>>(
+        &mut self,
+        cs: &mut CS,
+        index_resolved: &Integer,
+        array_len: u32,
+        span: &Span,
+    ) -> Result<(), ExpressionError> {
+        let bounds_check = evaluate_lt::<F, G, CS>(
+            cs,
+            ConstrainedValue::Integer(index_resolved.clone()),
+            ConstrainedValue::Integer(Integer::new(
+                &ConstInt::U32(array_len).cast_to(&index_resolved.get_type()),
+            )),
+            span,
+        )?;
+        let bounds_check = match bounds_check {
+            ConstrainedValue::Boolean(b) => b,
+            _ => unimplemented!("illegal non-Integer returned from lt"),
+        };
+        let namespace_string = format!("evaluate array access bounds {}:{}", span.line_start, span.col_start);
+        let mut unique_namespace = cs.ns(|| namespace_string);
+        bounds_check
+            .enforce_equal(&mut unique_namespace, &Boolean::Constant(true))
+            .map_err(|e| ExpressionError::cannot_enforce("array bounds check".to_string(), e, span))?;
+        Ok(())
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn enforce_array_access<CS: ConstraintSystem<F>>(
         &mut self,
@@ -65,21 +92,7 @@ impl<'a, F: PrimeField, G: GroupType<F>> ConstrainedProgram<'a, F, G> {
                     .len()
                     .try_into()
                     .map_err(|_| ExpressionError::array_length_out_of_bounds(span))?;
-                let bounds_check = evaluate_lt::<F, G, CS>(
-                    cs,
-                    ConstrainedValue::Integer(index_resolved.clone()),
-                    ConstrainedValue::Integer(Integer::new(&ConstInt::U32(array_len))),
-                    span,
-                )?;
-                let bounds_check = match bounds_check {
-                    ConstrainedValue::Boolean(b) => b,
-                    _ => unimplemented!("illegal non-Integer returned from lt"),
-                };
-                let namespace_string = format!("evaluate array access bounds {}:{}", span.line_start, span.col_start);
-                let mut unique_namespace = cs.ns(|| namespace_string);
-                bounds_check
-                    .enforce_equal(&mut unique_namespace, &Boolean::Constant(true))
-                    .map_err(|e| ExpressionError::cannot_enforce("array bounds check".to_string(), e, span))?;
+                self.array_bounds_check(cs, &&index_resolved, array_len, span)?;
             }
 
             let mut current_value = array.pop().unwrap();

compiler/src/function/function.rs

@@ -81,13 +81,8 @@ impl<'a, F: PrimeField, G: GroupType<F>> ConstrainedProgram<'a, F, G> {
                     .get(self_var.borrow().id)
                     .expect("no self variable found in mut self context")
                     .clone();
-                if let Some(assignable_target) = self.resolve_mut_ref(cs, target)? {
-                    if assignable_target.len() != 1 {
-                        panic!("found tuple as a self assignment target");
-                    }
-                    let assignable_target = assignable_target.into_iter().next().unwrap();
-                    *assignable_target = new_self;
-                } else {
+
+                if !self.resolve_mut_ref(cs, target, new_self, &indicator)? {
                     // todo: we should report a warning for calling a mutable function on an effectively copied self (i.e. wasn't assignable `tempStruct {x: 5}.myMutSelfFunction()`)
                 }
             }

compiler/src/function/mod.rs

@@ -28,5 +28,4 @@ pub use self::main_function::*;
 pub mod result;
 pub use self::result::*;
 
-pub mod mut_target;
-pub use self::mut_target::*;
+mod mut_target;

compiler/src/function/mut_target.rs

@@ -16,116 +16,96 @@
 
 //! Resolves assignees in a compiled Leo program.
 
-use crate::{
-    errors::StatementError,
-    program::ConstrainedProgram,
-    value::ConstrainedValue,
-    GroupType,
-    ResolvedAssigneeAccess,
-};
+use std::cell::Cell;
+
+use crate::{errors::StatementError, program::ConstrainedProgram, value::ConstrainedValue, GroupType};
 use leo_asg::{
     ArrayAccessExpression,
     ArrayRangeAccessExpression,
+    AssignAccess,
+    AssignOperation,
+    AssignStatement,
     CircuitAccessExpression,
     Expression,
     Node,
-    Span,
     TupleAccessExpression,
     Variable,
 };
 
 use snarkvm_fields::PrimeField;
+use snarkvm_gadgets::utilities::boolean::Boolean;
 use snarkvm_r1cs::ConstraintSystem;
 
 impl<'a, F: PrimeField, G: GroupType<F>> ConstrainedProgram<'a, F, G> {
-    fn prepare_mut_access<CS: ConstraintSystem<F>>(
-        &mut self,
-        cs: &mut CS,
+    fn prepare_mut_access(
+        out: &mut Vec<AssignAccess<'a>>,
         expr: &'a Expression<'a>,
-        span: &Span,
-        output: &mut Vec<ResolvedAssigneeAccess>,
-    ) -> Result<Option<Variable<'a>>, StatementError> {
+    ) -> Result<Option<&'a Variable<'a>>, StatementError> {
         match expr {
             Expression::ArrayRangeAccess(ArrayRangeAccessExpression { array, left, right, .. }) => {
-                let inner = self.prepare_mut_access(cs, array.get(), span, output)?;
-                let start_index = left
-                    .get()
-                    .map(|start| self.enforce_index(cs, start, span))
-                    .transpose()?
-                    .map(|x| {
-                        x.to_usize()
-                            .ok_or_else(|| StatementError::array_assign_index_const(span))
-                    })
-                    .transpose()?;
-                let stop_index = right
-                    .get()
-                    .map(|stop| self.enforce_index(cs, stop, span))
-                    .transpose()?
-                    .map(|x| {
-                        x.to_usize()
-                            .ok_or_else(|| StatementError::array_assign_index_const(span))
-                    })
-                    .transpose()?;
-
-                output.push(ResolvedAssigneeAccess::ArrayRange(start_index, stop_index));
+                let inner = Self::prepare_mut_access(out, array.get())?;
+
+                out.push(AssignAccess::ArrayRange(left.clone(), right.clone()));
+
                 Ok(inner)
             }
             Expression::ArrayAccess(ArrayAccessExpression { array, index, .. }) => {
-                let inner = self.prepare_mut_access(cs, array.get(), span, output)?;
-                let index = self
-                    .enforce_index(cs, index.get(), span)?
-                    .to_usize()
-                    .ok_or_else(|| StatementError::array_assign_index_const(span))?;
+                let inner = Self::prepare_mut_access(out, array.get())?;
 
-                output.push(ResolvedAssigneeAccess::ArrayIndex(index));
+                out.push(AssignAccess::ArrayIndex(index.clone()));
                 Ok(inner)
             }
             Expression::TupleAccess(TupleAccessExpression { tuple_ref, index, .. }) => {
-                let inner = self.prepare_mut_access(cs, tuple_ref.get(), span, output)?;
+                let inner = Self::prepare_mut_access(out, tuple_ref.get())?;
 
-                output.push(ResolvedAssigneeAccess::Tuple(*index, span.clone()));
+                out.push(AssignAccess::Tuple(*index));
                 Ok(inner)
             }
             Expression::CircuitAccess(CircuitAccessExpression { target, member, .. }) => {
                 if let Some(target) = target.get() {
-                    let inner = self.prepare_mut_access(cs, target, span, output)?;
+                    let inner = Self::prepare_mut_access(out, target)?;
 
-                    output.push(ResolvedAssigneeAccess::Member(member.clone()));
+                    out.push(AssignAccess::Member(member.clone()));
                     Ok(inner)
                 } else {
                     Ok(None)
                 }
             }
-            Expression::VariableRef(variable_ref) => Ok(Some(variable_ref.variable.clone())),
+            Expression::VariableRef(variable_ref) => Ok(Some(variable_ref.variable)),
             _ => Ok(None), // not a valid reference to mutable variable, we copy
         }
     }
 
     // resolve a mutable reference from an expression
-    // return Ok(None) if no valid mutable reference, or Err(_) on more critical error
+    // return false if no valid mutable reference, or Err(_) on more critical error
     pub fn resolve_mut_ref<CS: ConstraintSystem<F>>(
         &mut self,
         cs: &mut CS,
         assignee: &'a Expression<'a>,
-    ) -> Result<Option<Vec<&mut ConstrainedValue<'a, F, G>>>, StatementError> {
-        let span = &assignee.span().cloned().unwrap_or_default();
-
+        target_value: ConstrainedValue<'a, F, G>,
+        indicator: &Boolean,
+    ) -> Result<bool, StatementError> {
         let mut accesses = vec![];
-        let target = self.prepare_mut_access(cs, assignee, span, &mut accesses)?;
+        let target = Self::prepare_mut_access(&mut accesses, assignee)?;
         if target.is_none() {
-            return Ok(None);
+            return Ok(false);
         }
         let variable = target.unwrap();
-        let variable = variable.borrow();
 
-        let mut result = vec![match self.get_mut(variable.id) {
-            Some(value) => value,
-            None => return Err(StatementError::undefined_variable(variable.name.to_string(), span)),
-        }];
+        self.resolve_assign(
+            cs,
+            &AssignStatement {
+                parent: Cell::new(None),
+                span: assignee.span().cloned(),
+                operation: AssignOperation::Assign,
+                target_variable: Cell::new(variable),
+                target_accesses: accesses,
+                value: Cell::new(assignee),
+            },
+            target_value,
+            indicator,
+        )?;
 
-        for access in accesses {
-            result = Self::resolve_assignee_access(access, span, result)?;
-        }
-        Ok(Some(result))
+        Ok(true)
     }
 }

compiler/src/statement/assign/assign.rs

@@ -33,54 +33,13 @@ impl<'a, F: PrimeField, G: GroupType<F>> ConstrainedProgram<'a, F, G> {
     ) -> Result<(), StatementError> {
         // Get the name of the variable we are assigning to
         let new_value = self.enforce_expression(cs, statement.value.get())?;
-        let mut resolved_assignee = self.resolve_assign(cs, statement)?;
 
-        if resolved_assignee.len() == 1 {
-            let span = statement.span.clone().unwrap_or_default();
-
-            Self::enforce_assign_operation(
-                cs,
-                indicator,
-                format!("select {} {}:{}", new_value, &span.line_start, &span.col_start),
-                &statement.operation,
-                resolved_assignee[0],
-                new_value,
-                &span,
-            )?;
-        } else {
-            match new_value {
-                ConstrainedValue::Array(new_values) => {
-                    let span = statement.span.clone().unwrap_or_default();
-
-                    for (i, (old_ref, new_value)) in
-                        resolved_assignee.into_iter().zip(new_values.into_iter()).enumerate()
-                    {
-                        Self::enforce_assign_operation(
-                            cs,
-                            indicator,
-                            format!(
-                                "select-splice {} {} {}:{}",
-                                i, new_value, &span.line_start, &span.col_start
-                            ),
-                            &statement.operation,
-                            old_ref,
-                            new_value,
-                            &span,
-                        )?;
-                    }
-                }
-                _ => {
-                    return Err(StatementError::array_assign_range(
-                        &statement.span.clone().unwrap_or_default(),
-                    ));
-                }
-            };
-        }
+        self.resolve_assign(cs, statement, new_value, indicator)?;
 
         Ok(())
     }
 
-    fn enforce_assign_operation<CS: ConstraintSystem<F>>(
+    pub(super) fn enforce_assign_operation<CS: ConstraintSystem<F>>(
         cs: &mut CS,
         condition: &Boolean,
         scope: String,