From 808bae1bcf4624b63a427adf71f41e22684b66c6 Mon Sep 17 00:00:00 2001
From: Lukas Neubert <lukas.neubert@proton.me>
Date: Wed, 3 Jan 2024 23:31:43 +0100
Subject: [PATCH] fix: integer division assign possibly resulting in float
 (#163)

---
 CHANGELOG.md                                  |  3 +++
 lib/bait/parser/stmt.bt                       | 15 ++++++++++++--
 lib/bait/token/token.bt                       | 20 ++++++++++++++++---
 .../number_test.bt => number/float_test.bt}   |  0
 tests/number/int_division_test.bt             | 15 ++++++++++++++
 5 files changed, 48 insertions(+), 5 deletions(-)
 rename tests/{other/number_test.bt => number/float_test.bt} (100%)
 create mode 100644 tests/number/int_division_test.bt

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 63b711ca..05c6434d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,9 @@ _unreleased_
 ### Type Checks
 - Ensure function parameter type exists
 
+### General Fixes
+- Fix integer division assign possibly resulting in a float
+
 
 ## 0.0.6
 _26 December 2023_
diff --git a/lib/bait/parser/stmt.bt b/lib/bait/parser/stmt.bt
index c8b95bdf..871adb43 100644
--- a/lib/bait/parser/stmt.bt
+++ b/lib/bait/parser/stmt.bt
@@ -148,9 +148,20 @@ fun (mut p Parser) assign_stmt() ast.AssignStmt {
 }
 
 fun (mut p Parser) partial_assign_stmt(left ast.Expr) ast.AssignStmt {
-	op := p.tok
+	mut op := p.tok
 	p.next()
-	right := p.expr(.lowest)
+	mut right := p.expr(.lowest)
+
+	// Convert math assign to normal assign
+	if op.is_math_assign() {
+		right = ast.InfixExpr{
+			op = op.math_from_assign()
+			left = left
+			right = right
+		}
+		op = .assign
+	}
+
 	return ast.AssignStmt{
 		op = op
 		left = left
diff --git a/lib/bait/token/token.bt b/lib/bait/token/token.bt
index cfe91fad..00a859c9 100644
--- a/lib/bait/token/token.bt
+++ b/lib/bait/token/token.bt
@@ -110,6 +110,7 @@ pub fun kind_from_string(name string) Token {
 const COMPARE_KINDS := [Token.eq, .ne, .lt, .gt, .le, .ge, .key_is]
 const MATH_KINDS := [Token.plus, .minus, .mul, .div, .mod]
 const OTHER_INFIXES := [Token.key_and, .key_or]
+const MATH_ASSIGNS := [Token.plus_assign, .minus_assign, .mul_assign, .div_assign, .mod_assign]
 
 pub fun (kind Token) is_compare() bool {
 	return COMPARE_KINDS.contains(kind)
@@ -119,10 +120,23 @@ pub fun (kind Token) is_infix() bool {
 	return kind.is_compare() or MATH_KINDS.contains(kind) or OTHER_INFIXES.contains(kind)
 }
 
-const ASSIGN_KINDS := [Token.decl_assign, .assign, .plus_assign, .minus_assign, .mul_assign, .div_assign, .mod_assign]
-
 pub fun (kind Token) is_assign() bool {
-	return ASSIGN_KINDS.contains(kind)
+	return kind == .decl_assign or kind == .assign or kind.is_math_assign()
+}
+
+pub fun (tok Token) is_math_assign() bool {
+	return MATH_ASSIGNS.contains(tok)
+}
+
+pub fun (tok Token) math_from_assign() Token {
+	return match tok {
+		.plus_assign { .plus }
+		.minus_assign { .minus }
+		.mul_assign { .mul }
+		.div_assign { .div }
+		.mod_assign { .mod }
+		else { panic('invalid math assign token') }
+	}
 }
 
 pub fun (kind Token) js_repr() string {
diff --git a/tests/other/number_test.bt b/tests/number/float_test.bt
similarity index 100%
rename from tests/other/number_test.bt
rename to tests/number/float_test.bt
diff --git a/tests/number/int_division_test.bt b/tests/number/int_division_test.bt
new file mode 100644
index 00000000..1ab9dd45
--- /dev/null
+++ b/tests/number/int_division_test.bt
@@ -0,0 +1,15 @@
+// SPDX-FileCopyrightText: 2024-present Lukas Neubert <lukas.neubert@proton.me>
+// SPDX-License-Identifier: MPL-2.0
+
+fun test_integer_division() {
+	i := 10
+	res := i / 3
+	assert res == 3
+
+	mut n := 2
+	n /= 2
+	assert n == 1
+
+	n /= 2
+	assert n == 0
+}