From 33ca6151e44074754373a1eeebbd5f53e92f4fae 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: Wed, 4 Dec 2019 14:59:28 +0800
Subject: =?UTF-8?q?=E6=94=AF=E6=8C=81=E5=A4=9A=E9=87=8Dcall?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 script-beta/vm/eachRef.lua    | 28 +++++++++++++++++++++++++++-
 test-beta/references/init.lua | 26 ++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/script-beta/vm/eachRef.lua b/script-beta/vm/eachRef.lua
index d6ca819b..9ab9d6cf 100644
--- a/script-beta/vm/eachRef.lua
+++ b/script-beta/vm/eachRef.lua
@@ -254,8 +254,31 @@ local function asReturn(source, callback)
     local func = guide.getParentFunction(source)
     if func.type == 'main' then
     else
+        local index
+        for i = 1, #parent do
+            if parent[i] == source then
+                index = i
+                break
+            end
+        end
+        if not index then
+            return
+        end
         vm.eachRef(func, function (info)
-
+            local src = info.source
+            local call = src.parent
+            if not call or call.type ~= 'call' then
+                return
+            end
+            local recvs = getCallRecvs(call)
+            if recvs and recvs[index] then
+                vm.eachRef(recvs[index], callback)
+            elseif index == 1 then
+                callback {
+                    type   = 'call',
+                    source = call,
+                }
+            end
         end)
     end
 end
@@ -276,6 +299,7 @@ local function ofLocal(loc, callback)
                     mode     = 'get',
                 }
                 asValue(ref, callback)
+                asReturn(ref, callback)
             elseif ref.type == 'setlocal' then
                 callback {
                     source   = ref,
@@ -455,6 +479,8 @@ local function eachRef(source, callback)
     elseif stype == 'table'
     or     stype == 'function' then
         ofValue(source, callback)
+    elseif stype == 'call' then
+        ofCall(source.node, 1, callback)
     elseif stype == 'main' then
         ofMain(source, callback)
     end
diff --git a/test-beta/references/init.lua b/test-beta/references/init.lua
index 6ba1f5e4..4cd21e2b 100644
--- a/test-beta/references/init.lua
+++ b/test-beta/references/init.lua
@@ -141,8 +141,24 @@ local obj = setmetatable({}, mt)
 print(obj.<?x?>)
 ]]
 
+TEST [[
+local <!x!>
+local function f()
+    return <!x!>
+end
+local <?y?> = f()
+]]
+
 TEST [[
 local <?x?>
+local function f()
+    return <!x!>
+end
+local <!y!> = f()
+]]
+
+TEST [[
+local <!x!>
 local function f()
     return function ()
         return <!x!>
@@ -151,6 +167,16 @@ end
 local <?y?> = f()()
 ]]
 
+TEST [[
+local <?x?>
+local function f()
+    return function ()
+        return <!x!>
+    end
+end
+local <!y!> = f()()
+]]
+
 --TEST [[
 -----@class <!Class!>
 -----@type <?Class?>
-- 
cgit v1.2.3