support `@type` for `return`

#2144

3 files changed, 18 insertions, 2 deletions

diff --git a/script/parser/luadoc.lua b/script/parser/luadoc.lua
index 6a1e8fff..3454b241 100644
--- a/script/parser/luadoc.lua
+++ b/script/parser/luadoc.lua
@@ -1822,7 +1822,7 @@ local function bindDocsBetween(sources, binded, start, finish)
                 or src.type == 'setindex'
                 or src.type == 'setmethod'
                 or src.type == 'function'
-                or src.type == 'table'
+                or src.type == 'return'
                 or src.type == '...' then
                     if bindDoc(src, binded) then
                         ok = true
@@ -1934,7 +1934,7 @@ local bindDocAccept = {
     'local'     , 'setlocal'  , 'setglobal',
     'setfield'  , 'setmethod' , 'setindex' ,
     'tablefield', 'tableindex', 'self'     ,
-    'function'  , 'table'     , '...'      ,
+    'function'  , 'return'     , '...'      ,
 }
 
 local function bindDocs(state)
diff --git a/script/vm/compiler.lua b/script/vm/compiler.lua
index a4147d70..543ee0d4 100644
--- a/script/vm/compiler.lua
+++ b/script/vm/compiler.lua
@@ -766,6 +766,11 @@ function vm.selectNode(list, index)
     if not exp then
         return vm.createNode(vm.declareGlobal('type', 'nil')), nil
     end
+
+    if vm.bindDocs(list) then
+        return vm.compileNode(list), exp
+    end
+
     ---@type vm.node?
     local result
     if exp.type == 'call' then
diff --git a/test/type_inference/init.lua b/test/type_inference/init.lua
index 9f5c8493..aba6b0b4 100644
--- a/test/type_inference/init.lua
+++ b/test/type_inference/init.lua
@@ -4314,6 +4314,7 @@ repeat
 until <?x?>
 ]]
 
+-- #2144
 TEST 'A' [=[
 local function f()
     return {} --[[@as A]]
@@ -4321,3 +4322,13 @@ end
 
 local <?x?> = f()
 ]=]
+
+TEST 'A' [=[
+local function f()
+    ---@type A
+    return {}
+end
+
+local <?x?> = f()
+]=]
+--