重构搜索器

3 files changed, 84 insertions, 270 deletions

diff --git a/server-beta/src/core/definition.lua b/server-beta/src/core/definition.lua
index 6724ca89..afd368b1 100644
--- a/server-beta/src/core/definition.lua
+++ b/server-beta/src/core/definition.lua
@@ -1,69 +1,22 @@
 local guide    = require 'parser.guide'
 local engineer = require 'core.engineer'
 
-local m = {}
-
-function m.search(state, ast, source)
-    if not source then
-        return
-    end
-    if state.cache[source] then
-        return
-    end
-    state.cache[source] = true
-    local f = m['as' .. source.type]
-    if not f then
-        return
-    end
-    f(state, ast, source)
-end
-
-function m.aslocal(state, ast, source)
-    local searcher = engineer(ast)
-    searcher:eachLocalRef(source, function (src, mode)
-        if mode == 'local' or mode == 'set' then
-            state.callback(src)
-        end
-    end)
-    searcher:eachClass(source, function (src)
-        state.callback(src)
-    end)
-end
-
-m.asgetlocal = m.aslocal
-m.assetlocal = m.aslocal
-
-function m.globals(state, ast, source)
-    local searcher = engineer(ast)
-    searcher:eachGloablOfName(source[1], function (src, mode)
-        if mode == 'set' then
-            state.callback(src)
-        end
-    end)
-end
-
-m.assetglobal = m.globals
-m.asgetglobal = m.globals
-
 return function (ast, text, offset)
     local results = {}
-    local state = {
-        ast = ast,
-        cache = {},
-    }
-    function state.callback(target, uri)
-        if target.start == 0 then
-            return
-        end
-        results[#results+1] = {
-            uri    = uri or ast.uri,
-            source = state.source,
-            target = target,
-        }
-    end
+    local searcher = engineer(ast)
     guide.eachSource(ast.root, offset, function (source)
-        state.source = source
-        m.search(state, ast, source)
+        searcher:eachRef(source, function (src, mode)
+            if src.start == 0 then
+                return
+            end
+            if mode == 'local' or mode == 'set' then
+                results[#results+1] = {
+                    uri    = ast.uri,
+                    source = source,
+                    target = src,
+                }
+            end
+        end)
     end)
     if #results == 0 then
         return nil
diff --git a/server-beta/src/core/engineer.lua b/server-beta/src/core/engineer.lua
index a093a90a..45cff84e 100644
--- a/server-beta/src/core/engineer.lua
+++ b/server-beta/src/core/engineer.lua
@@ -3,6 +3,7 @@ local config = require 'config'
 
 local type         = type
 local setmetatable = setmetatable
+local ipairs       = ipairs
 
 _ENV = nil
 
@@ -11,238 +12,78 @@ local mt = {}
 mt.__index = mt
 mt.type = 'engineer'
 
