feat: infer function param type when overriding class function or method

Author: tomlau10

script/vm/compiler.lua

@@ -1117,26 +1117,56 @@ local function compileFunctionParam(func, source)
     end
     ---@cast aindex integer
 
-    -- local call ---@type fun(f: fun(x: number));call(function (x) end) --> x -> number
     local funcNode = vm.compileNode(func)
-    local found = false
-    for n in funcNode:eachObject() do
-        if n.type == 'doc.type.function' and n.args[aindex] then
-            local argNode = vm.compileNode(n.args[aindex])
-            for an in argNode:eachObject() do
-                if an.type ~= 'doc.generic.name' then
-                    vm.setNode(source, an)
+    if func.parent.type == 'callargs' then
+        -- local call ---@type fun(f: fun(x: number));call(function (x) end) --> x -> number
+        for n in funcNode:eachObject() do
+            if n.type == 'doc.type.function' and n.args[aindex] then
+                local argNode = vm.compileNode(n.args[aindex])
+                for an in argNode:eachObject() do
+                    if an.type ~= 'doc.generic.name' then
+                        vm.setNode(source, an)
+                    end
                 end
-            end
-            -- NOTE: keep existing behavior for local call which only set type based on the 1st match
-            if func.parent.type == 'callargs' then
+                -- NOTE: keep existing behavior for function as argument which only set type based on the 1st match
                 return true
             end
-            found = true
         end
-    end
-    if found then
-        return true
+    else
+        -- function declaration: use info from all `fun()`, also from the base function when overriding
+        --[[
+            ---@type fun(x: string)|fun(x: number)
+            local function f1(x) end --> x -> string|number
+
+            ---@overload fun(x: string)
+            ---@overload fun(x: number)
+            local function f2(x) end --> x -> string|number
+
+            ---@class A
+            local A = {}
+            ---@param x number
+            function A:f(x) end --> x -> number
+            ---@type A
+            local a = {}
+            function a:f(x) end --> x -> number
+        ]]
+        local found = false
+        for n in funcNode:eachObject() do
+            if (n.type == 'doc.type.function' or n.type == 'function')
+            and n.args[aindex] and n.args[aindex] ~= source
+            then
+                local argNode = vm.compileNode(n.args[aindex])
+                for an in argNode:eachObject() do
+                    if an.type ~= 'doc.generic.name' then
+                        vm.setNode(source, an)
+                    end
+                end
+                found = true
+            end
+        end
+        if found then
+            return true
+        end
     end
 
     local derviationParam = config.get(guide.getUri(func), 'Lua.type.inferParamType')

test/type_inference/common.lua

@@ -4441,3 +4441,15 @@ local B = {}
 
 function B:func(<?x?>) end
 ]]
+
+TEST 'number' [[
+---@class A
+local A = {}
+
+---@param x number
+function A:func(x) end
+
+---@type A
+local a = {}
+function a:func(<?x?>) end
+]]