2 files changed, 34 insertions, 28 deletions
diff --git a/script/core/diagnostics/missing-return.lua b/script/core/diagnostics/missing-return.lua
index 5806f121..435e8a96 100644
--- a/script/core/diagnostics/missing-return.lua
+++ b/script/core/diagnostics/missing-return.lua
@@ -4,25 +4,6 @@ local vm     = require 'vm'
 local lang   = require 'language'
 local await  = require 'await'
 
----@param uri uri
----@param func parser.object
-local function hasDocReturn(uri, func)
-    if not func.bindDocs then
-        return false
-    end
-    for _, doc in ipairs(func.bindDocs) do
-        if doc.type == 'doc.return' then
-            -- don't need return with only one `any`
-            local lastReturn = doc.returns[#doc.returns]
-            if lastReturn.returnIndex ~= 1
-            or vm.getInfer(lastReturn):view(uri) ~= 'any' then
-                return true
-            end
-        end
-    end
-    return false
-end
-
 ---@param block parser.object
 ---@return boolean
 local function hasReturn(block)
@@ -57,6 +38,17 @@ local function hasReturn(block)
     return false
 end
 
+---@param func parser.object
+---@return boolean
+local function isEmptyFunction(func)
+    if #func > 0 then
+        return false
+    end
+    local startRow  = guide.rowColOf(func.start)
+    local finishRow = guide.rowColOf(func.finish)
+    return finishRow - startRow <= 1
+end
+
 ---@async
 return function (uri, callback)
     local state = files.getState(uri)
@@ -67,21 +59,27 @@ return function (uri, callback)
     ---@async
     guide.eachSourceType(state.ast, 'function', function (source)
         -- check declare only
-        if #source == 0 then
+        if isEmptyFunction(source) then
             return
         end
         await.delay()
-        if not hasDocReturn(uri, source) then
+        if vm.countReturnsOfFunction(source) == 0 then
             return
         end
         if hasReturn(source) then
             return
         end
         local lastAction = source[#source]
-        local finish = lastAction.range or lastAction.finish
+        local pos
+        if lastAction then
+            pos = lastAction.range or lastAction.finish
+        else
+            local row = guide.rowColOf(source.finish)
+            pos = guide.positionOf(row - 1, 0)
+        end
         callback {
-            start   = finish,
-            finish  = finish,
+            start   = pos,
+            finish  = pos,
             message = lang.script('DIAG_MISSING_RETURN'),
         }
     end)
diff --git a/script/core/diagnostics/unused-local.lua b/script/core/diagnostics/unused-local.lua
index a637f427..8bff7dcb 100644
--- a/script/core/diagnostics/unused-local.lua
+++ b/script/core/diagnostics/unused-local.lua
@@ -63,16 +63,24 @@ local function isDocClass(source)
     return false
 end
 
+---@param func parser.object
+---@return boolean
+local function isEmptyFunction(func)
+    if #func > 0 then
+        return false
+    end
+    local startRow  = guide.rowColOf(func.start)
+    local finishRow = guide.rowColOf(func.finish)
+    return finishRow - startRow <= 1
+end
+
 ---@param source parser.object
 local function isDeclareFunctionParam(source)
     if source.parent.type ~= 'funcargs' then
         return false
     end
     local func = source.parent.parent
-    if #func > 0 then
-        return false
-    end
-    return true
+    return isEmptyFunction(func)
 end
 
 return function (uri, callback)