From f6e9e57760b003c2431d6f2ea33493b2f41c1888 Mon Sep 17 00:00:00 2001
From: Maksim Kurnikov <maxim.kurnikov@gmail.com>
Date: Mon, 17 Feb 2025 21:12:42 +0100
Subject: [PATCH] compound assignment expr should return unit type

---
 .../core/types/infer/TypeInferenceWalker.kt   | 25 ++++++++++---------
 .../move/lang/types/ExpressionTypesTest.kt    | 10 ++++++++
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/src/main/kotlin/org/move/lang/core/types/infer/TypeInferenceWalker.kt b/src/main/kotlin/org/move/lang/core/types/infer/TypeInferenceWalker.kt
index 064f2c571..af47bcffc 100644
--- a/src/main/kotlin/org/move/lang/core/types/infer/TypeInferenceWalker.kt
+++ b/src/main/kotlin/org/move/lang/core/types/infer/TypeInferenceWalker.kt
@@ -898,21 +898,22 @@ class TypeInferenceWalker(
     private fun inferBinaryExprTy(binaryExpr: MvBinaryExpr): Ty {
         return when (binaryExpr.binaryOp.op) {
             "<", ">", "<=", ">=" -> inferOrderingBinaryExprTy(binaryExpr)
-            "+", "-", "*", "/", "%" -> inferArithmeticBinaryExprTy(binaryExpr)
             "==", "!=" -> inferEqualityBinaryExprTy(binaryExpr)
             "||", "&&", "==>", "<==>" -> inferLogicBinaryExprTy(binaryExpr)
-            "^", "|", "&" -> inferBitOpsExprTy(binaryExpr)
-            "<<", ">>" -> inferBitShiftsExprTy(binaryExpr)
 
-            "+=", "-=", "*=", "/=", "%=" -> inferArithmeticBinaryExprTy(binaryExpr)
-            "|=", "^=", "&=" -> inferBitOpsExprTy(binaryExpr)
-            ">>=", "<<=" -> inferBitShiftsExprTy(binaryExpr)
+            "+", "-", "*", "/", "%" -> inferArithmeticBinaryExprTy(binaryExpr, false)
+            "^", "|", "&" -> inferBitOpsExprTy(binaryExpr, false)
+            "<<", ">>" -> inferBitShiftsExprTy(binaryExpr, false)
+
+            "+=", "-=", "*=", "/=", "%=" -> inferArithmeticBinaryExprTy(binaryExpr, true)
+            "|=", "^=", "&=" -> inferBitOpsExprTy(binaryExpr, true)
+            ">>=", "<<=" -> inferBitShiftsExprTy(binaryExpr, true)
 
             else -> TyUnknown
         }
     }
 
-    private fun inferArithmeticBinaryExprTy(binaryExpr: MvBinaryExpr): Ty {
+    private fun inferArithmeticBinaryExprTy(binaryExpr: MvBinaryExpr, compoundAssigment: Boolean): Ty {
         val leftExpr = binaryExpr.left
         val rightExpr = binaryExpr.right
         val op = binaryExpr.binaryOp.op
@@ -943,7 +944,7 @@ class TypeInferenceWalker(
                 typeErrorEncountered = true
             }
         }
-        return if (typeErrorEncountered) TyUnknown else leftTy
+        return if (typeErrorEncountered) TyUnknown else (if (compoundAssigment) TyUnit else leftTy)
     }
 
     private fun inferEqualityBinaryExprTy(binaryExpr: MvBinaryExpr): Ty {
@@ -1009,7 +1010,7 @@ class TypeInferenceWalker(
         return TyBool
     }
 
-    private fun inferBitOpsExprTy(binaryExpr: MvBinaryExpr): Ty {
+    private fun inferBitOpsExprTy(binaryExpr: MvBinaryExpr, compoundAssigment: Boolean): Ty {
         val leftExpr = binaryExpr.left
         val rightExpr = binaryExpr.right
 
@@ -1017,10 +1018,10 @@ class TypeInferenceWalker(
         if (rightExpr != null) {
             rightExpr.inferTypeCoercableTo(leftTy)
         }
-        return leftTy
+        return if (compoundAssigment) TyUnit else leftTy
     }
 
-    private fun inferBitShiftsExprTy(binaryExpr: MvBinaryExpr): Ty {
+    private fun inferBitShiftsExprTy(binaryExpr: MvBinaryExpr, compoundAssigment: Boolean): Ty {
         val leftExpr = binaryExpr.left
         val rightExpr = binaryExpr.right
 
@@ -1028,7 +1029,7 @@ class TypeInferenceWalker(
         if (rightExpr != null) {
             rightExpr.inferTypeCoercableTo(TyInteger.U8)
         }
-        return leftTy
+        return if (compoundAssigment) TyUnit else leftTy
     }
 
     private fun Ty.supportsArithmeticOp(): Boolean {
diff --git a/src/test/kotlin/org/move/lang/types/ExpressionTypesTest.kt b/src/test/kotlin/org/move/lang/types/ExpressionTypesTest.kt
index 39bebf3ec..9ea97cf85 100644
--- a/src/test/kotlin/org/move/lang/types/ExpressionTypesTest.kt
+++ b/src/test/kotlin/org/move/lang/types/ExpressionTypesTest.kt
@@ -2430,4 +2430,14 @@ module 0x1::main {
             }
         }        
     """)
+
+    fun `test compound assignment returns unit type`() = testExpr("""
+        module 0x1::m {
+            fun main() {
+                let a = 1;
+                (a += 1)
+              //^ ()  
+            }
+        }        
+    """)
 }