From 9601b43f5b92bf893159474507b6e71a4f03a700 Mon Sep 17 00:00:00 2001
From: Swastik <swstkbaranwal@gmail.com>
Date: Tue, 11 Jun 2024 01:05:32 +0530
Subject: [PATCH] checker: fix return type of alias primitive operator
 overloading

---
 vlib/v/checker/infix.v                        | 20 +++++++++++-
 vlib/v/checker/tests/method_op_alias_err.out  |  2 +-
 ...lias_primitive_operator_overloading_test.v | 31 +++++++++++++++++++
 .../alias_voidptr_operator_overloading_test.v | 30 ++++++++++++++++++
 4 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 vlib/v/tests/alias_primitive_operator_overloading_test.v
 create mode 100644 vlib/v/tests/alias_voidptr_operator_overloading_test.v

diff --git a/vlib/v/checker/infix.v b/vlib/v/checker/infix.v
index 0a2c7635b2a53e..008f65a61a253e 100644
--- a/vlib/v/checker/infix.v
+++ b/vlib/v/checker/infix.v
@@ -435,7 +435,25 @@ fn (mut c Checker) infix_expr(mut node ast.InfixExpr) ast.Type {
 					c.check_div_mod_by_zero(node.right, node.op)
 				}
 
-				return_type = promoted_type
+				left_sym = c.table.sym(unwrapped_left_type)
+				right_sym = c.table.sym(unwrapped_right_type)
+				if left_sym.info is ast.Alias
+					&& c.table.sym(left_sym.info.parent_type).is_primitive() {
+					if left_sym.has_method(node.op.str()) {
+						if method := left_sym.find_method(node.op.str()) {
+							return_type = method.return_type
+						}
+					}
+				} else if right_sym.info is ast.Alias
+					&& c.table.sym(right_sym.info.parent_type).is_primitive() {
+					if right_sym.has_method(node.op.str()) {
+						if method := left_sym.find_method(node.op.str()) {
+							return_type = method.return_type
+						}
+					}
+				} else {
+					return_type = promoted_type
+				}
 			}
 		}
 		.gt, .lt, .ge, .le {
diff --git a/vlib/v/checker/tests/method_op_alias_err.out b/vlib/v/checker/tests/method_op_alias_err.out
index b5e72b30424052..c92ff9c3dbaaeb 100644
--- a/vlib/v/checker/tests/method_op_alias_err.out
+++ b/vlib/v/checker/tests/method_op_alias_err.out
@@ -5,7 +5,7 @@ vlib/v/checker/tests/method_op_alias_err.vv:4:18: error: expected `Foo` not `Foo
       |                  ~~~~
     5 |     return Foo2(f + f1)
     6 | }
-vlib/v/checker/tests/method_op_alias_err.vv:5:17: error: infix expr: cannot use `string` (right expression) as `string`
+vlib/v/checker/tests/method_op_alias_err.vv:5:17: error: infix expr: cannot use `Foo2` (right expression) as `Foo`
     3 | 
     4 | fn (f Foo) + (f1 Foo2) Foo2 {
     5 |     return Foo2(f + f1)
diff --git a/vlib/v/tests/alias_primitive_operator_overloading_test.v b/vlib/v/tests/alias_primitive_operator_overloading_test.v
new file mode 100644
index 00000000000000..e4a755feea7833
--- /dev/null
+++ b/vlib/v/tests/alias_primitive_operator_overloading_test.v
@@ -0,0 +1,31 @@
+type Alias = u8
+
+fn new_alias() Alias {
+	return 0
+}
+
+fn (a Alias) add(b Alias) Alias {
+	return new_alias()
+}
+
+fn (a Alias) mul(b Alias) Alias {
+	return new_alias()
+}
+
+fn (a Alias) + (b Alias) Alias {
+	return a.add(b)
+}
+
+fn (a Alias) * (b Alias) Alias {
+	return a.mul(b)
+}
+
+fn test_alias_primitive_operator_overloading() {
+	a := new_alias()
+	b := new_alias()
+
+	c := a + b
+	d := a.add(b)
+	assert typeof(c).name == 'Alias'
+	assert typeof(c).name == 'Alias'
+}
diff --git a/vlib/v/tests/alias_voidptr_operator_overloading_test.v b/vlib/v/tests/alias_voidptr_operator_overloading_test.v
new file mode 100644
index 00000000000000..39caa380ad9580
--- /dev/null
+++ b/vlib/v/tests/alias_voidptr_operator_overloading_test.v
@@ -0,0 +1,30 @@
+type AF_ARRAY = voidptr
+
+fn (a AF_ARRAY) add(b AF_ARRAY) AF_ARRAY {
+	mut y := AF_ARRAY(0)
+	return y
+}
+
+fn (a AF_ARRAY) mul(b AF_ARRAY) AF_ARRAY {
+	mut y := AF_ARRAY(0)
+	return y
+}
+
+fn (a AF_ARRAY) + (b AF_ARRAY) AF_ARRAY {
+	return a.add(b)
+}
+
+fn (a AF_ARRAY) * (b AF_ARRAY) AF_ARRAY {
+	return a.mul(b)
+}
+
+fn test_alias_voidptr_operator_overloading() {
+	a := AF_ARRAY(0)
+	b := AF_ARRAY(0)
+
+	c := a + b
+	y := a * a
+
+	assert c == a.add(b)
+	assert y == a.mul(a)
+}