update

3 files changed, 95 insertions, 32 deletions

diff --git a/script/vm/compiler.lua b/script/vm/compiler.lua
index 4427a66c..c46519fd 100644
--- a/script/vm/compiler.lua
+++ b/script/vm/compiler.lua
@@ -407,6 +407,48 @@ local function selectNode(source, list, index)
     return nodeMgr.setNode(source, result)
 end
 
+local function setCallBackNode(source, call, callNode, fixIndex)
+    local valueMgr = require 'vm.value'
+    local myIndex
+    for i, arg in ipairs(call.args) do
+        if arg == source then
+            myIndex = i - fixIndex
+            break
+        end
+    end
+    local eventIndex = 1
+    local eventArg = call.args[eventIndex + fixIndex]
+    if eventArg and eventArg.dummy then
+        eventIndex = 2
+        eventArg = call.args[eventIndex + fixIndex]
+    end
+    local eventMap = valueMgr.getLiterals(eventArg)
+    for n in nodeMgr.eachNode(callNode) do
+        if n.type == 'function' then
+            local arg = n.args[myIndex]
+            for fn in nodeMgr.eachNode(m.compileNode(arg)) do
+                if fn.type == 'doc.type.function' then
+                    nodeMgr.setNode(source, fn)
+                end
+            end
+        end
+        if n.type == 'doc.type.function' then
+            local event = m.compileNode(n.args[eventIndex])
+            if not event
+            or not eventMap
+            or event.type ~= 'doc.type.enum'
+            or eventMap[event[1]] then
+                local arg = n.args[myIndex]
+                for fn in nodeMgr.eachNode(m.compileNode(arg)) do
+                    if fn.type == 'doc.type.function' then
+                        nodeMgr.setNode(source, fn)
+                    end
+                end
+            end
+        end
+    end
+end
+
 local compilerMap = util.switch()
     : case 'nil'
     : case 'boolean'
@@ -437,36 +479,13 @@ local compilerMap = util.switch()
         if source.parent.type == 'callargs' then
             local call = source.parent.parent
             local callNode = m.compileNode(call.node)
-            for n in nodeMgr.eachNode(callNode) do
-                if n.type == 'function' then
-                    for index, arg in ipairs(n.args) do
-                        if call.args[index] == source then
-                            for fn in nodeMgr.eachNode(m.compileNode(arg)) do
-                                if fn.type == 'doc.type.function' then
-                                    nodeMgr.setNode(source, fn)
-                                end
-                            end
-                        end
-                    end
-                end
-            end
+            setCallBackNode(source, call, callNode, 0)
+
             if call.node.special == 'pcall'
             or call.node.special == 'xpcall' then
                 local fixIndex = call.node.special == 'pcall' and 1 or 2
                 callNode = m.compileNode(call.args[1])
-                for n in nodeMgr.eachNode(callNode) do
-                    if n.type == 'function' then
-                        for index, arg in ipairs(n.args) do
-                            if call.args[index + fixIndex] == source then
-                                for fn in nodeMgr.eachNode(m.compileNode(arg)) do
-                                    if fn.type == 'doc.type.function' then
-                                        nodeMgr.setNode(source, fn)
-                                    end
-                                end
-                            end
-                        end
-                    end
-                end
+                setCallBackNode(source, call, callNode, fixIndex)
             end
         end
     end)
diff --git a/script/vm/value.lua b/script/vm/value.lua
index 43c5d4bd..549cee02 100644
--- a/script/vm/value.lua
+++ b/script/vm/value.lua
@@ -191,4 +191,37 @@ function m.getBoolean(v)
     return result
 end
 
+---@param v vm.node
+---@return table<any, boolean>?
+function m.getLiterals(v)
+    local map
+    local node = compiler.compileNode(v)
+    for n in nodeMgr.eachNode(node) do
+        local literal
+        if n.type == 'boolean'
+        or n.type == 'string'
+        or n.type == 'number'
+        or n.type == 'integer' then
+            literal = n[1]
+        end
+        if n.type == 'global' and n.cate == 'type' then
+            if n.name == 'true' then
+                literal = true
+            elseif n.name == 'false' then
+                literal = false
+            end
+        end
+        if n.type == 'doc.type.integer' then
+            literal = n[1]
+        end
+        if literal ~= nil then
+            if not map then
+                map = {}
+            end
+            map[literal] = true
+        end
+    end
+    return map
+end
+
 return m
diff --git a/test/type_inference/init.lua b/test/type_inference/init.lua
index e253cba5..7d3e3280 100644
--- a/test/type_inference/init.lua
+++ b/test/type_inference/init.lua
@@ -423,7 +423,7 @@ local <?x?>
 ]]
 
 TEST '"enum1"|"enum2"' [[
----@type '"enum1"' | '"enum2"'
+---@type 'enum1' | 'enum2'
 local <?x?>
 ]]
 
@@ -544,7 +544,7 @@ print(t.<?a?>)
 ]]
 
 TEST '"aaa"|"bbb"' [[
----@type table<string, '"aaa"'|'"bbb"'>
+---@type table<string, 'aaa'|'bbb'>
 local t = {}
 
 print(t.<?a?>)
@@ -1158,11 +1158,22 @@ end
 --f = function (<?x?>) end
 --]]
 
+TEST 'fun(i: integer)' [[
+--- @class Emit
+--- @field on fun(eventName: string, cb: function)
+--- @field on fun(eventName: 'died', cb: fun(i: integer))
+--- @field on fun(eventName: 'won', cb: fun(s: string))
+local emit = {}
+
+emit.on("died", <?function?> (i)
+end)
+]]
+
 TEST 'integer' [[
 --- @class Emit
 --- @field on fun(eventName: string, cb: function)
---- @field on fun(eventName: '"died"', cb: fun(i: integer))
---- @field on fun(eventName: '"won"', cb: fun(s: string))
+--- @field on fun(eventName: 'died', cb: fun(i: integer))
+--- @field on fun(eventName: 'won', cb: fun(s: string))
 local emit = {}
 
 emit.on("died", function (<?i?>)
@@ -1172,8 +1183,8 @@ end)
 TEST 'integer' [[
 --- @class Emit
 --- @field on fun(self: Emit, eventName: string, cb: function)
---- @field on fun(self: Emit, eventName: '"died"', cb: fun(i: integer))
---- @field on fun(self: Emit, eventName: '"won"', cb: fun(s: string))
+--- @field on fun(self: Emit, eventName: 'died', cb: fun(i: integer))
+--- @field on fun(self: Emit, eventName: 'won', cb: fun(s: string))
 local emit = {}
 
 emit:on("died", function (<?i?>)