From 96946f4022a45519fa4e512c652d322cf8232a73 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9C=80=E8=90=8C=E5=B0=8F=E6=B1=90?= <sumneko@hotmail.com>
Date: Mon, 8 Nov 2021 16:44:04 +0800
Subject: yieldable check linked

---
 script/core/diagnostics/not-yieldable.lua |  5 +--
 script/vm/getDocs.lua                     |  4 +++
 test/diagnostics/init.lua                 | 59 +++++++++++++++++++++++++++++--
 3 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/script/core/diagnostics/not-yieldable.lua b/script/core/diagnostics/not-yieldable.lua
index ed853380..0be92e82 100644
--- a/script/core/diagnostics/not-yieldable.lua
+++ b/script/core/diagnostics/not-yieldable.lua
@@ -43,10 +43,11 @@ return function (uri, callback)
         end
         for i, arg in ipairs(source.args) do
             if  vm.isAsync(arg, true)
+            and not vm.isLinkedCall(source.node, i)
             and not isYieldAble(defs, i) then
                 callback {
-                    start   = source.node.start,
-                    finish  = source.node.finish,
+                    start   = arg.start,
+                    finish  = arg.finish,
                     message = lang.script('DIAG_NOT_YIELDABLE', i),
                 }
             end
diff --git a/script/vm/getDocs.lua b/script/vm/getDocs.lua
index f5c4e4e1..748d17c2 100644
--- a/script/vm/getDocs.lua
+++ b/script/vm/getDocs.lua
@@ -291,6 +291,10 @@ local function isLinkedCall(node, index)
     return false
 end
 
+function vm.isLinkedCall(node, index)
+    return isLinkedCall(node, index)
+end
+
 function vm.isAsyncCall(call)
     if vm.isAsync(call.node, true) then
         return true
diff --git a/test/diagnostics/init.lua b/test/diagnostics/init.lua
index 96e9e16d..e8e24141 100644
--- a/test/diagnostics/init.lua
+++ b/test/diagnostics/init.lua
@@ -7,7 +7,6 @@ local catch  = require 'catch'
 config.get 'Lua.diagnostics.neededFileStatus'['deprecated']    = 'Any'
 config.get 'Lua.diagnostics.neededFileStatus'['type-check']    = 'Any'
 config.get 'Lua.diagnostics.neededFileStatus'['await-in-sync'] = 'Any'
-config.get 'Lua.diagnostics.neededFileStatus'['not-yieldable'] = 'Any'
 
 rawset(_G, 'TEST', true)
 
@@ -1419,7 +1418,6 @@ end
 <!f()!>
 ]]
 
-
 TEST [[
 ---@nodiscard
 local function f()
@@ -1428,3 +1426,60 @@ end
 
 X = f()
 ]]
+
+config.get 'Lua.diagnostics.neededFileStatus'['not-yieldable'] = 'Any'
+TEST [[
+local function f(cb)
+    return cb
+end
+
+---@async
+local function af()
+    return nil
+end
+
+f(<!af!>)
+]]
+
+TEST [[
+---@param cb async fun()
+local function f(cb)
+    return cb
+end
+
+---@async
+local function af()
+    return nil
+end
+
+f(af)
+]]
+
+TEST [[
+local function f(cb)
+    cb()
+end
+
+local function af()
+    <!f!>(function () ---@async
+        return nil
+    end)
+end
+
+return af
+]]
+
+TEST [[
+local function f(cb)
+    cb()
+end
+
+---@async
+local function af()
+    f(function () ---@async
+        return nil
+    end)
+end
+
+return af
+]]
-- 
cgit v1.2.3