support @type for return

#2144

Author: sumneko

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)

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

test/type_inference/init.lua

@@ -4314,10 +4314,21 @@ repeat
 until <?x?>
 ]]
 
+-- #2144
 TEST 'A' [=[
 local function f()
     return {} --[[@as A]]
 end
 
 local <?x?> = f()
 ]=]
+
+TEST 'A' [=[
+local function f()
+    ---@type A
+    return {}
+end
+
+local <?x?> = f()
+]=]
+--