From de77114593b93334f24e9ffbe3d78f82431d6035 Mon Sep 17 00:00:00 2001
From: Delyan Angelov <delian66@gmail.com>
Date: Thu, 7 Jul 2022 15:13:22 +0300
Subject: [PATCH] eval: support assignment operators like +=, make `for a in
 0..10 {` more robust

---
 vlib/v/eval/expr.v                            | 14 ++++++++++++
 vlib/v/eval/stmt.v                            | 22 +++++++++++++++++--
 vlib/v/eval/testdata/range_and_assign_ops.out |  1 +
 vlib/v/eval/testdata/range_and_assign_ops.vv  |  5 +++++
 vlib/v/token/token.v                          | 17 ++++++++++++++
 5 files changed, 57 insertions(+), 2 deletions(-)
 create mode 100644 vlib/v/eval/testdata/range_and_assign_ops.out
 create mode 100644 vlib/v/eval/testdata/range_and_assign_ops.vv

diff --git a/vlib/v/eval/expr.v b/vlib/v/eval/expr.v
index 9c60ad46c4..a91b394c03 100644
--- a/vlib/v/eval/expr.v
+++ b/vlib/v/eval/expr.v
@@ -6,6 +6,20 @@ import v.util
 import math
 import strconv
 
+fn (o Object) as_i64() !i64 {
+	match o {
+		i64 {
+			return o
+		}
+		Int {
+			return o.val
+		}
+		else {
+			return error('can not cast object to i64')
+		}
+	}
+}
+
 pub fn (mut e Eval) expr(expr ast.Expr, expecting ast.Type) Object {
 	match expr {
 		ast.CallExpr {
diff --git a/vlib/v/eval/stmt.v b/vlib/v/eval/stmt.v
index e0781abfb5..fa27622a49 100644
--- a/vlib/v/eval/stmt.v
+++ b/vlib/v/eval/stmt.v
@@ -1,6 +1,7 @@
 module eval
 
 import v.ast
+import v.token
 
 pub fn (mut e Eval) stmts(stmts []ast.Stmt) {
 	e.open_scope()
@@ -40,6 +41,16 @@ pub fn (mut e Eval) stmt(stmt ast.Stmt) {
 						e.set(left, rights[i], false, stmt.left_types[i])
 					}
 				}
+				.plus_assign, .minus_assign, .mult_assign, .div_assign, .xor_assign, .mod_assign,
+				.or_assign, .and_assign, .right_shift_assign, .unsigned_right_shift_assign,
+				.left_shift_assign {
+					infix_op := token.assign_op_to_infix_op(stmt.op)
+					for i, left in stmt.left {
+						res := e.infix_expr(e.expr(left, stmt.left_types[i]), rights[i],
+							infix_op, stmt.left_types[i])
+						e.set(left, res, false, stmt.left_types[i])
+					}
+				}
 				else {
 					e.error('unknown assign statment: $stmt.op')
 				}
@@ -64,8 +75,15 @@ pub fn (mut e Eval) stmt(stmt ast.Stmt) {
 			if !underscore {
 				e.set(ast.Ident{ name: stmt.val_var, scope: 0 }, Int{-1, 32}, true, stmt.val_type)
 			}
-			for i in (e.expr(stmt.cond, ast.int_type_idx) as Int).val .. (e.expr(stmt.high,
-				ast.int_type_idx) as Int).val {
+			fstart := e.expr(stmt.cond, ast.int_type_idx).as_i64() or {
+				e.error('invalid integer for start of range')
+				0
+			}
+			fend := e.expr(stmt.high, ast.int_type_idx).as_i64() or {
+				e.error('invalid integer for end of range')
+				0
+			}
+			for i in fstart .. fend {
 				if !underscore {
 					e.set(ast.Ident{ name: stmt.val_var, scope: 0 }, Int{i, 32}, false,
 						stmt.val_type)
diff --git a/vlib/v/eval/testdata/range_and_assign_ops.out b/vlib/v/eval/testdata/range_and_assign_ops.out
new file mode 100644
index 0000000000..ea90ee3198
--- /dev/null
+++ b/vlib/v/eval/testdata/range_and_assign_ops.out
@@ -0,0 +1 @@
+45
diff --git a/vlib/v/eval/testdata/range_and_assign_ops.vv b/vlib/v/eval/testdata/range_and_assign_ops.vv
new file mode 100644
index 0000000000..0fe40d9645
--- /dev/null
+++ b/vlib/v/eval/testdata/range_and_assign_ops.vv
@@ -0,0 +1,5 @@
+mut sum := 0
+for a in 0..10 {
+	sum += a
+}
+println(sum)
diff --git a/vlib/v/token/token.v b/vlib/v/token/token.v
index ad7045482b..625b128e97 100644
--- a/vlib/v/token/token.v
+++ b/vlib/v/token/token.v
@@ -743,3 +743,20 @@ pub fn kind_from_string(s string) ?Kind {
 		else { error('unknown') }
 	}
 }
+
+pub fn assign_op_to_infix_op(op Kind) Kind {
+	return match op {
+		.plus_assign { .plus }
+		.minus_assign { .minus }
+		.mult_assign { .mul }
+		.div_assign { .div }
+		.xor_assign { .xor }
+		.mod_assign { .mod }
+		.or_assign { .pipe }
+		.and_assign { .amp }
+		.right_shift_assign { .right_shift }
+		.unsigned_right_shift_assign { .unsigned_right_shift }
+		.left_shift_assign { .left_shift }
+		else { ._end_ }
+	}
+}