stash

2 files changed, 44 insertions, 7 deletions

diff --git a/script/core/diagnostics/missing-return-value.lua b/script/core/diagnostics/missing-return-value.lua
index 2156d66c..3aa7bd55 100644
--- a/script/core/diagnostics/missing-return-value.lua
+++ b/script/core/diagnostics/missing-return-value.lua
@@ -4,6 +4,7 @@ local vm     = require 'vm'
 local lang   = require 'language'
 local await  = require 'await'
 
+---@param func parser.object
 local function hasDocReturn(func)
     if not func.bindDocs then
         return false
@@ -12,10 +13,36 @@ local function hasDocReturn(func)
         if doc.type == 'doc.return' then
             return true
         end
+        if doc.type == 'doc.overload' then
+            if #doc.overload.returns > 0 then
+                return true
+            end
+        end
     end
     return false
 end
 
+---@param func parser.object
+---@return integer
+local function getReturnsMin(func)
+    local min = vm.countReturnsOfFunction(func)
+    if min == 0 then
+        return 0
+    end
+    for _, doc in ipairs(func.bindDocs) do
+        if doc.type == 'doc.overload' then
+            local n = vm.countReturnsOfFunction(doc.overload)
+            if n == 0 then
+                return 0
+            end
+            if n < min then
+                min = n
+            end
+        end
+    end
+    return min
+end
+
 ---@async
 return function (uri, callback)
     local state = files.getState(uri)
@@ -26,15 +53,15 @@ return function (uri, callback)
     ---@async
     guide.eachSourceType(state.ast, 'function', function (source)
         await.delay()
-        if not hasDocReturn(source) then
+        local returns = source.returns
+        if not returns then
             return
         end
-        local min = vm.countReturnsOfFunction(source)
-        if min == 0 then
+        if not hasDocReturn(source) then
             return
         end
-        local returns = source.returns
-        if not returns then
+        local min = getReturnsMin(source)
+        if min == 0 then
             return
         end
         for _, ret in ipairs(returns) do
diff --git a/script/core/diagnostics/return-type-mismatch.lua b/script/core/diagnostics/return-type-mismatch.lua
index cce4aad8..2ff8a909 100644
--- a/script/core/diagnostics/return-type-mismatch.lua
+++ b/script/core/diagnostics/return-type-mismatch.lua
@@ -3,6 +3,7 @@ local lang  = require 'language'
 local guide = require 'parser.guide'
 local vm    = require 'vm'
 local await = require 'await'
+local util  = require 'utility'
 
 ---@param func parser.object
 ---@return vm.node[]?
@@ -10,14 +11,23 @@ local function getDocReturns(func)
     if not func.bindDocs then
         return nil
     end
-    local returns = {}
+    ---@type table<integer, vm.node>
+    local returns = util.defaultTable(function ()
+        return vm.createNode()
+    end)
     for _, doc in ipairs(func.bindDocs) do
         if doc.type == 'doc.return' then
             for _, ret in ipairs(doc.returns) do
-                returns[ret.returnIndex] = vm.compileNode(ret)
+                returns[ret.returnIndex]:merge(vm.compileNode(ret))
+            end
+        end
+        if doc.type == 'doc.overload' then
+            for i, ret in ipairs(doc.overload.returns) do
+                returns[i]:merge(vm.compileNode(ret))
             end
         end
     end
+    setmetatable(returns, nil)
     if #returns == 0 then
         return nil
     end