cgen: implement argument operator overloading (#8067)

CHANGELOG.md

@@ -11,6 +11,7 @@
 - Treating `enum` as `int` and operations on `enum` except `==` and `!=` are removed for strict type checking.
 - Support `[manualfree] fn f1(){}` and `[manualfree] module m1`, for functions doing their own memory management.
 - Allow usage of `<` and `>` operators for struct in `.sort` method for arrays, i.e. `arr.sort(a < b)`.
+- Auto generate assignment operators like `+=`, `-=`, `*=`, `/=` and `%=` if the operators are defined.
 
 ## V 0.2.1
 *30 Dec 2020*

doc/docs.md

@@ -3216,8 +3216,11 @@ fn (a Vec) - (b Vec) Vec {
 fn main() {
 	a := Vec{2, 3}
 	b := Vec{4, 5}
+	mut c := Vec{1, 2}
 	println(a + b) // "{6, 8}"
 	println(a - b) // "{-2, -2}"
+	c += a
+	println(c) // "{3, 5}"
 }
 ```
 
@@ -3235,6 +3238,8 @@ To improve safety and maintainability, operator overloading is limited:
 - Operator functions can't modify their arguments.
 - When using `<`, `>`, `>=`, `<=`, `==` and `!=` operators, the return type must be `bool`.
 - Both arguments must have the same type (just like with all operators in V).
+- Assignment operators (`*=`, `+=`, `/=`, etc) 
+are auto generated when the operators are defined though they must return the same type.
 
 ## Inline assembly
 

vlib/v/checker/checker.v

@@ -2494,10 +2494,10 @@ pub fn (mut c Checker) assign_stmt(mut assign_stmt ast.AssignStmt) {
 						c.error('invalid right operand: $left_sym.name $assign_stmt.op $right_sym.name',
 							right.position())
 					}
-				} else if !left_sym.is_number() && left_sym.kind !in [.byteptr, .charptr] {
+				} else if !left_sym.is_number() && left_sym.kind !in [.byteptr, .charptr, .struct_] {
 					c.error('operator `$assign_stmt.op` not defined on left operand type `$left_sym.name`',
 						left.position())
-				} else if !right_sym.is_number() && left_sym.kind !in [.byteptr, .charptr] {
+				} else if !right_sym.is_number() && left_sym.kind !in [.byteptr, .charptr, .struct_] {
 					c.error('invalid right operand: $left_sym.name $assign_stmt.op $right_sym.name',
 						right.position())
 				} else if right is ast.IntegerLiteral {
@@ -2513,11 +2513,11 @@ pub fn (mut c Checker) assign_stmt(mut assign_stmt ast.AssignStmt) {
 			}
 			.mult_assign, .div_assign {
 				if !left_sym.is_number() &&
-					!c.table.get_final_type_symbol(left_type_unwrapped).is_int() {
+					!c.table.get_final_type_symbol(left_type_unwrapped).is_int() && left_sym.kind != .struct_ {
 					c.error('operator $assign_stmt.op.str() not defined on left operand type `$left_sym.name`',
 						left.position())
 				} else if !right_sym.is_number() &&
-					!c.table.get_final_type_symbol(left_type_unwrapped).is_int() {
+					!c.table.get_final_type_symbol(left_type_unwrapped).is_int() && left_sym.kind != .struct_ {
 					c.error('operator $assign_stmt.op.str() not defined on right operand type `$right_sym.name`',
 						right.position())
 				}
@@ -2535,6 +2535,33 @@ pub fn (mut c Checker) assign_stmt(mut assign_stmt ast.AssignStmt) {
 			}
 			else {}
 		}
+		if assign_stmt.op in
+			[.plus_assign, .minus_assign, .mod_assign, .mult_assign, .div_assign] &&
+			left_sym.kind == .struct_ && right_sym.kind == .struct_ {
+			extracted_op := match assign_stmt.op {
+				.plus_assign { '+' }
+				.minus_assign { '-' }
+				.div_assign { '/' }
+				.mod_assign { '%' }
+				.mult_assign { '*' }
+				else { 'unknown op' }
+			}
+			left_name := c.table.type_to_str(left_type)
+			if method := left_sym.find_method(extracted_op) {
+				if method.return_type != left_type {
+					c.error('operator `$extracted_op` must return `$left_name` to be used as an assignment operator',
+						assign_stmt.pos)
+				}
+			} else {
+				right_name := c.table.type_to_str(right_type)
+				if left_name == right_name {
+					c.error('operation `$left_name` $extracted_op `$right_name` does not exist, please define it',
+						assign_stmt.pos)
+				} else {
+					c.error('mismatched types `$left_name` and `$right_name`', assign_stmt.pos)
+				}
+			}
+		}
 		if !is_blank_ident && right_sym.kind != .placeholder && left_sym.kind != .interface_ {
 			// Dual sides check (compatibility check)
 			c.check_expected(right_type_unwrapped, left_type_unwrapped) or {

vlib/v/checker/tests/method_op_err.out

@@ -33,23 +33,44 @@ vlib/v/checker/tests/method_op_err.vv:26:1: error: argument cannot be `mut` for
       | ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    27 |     return User{}
    28 | }
-vlib/v/checker/tests/method_op_err.vv:32:24: error: infix expr: cannot use `Foo` (right expression) as `User`
-   30 | fn main() {
-   31 |     println(User{3, 4})
-   32 |     println(User{3, 4} - Foo{3, 3})
+vlib/v/checker/tests/method_op_err.vv:36:24: error: infix expr: cannot use `Foo` (right expression) as `User`
+   34 | fn main() {
+   35 |     println(User{3, 4})
+   36 |     println(User{3, 4} - Foo{3, 3})
       |                        ^
-   33 |     println(User{3, 2} < User{2, 4})
-   34 |     println(User{3, 4} < Foo{3, 4})
-vlib/v/checker/tests/method_op_err.vv:33:24: error: operation `User` < `User` does not exist, please define it
-   31 |     println(User{3, 4})
-   32 |     println(User{3, 4} - Foo{3, 3})
-   33 |     println(User{3, 2} < User{2, 4})
+   37 |     println(User{3, 2} < User{2, 4})
+   38 |     println(User{3, 4} < Foo{3, 4})
+vlib/v/checker/tests/method_op_err.vv:37:24: error: operation `User` < `User` does not exist, please define it
+   35 |     println(User{3, 4})
+   36 |     println(User{3, 4} - Foo{3, 3})
+   37 |     println(User{3, 2} < User{2, 4})
       |                        ^
-   34 |     println(User{3, 4} < Foo{3, 4})
-   35 | }
-vlib/v/checker/tests/method_op_err.vv:34:24: error: mismatched types `User` and `Foo`
-   32 |     println(User{3, 4} - Foo{3, 3})
-   33 |     println(User{3, 2} < User{2, 4})
-   34 |     println(User{3, 4} < Foo{3, 4})
+   38 |     println(User{3, 4} < Foo{3, 4})
+   39 |     mut u := User{3, 4}
+vlib/v/checker/tests/method_op_err.vv:38:24: error: mismatched types `User` and `Foo`
+   36 |     println(User{3, 4} - Foo{3, 3})
+   37 |     println(User{3, 2} < User{2, 4})
+   38 |     println(User{3, 4} < Foo{3, 4})
       |                        ^
-   35 | }
+   39 |     mut u := User{3, 4}
+   40 |     u += 12
+vlib/v/checker/tests/method_op_err.vv:40:10: error: cannot assign to `u`: expected `User`, not `int literal`
+   38 |     println(User{3, 4} < Foo{3, 4})
+   39 |     mut u := User{3, 4}
+   40 |     u += 12
+      |          ~~
+   41 |     u %= User{1, 3}
+   42 |     u += User{2, 3}
+vlib/v/checker/tests/method_op_err.vv:41:5: error: operator %= not defined on left operand type `User`
+   39 |     mut u := User{3, 4}
+   40 |     u += 12
+   41 |     u %= User{1, 3}
+      |     ^
+   42 |     u += User{2, 3} 
+   43 | }
+vlib/v/checker/tests/method_op_err.vv:42:7: error: operator `+` must return `User` to be used as an assignment operator
+   40 |     u += 12
+   41 |     u %= User{1, 3}
+   42 |     u += User{2, 3} 
+      |       ~~
+   43 | }

vlib/v/checker/tests/method_op_err.vv

@@ -27,9 +27,17 @@ fn (u User) / (mut u1 User) User {
     return User{}
 }
 
+fn (u User) + (u1 User) Foo {
+    return Foo{a: u.a + u1.a, b: u.b + u1.b}
+}
+
 fn main() {
     println(User{3, 4})
     println(User{3, 4} - Foo{3, 3})
     println(User{3, 2} < User{2, 4})
     println(User{3, 4} < Foo{3, 4})
+    mut u := User{3, 4}
+    u += 12
+    u %= User{1, 3}
+    u += User{2, 3} 
 }

vlib/v/gen/cgen.v

@@ -1919,6 +1919,7 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 				''
 			}
 			mut str_add := false
+			mut op_overloaded := false
 			if var_type == table.string_type_idx && assign_stmt.op == .plus_assign {
 				if left is ast.IndexExpr {
 					// a[0] += str => `array_set(&a, 0, &(string[]) {string_add(...))})`
@@ -1933,6 +1934,22 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 				g.is_assign_rhs = true
 				str_add = true
 			}
+			// Assignment Operator Overloading
+			if left_sym.kind == .struct_ &&
+				right_sym.kind == .struct_ && assign_stmt.op in
+				[.plus_assign, .minus_assign, .div_assign, .mult_assign, .mod_assign] {
+				g.expr(left)
+				extracted_op := match assign_stmt.op {
+					.plus_assign { '+' }
+					.minus_assign { '-' }
+					.div_assign { '/' }
+					.mod_assign { '%' }
+					.mult_assign { '*' }
+					else { 'unknown op' }
+				}
+				g.write(' = ${styp}_${util.replace_op(extracted_op)}(')
+				op_overloaded = true
+			}
 			if right_sym.kind == .function && is_decl {
 				if is_inside_ternary && is_decl {
 					g.out.write(tabs[g.indent - g.inside_ternary])
@@ -1969,9 +1986,9 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 				if is_decl {
 					g.writeln(';')
 				}
-			} else if !g.is_array_set && !str_add {
+			} else if !g.is_array_set && !str_add && !op_overloaded {
 				g.write(' $op ')
-			} else if str_add {
+			} else if str_add || op_overloaded {
 				g.write(', ')
 			}
 			mut cloned := false
@@ -2049,7 +2066,7 @@ fn (mut g Gen) gen_assign_stmt(assign_stmt ast.AssignStmt) {
 					g.write('.data')
 				}
 			}
-			if str_add {
+			if str_add || op_overloaded {
 				g.write(')')
 			}
 			if g.is_array_set {

vlib/v/tests/operator_overloading_with_string_interpolation_test.v

@@ -86,4 +86,14 @@ fn test_operator_overloading_with_string_interpolation() {
 	assert e.str() == '{8, 15}'
 	assert f.str() == '{0, 0}'
 	assert g.str() == '{2, 3}'
+	///// /// //
+	mut ad := Vec{2, 4}
+	ad += Vec{3, 6}
+	assert ad.str() == '{5, 10}'
+	ad -= Vec{1, 1}
+	assert ad.str() == '{4, 9}'
+	ad *= Vec{2, 2}
+	assert ad.str() == '{8, 18}'
+	ad /= Vec{2, 2}
+	assert ad.str() == '{4, 9}' 
 }

vlib/v/token/token.v

@@ -53,11 +53,11 @@ pub enum Kind {
 	decl_assign // :=
 	plus_assign // +=
 	minus_assign // -=
-	div_assign
-	mult_assign
-	xor_assign
-	mod_assign
-	or_assign
+	div_assign // /=
+	mult_assign // *=
+	xor_assign // ^=
+	mod_assign // %=
+	or_assign // |=
 	and_assign
 	right_shift_assign
 	left_shift_assign // {}  () []