resolve #871 infer called function by params num

Author: sumneko

changelog.md

@@ -6,6 +6,17 @@
   local function f() end
   local x = f() -- `x` is `nil` instead of `unknown`
   ```
+* `CHG` infer called function by params num
+  ```lua
+  ---@overload fun(number, number):string
+  ---@overload fun(number):number
+  ---@return boolean
+  local function f() end
+
+  local n1 = f()     -- `n1` is `boolean`
+  local n2 = f(0)    -- `n2` is `number`
+  local n3 = f(0, 0) -- `n3` is `string`
+  ```
 
 ## 3.3.1
 `2022-6-17`

script/vm/compiler.lua

@@ -554,32 +554,29 @@ local function getReturn(func, index, args)
         end
         return vm.compileNode(ast)
     end
-    local node = vm.compileNode(func)
+    local funcs = vm.getMatchedFunctions(func, args)
     ---@type vm.node?
     local result
-    for cnode in node:eachObject() do
-        if cnode.type == 'function'
-        or cnode.type == 'doc.type.function' then
-            local returnObject = vm.getReturnOfFunction(cnode, index)
-            if returnObject then
-                local returnNode = vm.compileNode(returnObject)
+    for _, mfunc in ipairs(funcs) do
+        local returnObject = vm.getReturnOfFunction(mfunc, index)
+        if returnObject then
+            local returnNode = vm.compileNode(returnObject)
+            for rnode in returnNode:eachObject() do
+                if rnode.type == 'generic' then
+                    returnNode = rnode:resolve(guide.getUri(func), args)
+                    break
+                end
+            end
+            if returnNode then
                 for rnode in returnNode:eachObject() do
-                    if rnode.type == 'generic' then
-                        returnNode = rnode:resolve(guide.getUri(func), args)
-                        break
+                    -- TODO: narrow type
+                    if rnode.type ~= 'doc.generic.name' then
+                        result = result or vm.createNode()
+                        result:merge(rnode)
                     end
                 end
-                if returnNode then
-                    for rnode in returnNode:eachObject() do
-                        -- TODO: narrow type
-                        if rnode.type ~= 'doc.generic.name' then
-                            result = result or vm.createNode()
-                            result:merge(rnode)
-                        end
-                    end
-                    if result and returnNode:isOptional() then
-                        result:addOptional()
-                    end
+                if result and returnNode:isOptional() then
+                    result:addOptional()
                 end
             end
         end

script/vm/function.lua

@@ -0,0 +1,127 @@
+---@class vm
+local vm = require 'vm.vm'
+
+---@param func parser.object
+---@return integer min
+---@return integer max
+function vm.countParamsOfFunction(func)
+    local min = 0
+    local max = 0
+    if func.type == 'function'
+    or func.type == 'doc.type.function' then
+        if func.args then
+            max = #func.args
+            min = max
+            for i = #func.args, 1, -1 do
+                local arg = func.args[i]
+                if arg.type == '...'
+                or (arg.name and arg.name[1] =='...') then
+                    max = math.huge
+                elseif not vm.compileNode(arg):isNullable() then
+                    min = i
+                    break
+                end
+            end
+        end
+    end
+    return min, max
+end
+
+---@param func parser.object
+---@return integer min
+---@return integer max
+function vm.countReturnsOfFunction(func)
+    if func.type == 'function' then
+        if not func.returns then
+            return 0, 0
+        end
+        local min, max
+        for _, ret in ipairs(func.returns) do
+            local rmin, rmax = vm.countList(ret)
+            if not min or rmin < min then
+                min = rmin
+            end
+            if not max or rmax > max then
+                max = rmax
+            end
+        end
+        return min, max
+    end
+    if func.type == 'doc.type.function' then
+        return vm.countList(func.returns)
+    end
+    return 0, 0
+end
+
+---@param func parser.object
+---@return integer min
+---@return integer max
+function vm.countReturnsOfCall(func, args)
+    local funcs = vm.getMatchedFunctions(func, args)
+    local min
+    local max
+    for _, f in ipairs(funcs) do
+        local rmin, rmax = vm.countReturnsOfFunction(f)
+        if not min or rmin < min then
+            min = rmin
+        end
+        if not max or rmax > max then
+            max = rmax
+        end
+    end
+    return min or 0, max or 0
+end
+
+---@param list parser.object[]?
+---@return integer min
+---@return integer max
+function vm.countList(list)
+    if not list then
+        return 0, 0
+    end
+    local lastArg = list[#list]
+    if not lastArg then
+        return 0, 0
+    end
+    if lastArg.type == '...' then
+        return #list - 1, math.huge
+    end
+    if lastArg.type == 'call' then
+        local rmin, rmax = vm.countReturnsOfCall(lastArg.node, lastArg.args)
+        return #list - 1 + rmin, #list - 1 + rmax
+    end
+    return #list, #list
+end
+
+---@param func parser.object
+---@param args parser.object[]?
+---@return parser.object[]
+function vm.getMatchedFunctions(func, args)
+    local funcs = {}
+    local node = vm.compileNode(func)
+    for n in node:eachObject() do
+        if n.type == 'function'
+        or n.type == 'doc.type.function' then
+            funcs[#funcs+1] = n
+        end
+    end
+    if #funcs <= 1 then
+        return funcs
+    end
+
+    local amin, amax = vm.countList(args)
+
+    local matched = {}
+    for _, n in ipairs(funcs) do
+        local min, max = vm.countParamsOfFunction(n)
+        if amin >= min and amax <= max then
+            matched[#matched+1] = n
+        end
+    end
+
+    if #matched == 0 then
+        return funcs
+    else
+        return matched
+    end
+end

script/vm/init.lua

@@ -17,4 +17,5 @@ require 'vm.generic'
 require 'vm.sign'
 require 'vm.local-id'
 require 'vm.global'
+require 'vm.function'
 return vm

test/type_inference/init.lua

@@ -2835,3 +2835,128 @@ local b
 
 local <?x?> = echo(b)
 ]]
+
+TEST 'boolean' [[
+---@return boolean
+function f()
+end
+
+---@param x integer
+---@return number
+function f(x)
+end
+
+local <?x?> = f()
+]]
+
+TEST 'number' [[
+---@return boolean
+function f()
+end
+
+---@param x integer
+---@return number
+function f(x)
+end
+
+local <?x?> = f(1)
+]]
+
+TEST 'boolean' [[
+---@return boolean
+function f()
+end
+
+---@param x integer
+---@return number
+function f(x)
+end
+
+function r0()
+    return
+end
+
+local <?x?> = f(r0())
+]]
+
+TEST 'number' [[
+---@return boolean
+function f()
+end
+
+---@param x integer
+---@return number
+function f(x)
+end
+
+function r1()
+    return 1
+end
+
+local <?x?> = f(r1())
+]]
+
+TEST 'boolean' [[
+---@return boolean
+function f()
+end
+
+---@param x integer
+---@return number
+function f(x)
+end
+
+---@type fun()
+local r0
+
+local <?x?> = f(r0())
+]]
+
+TEST 'number' [[
+---@return boolean
+function f()
+end
+
+---@param x integer
+---@return number
+function f(x)
+end
+
+---@type fun():integer
+local r1
+
+local <?x?> = f(r1())
+]]
+
+TEST 'boolean' [[
+---@overload fun(number, number):string
+---@overload fun(number):number
+---@return boolean
+local function f() end
+
+local <?n1?> = f()
+local n2 = f(0)
+local n3 = f(0, 0)
+]]
+
+TEST 'number' [[
+---@overload fun(number, number):string
+---@overload fun(number):number
+---@return boolean
+local function f() end
+
+local n1 = f()
+local <?n2?> = f(0)
+local n3 = f(0, 0)
+]]
+
+TEST 'string' [[
+---@overload fun(number, number):string
+---@overload fun(number):number
+---@return boolean
+local function f() end
+
+local n1 = f()
+local n2 = f(0)
+local <?n3?> = f(0, 0)
+]]