2 files changed, 53 insertions, 0 deletions
diff --git a/script/vm/compiler.lua b/script/vm/compiler.lua
index 390387b4..c7f5ad84 100644
--- a/script/vm/compiler.lua
+++ b/script/vm/compiler.lua
@@ -432,6 +432,7 @@ local compilerMap = util.switch()
                 nodeMgr.setNode(source, m.compileNode(source.value))
             end
         end
+        local hasMarkParam
         -- function x.y(self, ...) --> function x:y(...)
         if  source[1] == 'self'
         and not hasMarkDoc
@@ -439,9 +440,13 @@ local compilerMap = util.switch()
         and source.parent[1] == source then
             local setfield = source.parent.parent.parent
             if setfield.type == 'setfield' then
+                hasMarkParam = true
                 nodeMgr.setNode(source, m.compileNode(setfield.node))
             end
         end
+        if source.parent.type == 'funcargs' and not hasMarkDoc and not hasMarkParam then
+            nodeMgr.setNode(source, globalMgr.getGlobal('type', 'any'))
+        end
         -- for x in ... do
         if source.parent.type == 'in' then
             m.compileNode(source.parent)
@@ -601,6 +606,14 @@ local compilerMap = util.switch()
             nodeMgr.setNode(source, m.compileNode(type))
         end
     end)
+    : case 'doc.type.arg'
+    : call(function (source)
+        if source.extends then
+            nodeMgr.setNode(source, m.compileNode(source.extends))
+        else
+            nodeMgr.setNode(source, globalMgr.getGlobal('type', 'any'))
+        end
+    end)
     : case 'unary'
     : call(function (source)
         local valueMgr = require 'vm.value'
diff --git a/script/vm/infer.lua b/script/vm/infer.lua
index 80558eaf..bcbd2935 100644
--- a/script/vm/infer.lua
+++ b/script/vm/infer.lua
@@ -28,6 +28,14 @@ local viewNodeMap = util.switch()
     : call(function (source, options)
         options['hasTable'] = true
     end)
+    : case 'local'
+    : call(function (source, options)
+        if source.parent == 'funcargs' then
+            options['isParam'] = true
+        else
+            options['isLocal'] = true
+        end
+    end)
     : case 'global'
     : call(function (source, options)
         if source.cate == 'type' then
@@ -52,6 +60,33 @@ local viewNodeMap = util.switch()
     : call(function (source, options)
         return ('%q'):format(source[1])
     end)
+    : case 'doc.type.function'
+    : call(function (source, options)
+        local args = {}
+        local rets = {}
+        local argView = ''
+        local regView = ''
+        for i, arg in ipairs(source.args) do
+            args[i] = string.format('%s%s: %s'
+                , arg.name[1]
+                , arg.optional and '?' or ''
+                , m.viewType(arg)
+            )
+        end
+        if #args > 0 then
+            argView = table.concat(args, ', ')
+        end
+        for i, ret in ipairs(source.returns) do
+            rets[i] = string.format('%s%s'
+                , m.viewType(ret)
+                , ret.optional and '?' or ''
+            )
+        end
+        if #rets > 0 then
+            regView = ':' .. table.concat(rets, ', ')
+        end
+        return ('fun(%s)%s'):format(argView, regView)
+    end)
     : getMap()
 
 ---@param node vm.node
@@ -119,6 +154,11 @@ end
 ---@return string
 function m.viewType(source)
     local views = m.getViews(source)
+
+    if views['any'] then
+        return 'any'
+    end
+
     if not next(views) then
         return 'unknown'
     end