1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107

// Copyright (C) 2019-2021 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/>.

//! Enforces an assign statement in a compiled Leo program.

use crate::{arithmetic::*, errors::StatementError, program::ConstrainedProgram, value::ConstrainedValue, GroupType};
use leo_asg::{AssignOperation, AssignStatement, Span};

use snarkvm_models::{
    curves::{Field, PrimeField},
    gadgets::{
        r1cs::ConstraintSystem,
        utilities::{boolean::Boolean, select::CondSelectGadget},
    },
};

impl<F: Field + PrimeField, G: GroupType<F>> ConstrainedProgram<F, G> {
    #[allow(clippy::too_many_arguments)]
    pub fn enforce_assign_statement<CS: ConstraintSystem<F>>(
        &mut self,
        cs: &mut CS,
        indicator: &Boolean,
        statement: &AssignStatement,
    ) -> Result<(), StatementError> {
        // Get the name of the variable we are assigning to
        let new_value = self.enforce_expression(cs, &statement.value)?;
        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, &span.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, &span.start),
                            &statement.operation,
                            old_ref,
                            new_value,
                            &span,
                        )?;
                    }
                }
                _ => {
                    return Err(StatementError::array_assign_range(
                        statement.span.clone().unwrap_or_default(),
                    ));
                }
            };
        }

        Ok(())
    }

    fn enforce_assign_operation<CS: ConstraintSystem<F>>(
        cs: &mut CS,
        condition: &Boolean,
        scope: String,
        operation: &AssignOperation,
        target: &mut ConstrainedValue<F, G>,
        new_value: ConstrainedValue<F, G>,
        span: &Span,
    ) -> Result<(), StatementError> {
        let new_value = match operation {
            AssignOperation::Assign => new_value,
            AssignOperation::Add => enforce_add(cs, target.clone(), new_value, span)?,
            AssignOperation::Sub => enforce_sub(cs, target.clone(), new_value, span)?,
            AssignOperation::Mul => enforce_mul(cs, target.clone(), new_value, span)?,
            AssignOperation::Div => enforce_div(cs, target.clone(), new_value, span)?,
            AssignOperation::Pow => enforce_pow(cs, target.clone(), new_value, span)?,
        };
        let selected_value = ConstrainedValue::conditionally_select(cs.ns(|| scope), condition, &new_value, target)
            .map_err(|_| StatementError::select_fail(new_value.to_string(), target.to_string(), span.clone()))?;

        *target = selected_value;
        Ok(())
    }
}