From f5fcab05b5a30373d1b35689d456090c8299c4c2 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: Fri, 11 Aug 2023 18:12:04 +0800
Subject: support `---@type` and `--[[@as]]` for return

resolve #2144
---
 changelog.md                        |  1 +
 script/vm/compiler.lua              | 10 ++++++++++
 test/diagnostics/missing-fields.lua |  8 ++++++++
 3 files changed, 19 insertions(+)

diff --git a/changelog.md b/changelog.md
index 78810cb8..8cbfc1d2 100644
--- a/changelog.md
+++ b/changelog.md
@@ -1,6 +1,7 @@
 # changelog
 
 ## 3.7.0
+* `NEW` support `---@type` and `--[[@as]]` for return statement
 * `FIX` wrong hover and signature for method with varargs and overloads
 * `FIX` [#2224]
 
diff --git a/script/vm/compiler.lua b/script/vm/compiler.lua
index 543ee0d4..7faa4f47 100644
--- a/script/vm/compiler.lua
+++ b/script/vm/compiler.lua
@@ -1168,6 +1168,16 @@ local compilerSwitch = util.switch()
             vm.compileCallArg(source, call)
         end
 
+        if source.parent.type == 'return' then
+            local myIndex = util.arrayIndexOf(source.parent, source)
+            ---@cast myIndex -?
+            local parentNode = vm.selectNode(source.parent, myIndex)
+            if not parentNode:isEmpty() then
+                vm.setNode(source, parentNode)
+                return
+            end
+        end
+
         if source.parent.type == 'setglobal'
         or source.parent.type == 'local'
         or source.parent.type == 'setlocal'
diff --git a/test/diagnostics/missing-fields.lua b/test/diagnostics/missing-fields.lua
index d8731775..f5fdd35c 100644
--- a/test/diagnostics/missing-fields.lua
+++ b/test/diagnostics/missing-fields.lua
@@ -197,3 +197,11 @@ local b = {
     c = 3,
 }
 ]]
+
+TEST [[
+---@class A
+---@field x integer
+
+---@type A
+return <!{}!>
+]]
-- 
cgit v1.2.3