-function mt:call(method, obj, ...)
-    self.step = self.step + 1
-    if self.step > 100 then
-        return nil
-    end
-    if not obj then
-        return nil
-    end
-    if ... == nil and obj['_'..method] ~= nil then
-        return obj['_'..method]
-    end
-    local res = self[method](self, obj, ...)
-    self.step = self.step - 1
-    if ... == nil then
-        obj['_'..method] = res
-    end
-    return res
-end
-
---- 根据变量名，遍历全局变量
-function mt:eachGloablOfName(name, callback)
-    if type(name) ~= 'string' then
-        return
-    end
-    guide.eachSourceOf(self.ast.root, {
-        ['setglobal'] = function (source)
-            if source[1] == name then
-                callback(source, 'set')
-            end
-        end,
-        ['getglobal'] = function (source)
-            if source[1] == name then
-                callback(source, 'get')
-            end
-        end,
-        ['setfield'] = function (source)
-            if source.field[1] ~= name then
-                return
-            end
-            if self:call('isGlobalField', source) then
-                callback(source, 'set')
-            end
-        end,
-        ['getfield'] = function (source)
-            if source.field[1] ~= name then
-                return
-            end
-            if self:call('isGlobalField', source) then
-                callback(source, 'get')
-            end
-        end,
-        ['call'] = function (source)
-            local d = self:call('asRawSet', source)
-            if d then
-                if self:call('getLiteral', d.k) == name then
-                    callback(source, 'set')
-                end
-            end
-            local d = self:call('asRawGet', source)
-            if d then
-                if self:call('getLiteral', d.k) == name then
-                    callback(source, 'get')
-                end
-            end
-        end,
-    })
-end
-
---- 是否是全局变量
-function mt:isGlobal(obj)
-    if obj.type == 'getglobal' then
-        return true
-    end
-    if obj.type == 'getfield' then
-        return self:call('isGlobalField', obj)
-    end
-    if obj.type == 'call' then
-        local d = self:call('asRawGet', obj)
-        return not not d
-    end
-    return false
-end
-
---- 是否是指定名称的全局变量
-function mt:isGlobalOfName(obj, name)
-    if not self:call('isGlobal', obj) then
-        return false
-    end
-    return self:call('getName', obj) == name
-end
-
---- 获取名称
-function mt:getName(obj)
-    if obj.type == 'setglobal' or obj.type == 'getglobal' then
-        return obj[1]
-    elseif obj.type == 'setfield' or obj.type == 'getfield' then
-        return obj.field[1]
-    elseif obj.type == 'local' or obj.type == 'setlocal' or obj.type == 'getlocal' then
-        return obj[1]
-    end
-    return false
-end
-
---- 获取字面量值
-function mt:getLiteral(obj)
-    if obj.type == 'number' then
-        return obj[1]
-    elseif obj.type == 'boolean' then
-        return obj[1]
-    elseif obj.type == 'string' then
+--- 获取对象作为域时的名字
+function mt:getFieldName(obj)
+    if obj.type == 'getglobal' or obj.type == 'setglobal' then
         return obj[1]
     end
     return nil
 end
 
---- 是否是全局field
----|_G.xxx
----|_ENV.xxx
----|_ENV._G.xxx
-function mt:isGlobalField(obj)
-    local node = self.ast.root[obj.node]
-    if self:call('isG', node) then
-        return true
-    end
-    if self:call('isENV', node) then
-        return true
+--- 查找所有局部变量引用
+function mt:eachRefAsLocal(obj, callback)
+    callback(obj, 'local')
+    if obj.ref then
+        for _, ref in ipairs(obj.ref) do
+            local refObj = self.root[ref]
+            if refObj.type == 'setlocal' then
+                callback(self.root[ref], 'set')
+            elseif refObj.type == 'getlocal' then
+                callback(self.root[ref], 'get')
+            elseif refObj.type == 'setglobal' then
+                callback(self.root[ref], 'settable', refObj)
+            elseif refObj.type == 'getglobal' then
+                callback(self.root[ref], 'gettable', refObj)
+            end
+        end
     end
-    return false
 end
 
---- 是否是_ENV
-function mt:isENV(obj)
+--- 查找所有全局变量引用
+function mt:eachRefAsGlobal(obj, callback)
     local version = config.config.runtime.version
-    if version == 'Lua 5.1' or version == 'LuaJIT' then
-        return false
-    end
-    if self:isGlobalOfName(obj, '_ENV') then
-        return true
-    end
-    return false
-end
-
---- 是否是_G
-function mt:isG(obj)
-    if self:isGlobalOfName(obj, '_G') then
-        return true
-    end
-    return false
-end
-
---- 获取call的参数
-function mt:getCallArg(obj, i)
-    local args = self.ast.root[obj.args]
-    if not args then
-        return nil
-    end
-    return self.ast.root[args[i]]
-end
-
---- 获取rawset信息
-function mt:asRawSet(obj)
-    local node = self.ast.root[obj.node]
-    if not self:isGlobalOfName(node, 'rawset') then
-        return false
+    if version ~= 'Lua 5.1' and version ~= 'LuaJIT' then
+        local env   = guide.getLocal(self.root, obj, '_ENV', obj.start)
+        local field = self:getFieldName(obj)
+        self:eachRefAsField(env, field, callback)
+        return
     end
