From f74d947c6b850e377a0ee55b6f79f3c171beaa6d Mon Sep 17 00:00:00 2001
From: Stephen Chung <Stephen.Chung@intexact.com>
Date: Tue, 3 Nov 2020 13:08:19 +0800
Subject: [PATCH] Fix constant assignment.

---
 src/optimize.rs    | 25 ++++++++++++++++++++-----
 src/parser.rs      |  4 ++--
 src/result.rs      |  6 ++++--
 tests/constants.rs | 19 +++++++++++++++++--
 4 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/src/optimize.rs b/src/optimize.rs
index 3895fd52..6d9bc23d 100644
--- a/src/optimize.rs
+++ b/src/optimize.rs
@@ -156,11 +156,16 @@ fn call_fn_with_constant_arguments(
 /// Optimize a statement.
 fn optimize_stmt(stmt: Stmt, state: &mut State, preserve_result: bool) -> Stmt {
     match stmt {
-        // id op= expr
-        Stmt::Assignment(x, pos) => Stmt::Assignment(
-            Box::new((optimize_expr(x.0, state), x.1, optimize_expr(x.2, state))),
-            pos,
-        ),
+        // expr op= expr
+        Stmt::Assignment(x, pos) => match x.0 {
+            Expr::Variable(_) => {
+                Stmt::Assignment(Box::new((x.0, x.1, optimize_expr(x.2, state))), pos)
+            }
+            _ => Stmt::Assignment(
+                Box::new((optimize_expr(x.0, state), x.1, optimize_expr(x.2, state))),
+                pos,
+            ),
+        },
         // if false { if_block } -> Noop
         Stmt::IfThenElse(Expr::False(pos), x, _) if x.1.is_none() => {
             state.set_dirty();
@@ -462,6 +467,11 @@ fn optimize_expr(expr: Expr, state: &mut State) -> Expr {
                     .map(|(_, mut expr)| { expr.set_position(pos); expr })
                     .unwrap_or_else(|| Expr::Unit(pos))
             }
+            // var.rhs
+            (lhs @ Expr::Variable(_), rhs) => Expr::Dot(Box::new(BinaryExpr {
+                lhs,
+                rhs: optimize_expr(rhs, state),
+            }), dot_pos),
             // lhs.rhs
             (lhs, rhs) => Expr::Dot(Box::new(BinaryExpr {
                 lhs: optimize_expr(lhs, state),
@@ -498,6 +508,11 @@ fn optimize_expr(expr: Expr, state: &mut State) -> Expr {
                 state.set_dirty();
                 Expr::CharConstant(s.name.chars().nth(i as usize).unwrap(), s.pos)
             }
+            // var[rhs]
+            (lhs @ Expr::Variable(_), rhs) => Expr::Index(Box::new(BinaryExpr {
+                lhs,
+                rhs: optimize_expr(rhs, state),
+            }), idx_pos),
             // lhs[rhs]
             (lhs, rhs) => Expr::Index(Box::new(BinaryExpr {
                 lhs: optimize_expr(lhs, state),
diff --git a/src/parser.rs b/src/parser.rs
index df7a28ad..bd6d613c 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -1087,7 +1087,7 @@ fn make_assignment_stmt<'a>(
 ) -> Result<Stmt, ParseError> {
     match &lhs {
         // var (non-indexed) = rhs
-        Expr::Variable(x) if x.1.is_none() => {
+        Expr::Variable(x) if x.0.is_none() => {
             Ok(Stmt::Assignment(Box::new((lhs, fn_name.into(), rhs)), pos))
         }
         // var (indexed) = rhs
@@ -1114,7 +1114,7 @@ fn make_assignment_stmt<'a>(
         // xxx[???] = rhs, xxx.??? = rhs
         Expr::Index(x, _) | Expr::Dot(x, _) => match &x.lhs {
             // var[???] (non-indexed) = rhs, var.??? (non-indexed) = rhs
-            Expr::Variable(x) if x.1.is_none() => {
+            Expr::Variable(x) if x.0.is_none() => {
                 Ok(Stmt::Assignment(Box::new((lhs, fn_name.into(), rhs)), pos))
             }
             // var[???] (indexed) = rhs, var.??? (indexed) = rhs
diff --git a/src/result.rs b/src/result.rs
index 79fb519b..8153a2c5 100644
--- a/src/result.rs
+++ b/src/result.rs
@@ -127,7 +127,7 @@ impl EvalAltResult {
             Self::ErrorVariableNotFound(_, _) => "Variable not found",
             Self::ErrorModuleNotFound(_, _) => "Module not found",
             Self::ErrorDataRace(_, _) => "Data race detected when accessing variable",
-            Self::ErrorAssignmentToConstant(_, _) => "Assignment to a constant variable",
+            Self::ErrorAssignmentToConstant(_, _) => "Cannot assign to a constant",
             Self::ErrorMismatchOutputType(_, _, _) => "Output type is incorrect",
             Self::ErrorInExpr(_) => "Malformed 'in' expression",
             Self::ErrorDotExpr(_, _) => "Malformed dot expression",
@@ -194,7 +194,9 @@ impl fmt::Display for EvalAltResult {
             Self::ErrorRuntime(d, _) if d.is::<()>() => f.write_str(desc)?,
             Self::ErrorRuntime(d, _) => write!(f, "{}: {}", desc, d)?,
 
-            Self::ErrorAssignmentToConstant(s, _) => write!(f, "{}: '{}'", desc, s)?,
+            Self::ErrorAssignmentToConstant(s, _) => {
+                write!(f, "Cannot assign to constant '{}'", s)?
+            }
             Self::ErrorMismatchOutputType(r, s, _) => {
                 write!(f, "Output type is incorrect: {} (expecting {})", r, s)?
             }
diff --git a/tests/constants.rs b/tests/constants.rs
index 4569cda0..2f516a1b 100644
--- a/tests/constants.rs
+++ b/tests/constants.rs
@@ -1,4 +1,4 @@
-use rhai::{Engine, EvalAltResult, ParseErrorType, INT};
+use rhai::{Engine, EvalAltResult, ParseErrorType, Scope, INT};
 
 #[test]
 fn test_constant() -> Result<(), Box<EvalAltResult>> {
@@ -15,13 +15,28 @@ fn test_constant() -> Result<(), Box<EvalAltResult>> {
 
     #[cfg(not(feature = "no_index"))]
     assert!(matches!(
-        *engine.eval::<INT>("const x = [1, 2, 3, 4, 5]; x[2] = 42;").expect_err("expects error"),
+        *engine.consume("const x = [1, 2, 3, 4, 5]; x[2] = 42;").expect_err("expects error"),
         EvalAltResult::ErrorParsing(ParseErrorType::AssignmentToConstant(x), _) if x == "x"
     ));
 
     Ok(())
 }
 
+#[test]
+fn test_constant_scope() -> Result<(), Box<EvalAltResult>> {
+    let engine = Engine::new();
+
+    let mut scope = Scope::new();
+    scope.push_constant("x", 42 as INT);
+
+    assert!(matches!(
+        *engine.consume_with_scope(&mut scope, "x = 1").expect_err("expects error"),
+        EvalAltResult::ErrorAssignmentToConstant(x, _) if x == "x"
+    ));
+
+    Ok(())
+}
+
 #[test]
 fn test_var_is_def() -> Result<(), Box<EvalAltResult>> {
     let engine = Engine::new();