binary operator

2 files changed, 144 insertions, 4 deletions

diff --git a/script/vm/operator.lua b/script/vm/operator.lua
index 67d039e4..e71911f9 100644
--- a/script/vm/operator.lua
+++ b/script/vm/operator.lua
@@ -25,6 +25,22 @@ vm.BINARY_OP = {
     'concat',
 }
 
+local binaryMap = {
+    ['+']  = 'add',
+    ['-']  = 'sub',
+    ['*']  = 'mul',
+    ['/']  = 'div',
+    ['%']  = 'mod',
+    ['^']  = 'pow',
+    ['//'] = 'idiv',
+    ['&']  = 'band',
+    ['|']  = 'bor',
+    ['~']  = 'bxor',
+    ['<<'] = 'shl',
+    ['>>'] = 'shr',
+    ['..'] = 'concat',
+}
+
 ---@param operators parser.object[]
 ---@param op string
 ---@param value? parser.object
@@ -184,8 +200,8 @@ vm.binarySwitch = util.switch()
     : call(function (source)
         local a = vm.getInteger(source[1])
         local b = vm.getInteger(source[2])
+        local op = source.op.type
         if a and b then
-            local op = source.op.type
             local result = op.type == '<<' and a << b
                         or op.type == '>>' and a >> b
                         or op.type == '&'  and a &  b
@@ -199,7 +215,8 @@ vm.binarySwitch = util.switch()
                 [1]    = result,
             })
         else
-            vm.setNode(source, vm.declareGlobal('type', 'integer'))
+            local node = vm.runOperator(binaryMap[op], source[1], source[2])
+            vm.setNode(source, node or vm.declareGlobal('type', 'integer'))
         end
     end)
     : case '+'
@@ -234,6 +251,11 @@ vm.binarySwitch = util.switch()
                 [1]    = result,
             })
         else
+            local node = vm.runOperator(binaryMap[op], source[1], source[2])
+            if node then
+                vm.setNode(source, node)
+                return
+            end
             if op == '+'
             or op == '-'
             or op == '*'
@@ -251,7 +273,7 @@ vm.binarySwitch = util.switch()
                     end
                 end
             end
-            vm.setNode(source, vm.declareGlobal('type', 'number'))
+            vm.setNode(source, node or vm.declareGlobal('type', 'number'))
         end
     end)
     : case '..'
@@ -287,7 +309,8 @@ vm.binarySwitch = util.switch()
                 [1]    = a .. b,
             })
         else
-            vm.setNode(source, vm.declareGlobal('type', 'string'))
+            local node = vm.runOperator(binaryMap[source.op.type], source[1], source[2])
+            vm.setNode(source, node or vm.declareGlobal('type', 'string'))
         end
     end)
     : case '>'
diff --git a/test/type_inference/init.lua b/test/type_inference/init.lua
index f31a8777..7bb39f19 100644
--- a/test/type_inference/init.lua
+++ b/test/type_inference/init.lua
@@ -3418,3 +3418,120 @@ TEST 'A' [[
 local a
 local <?b?> = #a
 ]]
+
+TEST 'A' [[
+---@class A
+---@operator add: A
+
+---@type A
+local a
+local <?b?> = a + 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator sub: A
+
+---@type A
+local a
+local <?b?> = a - 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator mul: A
+
+---@type A
+local a
+local <?b?> = a * 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator div: A
+
+---@type A
+local a
+local <?b?> = a / 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator mod: A
+
+---@type A
+local a
+local <?b?> = a % 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator pow: A
+
+---@type A
+local a
+local <?b?> = a ^ 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator idiv: A
+
+---@type A
+local a
+local <?b?> = a // 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator band: A
+
+---@type A
+local a
+local <?b?> = a & 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator bor: A
+
+---@type A
+local a
+local <?b?> = a | 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator bxor: A
+
+---@type A
+local a
+local <?b?> = a ~ 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator shl: A
+
+---@type A
+local a
+local <?b?> = a << 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator shr: A
+
+---@type A
+local a
+local <?b?> = a >> 1
+]]
+
+TEST 'A' [[
+---@class A
+---@operator concat: A
+
+---@type A
+local a
+local <?b?> = a .. 1
+]]