2 files changed, 42 insertions, 36 deletions
diff --git a/server-beta/src/core/definition.lua b/server-beta/src/core/definition.lua
index 2f4a2da3..bd8124d2 100644
--- a/server-beta/src/core/definition.lua
+++ b/server-beta/src/core/definition.lua
@@ -3,16 +3,14 @@ local engineer = require 'core.engineer'
 
 return function (ast, text, offset)
     local results = {}
-    local searcher = engineer(ast)
     guide.eachSource(ast.ast, offset, function (source)
-        searcher:eachRef(source, function (src, mode)
-            if mode == 'set' or mode == 'local' then
-                results[#results+1] = {
-                    uri    = ast.uri,
-                    source = source,
-                    target = src,
-                }
-            end
+        local searcher = engineer(ast)
+        searcher:eachRef(source, 'def', function (src, mode)
+            results[#results+1] = {
+                uri    = ast.uri,
+                source = source,
+                target = src,
+            }
         end)
     end)
     if #results == 0 then
diff --git a/server-beta/src/core/engineer.lua b/server-beta/src/core/engineer.lua
index a9251996..74c6ca33 100644
--- a/server-beta/src/core/engineer.lua
+++ b/server-beta/src/core/engineer.lua
@@ -12,41 +12,49 @@ local mt = {}
 mt.__index = mt
 mt.type = 'engineer'
 
---- 查找所有局部变量引用
-function mt:eachRefAsLocal(obj, callback)
-    callback(obj, 'local')
-    if obj.ref then
-        for _, ref in ipairs(obj.ref) do
-            if ref.type == 'setlocal' then
-                callback(ref, 'set')
-            elseif ref.type == 'getlocal' then
-                callback(ref, 'get')
+mt['local'] = function (self, source, mode, callback)
+    if mode == 'def' then
+        callback(source, 'local')
+        if source.ref then
+            for _, ref in ipairs(source.ref) do
+                if ref.type == 'setlocal' then
+                    callback(ref, 'set')
+                end
             end
         end
     end
 end
-
---- 查找所有域的引用
-function mt:eachRefAsField(obj, callback)
-    local node = obj.node
-    guide.eachChildRef(node, function (ref)
-        if ref.type == 'setglobal' or ref.type == 'setfield' then
-            callback(ref, 'set')
-        elseif ref.type == 'getglobal' or ref.type == 'getfield' then
-            callback(ref, 'get')
+mt['getlocal'] = function (self, source, mode, callback)
+    self:search(source.loc, 'local', mode, callback)
+end
+mt['setlocal'] = mt['getlocal']
+mt['_ENV'] = function (self, source, mode, callback)
+    if mode == 'def' then
+        if source.ref then
+            for _, ref in ipairs(source.ref) do
+                if ref.type == 'setglobal' then
+                    callback(ref, 'set')
+                end
+            end
         end
-    end)
+    end
 end
+mt['getglobal'] = function (self, source, mode, callback)
+    self:search(source.node, '_ENV', mode, callback)
+end
+mt['setglobal'] = mt['getglobal']
 
---- 查找所有引用
-function mt:eachRef(obj, callback)
-    if obj.type == 'local' then
-        self:eachRefAsLocal(obj, callback)
-    elseif obj.type == 'getlocal' or obj.type == 'setlocal' then
-        self:eachRefAsLocal(obj.loc, callback)
-    elseif obj.type == 'setglobal' or obj.type == 'getglobal' then
-        self:eachRefAsField(obj, callback)
+function mt:search(source, method, mode, callback)
+    local f = mt[method]
+    if not f then
+        return
     end
+    f(self, source, mode, callback)
+end
+
+function mt:eachRef(source, mode, callback)
+    local tp = source.type
+    self:search(source, tp, mode, callback)
 end
 
 return function (ast)