-    return {
-        t = self:getCallArg(obj, 1),
-        k = self:getCallArg(obj, 2),
-        v = self:getCallArg(obj, 3),
-    }
 end
 
---- 获取rawget信息
-function mt:asRawGet(obj)
-    local node = self.ast.root[obj.node]
-    if not self:isGlobalOfName(node, 'rawget') then
-        return false
-    end
-    return {
-        t = self:getCallArg(obj, 1),
-        k = self:getCallArg(obj, 2),
-    }
+--- 查找所有域引用
+function mt:eachRefAsField(parent, field, callback)
+    self:eachRef(parent, function (src, mode)
+        if self:getFieldName(src) ~= field then
+            return
+        end
+        if mode == 'settable' then
+            callback(src, 'set')
+        elseif mode == 'gettable' then
+            callback(src, 'get')
+        end
+    end)
 end
 
---- 根据指定的局部变量，遍历局部变量引用
-function mt:eachLocalRef(obj, callback)
-    if not obj then
-        return
-    end
-    local src
+--- 查找所有引用
+function mt:eachRef(obj, callback)
     if obj.type == 'local' then
-        src = obj
+        self:eachRefAsLocal(obj, callback)
     elseif obj.type == 'getlocal' or obj.type == 'setlocal' then
-        src = self.ast.root[obj.loc]
-    else
-        return
-    end
-    callback(src, 'local')
-    if src.ref then
-        for i = 1, #src.ref do
-            local ref = src.ref[i]
-            local refObj = self.ast.root[ref]
-            if refObj.type == 'setlocal' then
-                callback(refObj, 'set')
-            elseif refObj.type == 'getlocal' then
-                callback(refObj, 'get')
-            end
-        end
-    end
-end
-
---- 遍历class
-function mt:eachClass(obj, callback)
-    local root = self.ast.root
-    if obj.type == 'setlocal' or obj.type == 'getlocal' then
-        local loc = root[obj.loc]
-        local setmethod = root[loc.method]
-        if setmethod then
-            local node = root[setmethod.node]
-            self:call('eachLocalRef', node, function (src, mode)
-                if mode == 'local' or mode == 'set' then
-                    callback(src)
-                end
-            end)
-        end
+        local loc = self.root[obj.loc]
+        self:eachRefAsLocal(loc, callback)
+    elseif obj.type == 'setglobal' or obj.type == 'getglobal' then
+        self:eachRefAsGlobal(obj, callback)
     end
 end
 
 return function (ast)
+    if not ast.vm then
+        ast.vm = {}
+    end
     local self = setmetatable({
         step = 0,
-        ast  = ast,
+        root = ast.root,
+        vm   = ast.vm,
     }, mt)
     return self
 end
diff --git a/server-beta/src/parser/compile.lua b/server-beta/src/parser/compile.lua
index 407b778f..8cfcd7a2 100644
--- a/server-beta/src/parser/compile.lua
+++ b/server-beta/src/parser/compile.lua
@@ -39,6 +39,16 @@ local vmMap = {
             loc.ref[#loc.ref+1] = id
         else
             obj.type = 'getglobal'
+            if ENVMode == '_ENV' then
+                local node = guide.getLocal(Root, obj, '_ENV', obj.start)
+                if node then
+                    if not node.ref then
+                        node.ref = {}
+                    end
+                    node.ref[#node.ref+1] = id
+                    obj.node = Cache[node]
+                end
+            end
         end
         return id
     end,
@@ -267,6 +277,16 @@ local vmMap = {
             loc.ref[#loc.ref+1] = id
         else
             obj.type = 'setglobal'
+            if ENVMode == '_ENV' then
+                local node = guide.getLocal(Root, obj, '_ENV', obj.start)
+                if node then
+                    if not node.ref then
+                        node.ref = {}
+                    end
+                    node.ref[#node.ref+1] = id
+                    obj.node = Cache[node]
+                end
+            end
         end
         return id
     end,