update

3 files changed, 55 insertions, 10 deletions

diff --git a/script/vm/compiler.lua b/script/vm/compiler.lua
index 367fb334..7d622f37 100644
--- a/script/vm/compiler.lua
+++ b/script/vm/compiler.lua
@@ -425,6 +425,24 @@ local compilerMap = util.switch()
                 end
             end
         end
+
+        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
+        end
     end)
     : case 'paren'
     : call(function (source)
@@ -467,7 +485,22 @@ local compilerMap = util.switch()
             end
         end
         if source.parent.type == 'funcargs' and not hasMarkDoc and not hasMarkParam then
-            nodeMgr.setNode(source, globalMgr.getGlobal('type', 'any'))
+            local func = source.parent.parent
+            local funcNode = m.compileNode(func)
+            local hasDocArg
+            for n in nodeMgr.eachNode(funcNode) do
+                if n.type == 'doc.type.function' then
+                    for index, arg in ipairs(n.args) do
+                        if func.args[index] == source then
+                            nodeMgr.setNode(source, m.compileNode(arg))
+                            hasDocArg = true
+                        end
+                    end
+                end
+            end
+            if not hasDocArg then
+                nodeMgr.setNode(source, globalMgr.getGlobal('type', 'any'))
+            end
         end
         -- for x in ... do
         if source.parent.type == 'in' then
diff --git a/script/vm/infer.lua b/script/vm/infer.lua
index 564c732c..b78be013 100644
--- a/script/vm/infer.lua
+++ b/script/vm/infer.lua
@@ -85,6 +85,7 @@ local viewNodeMap = util.switch()
     end)
     : case 'doc.type.function'
     : call(function (source, options)
+        options['hasDocFunction'] = true
         local args = {}
         local rets = {}
         local argView = ''
@@ -167,6 +168,9 @@ function m.getViews(source)
     if options['hasNumber'] then
         views['integer'] = nil
     end
+    if options['hasDocFunction'] then
+        views['function'] = nil
+    end
     if options['hasTable'] and not options['hasClass'] then
         views['table'] = true
     end
diff --git a/test/type_inference/init.lua b/test/type_inference/init.lua
index a4638468..223b2e44 100644
--- a/test/type_inference/init.lua
+++ b/test/type_inference/init.lua
@@ -8,17 +8,20 @@ rawset(_G, 'TEST', true)
 
 local function getSource(pos)
     local ast = files.getState('')
-    return guide.eachSourceContain(ast.ast, pos, function (source)
+    local result
+    guide.eachSourceContain(ast.ast, pos, function (source)
         if source.type == 'local'
         or source.type == 'getlocal'
         or source.type == 'setlocal'
         or source.type == 'setglobal'
         or source.type == 'getglobal'
         or source.type == 'field'
-        or source.type == 'method' then
-            return source
+        or source.type == 'method'
+        or source.type == 'function' then
+            result = source
         end
     end)
+    return result
 end
 
 function TEST(wanted)
@@ -1059,13 +1062,18 @@ TEST 'fun():number, boolean' [[
 local <?t?>
 ]]
 
---[[
-l:value
-l:work|&1|&1
-f:|&1|&1
-dfun:|&1
-dn:Class
+
+TEST 'fun(value: Class)' [[
+---@class Class
+
+---@param callback fun(value: Class)
+function work(callback)
+end
+
+work(<?function?> (value)
+end)
 ]]
+
 TEST 'Class' [[
 ---@class Class