cleanup

9 files changed, 307 insertions, 172 deletions

diff --git a/script/vm/getDef.lua b/script/vm/getDef.lua
index 339053e4..2741ae46 100644
--- a/script/vm/getDef.lua
+++ b/script/vm/getDef.lua
@@ -3,6 +3,7 @@ local vm       = require 'vm.vm'
 local util     = require 'utility'
 local compiler = require 'vm.node.compiler'
 local guide    = require 'parser.guide'
+local localID  = require 'vm.local-id'
 
 local simpleMap
 
@@ -42,15 +43,17 @@ simpleMap = util.switch()
     end)
     : case 'field'
     : call(function (source, results)
-        local node = source.parent.node
-        if node.type == 'getlocal' then
-            searchGetLocal(source, node, results)
-        end
+        local parent = source.parent
+        simpleMap[parent.type](parent, results)
     end)
     : case 'setfield'
     : case 'getfield'
     : call(function (source, results)
-        simpleMap['field'](source.field, results)
+        local node = source.node
+        if node.type == 'getlocal' then
+            searchGetLocal(source, node, results)
+            return
+        end
     end)
     : case 'getindex'
     : case 'setindex'
@@ -76,11 +79,11 @@ local searchFieldMap = util.switch()
     end)
     : getMap()
 
-local compiledMap;compiledMap = util.switch()
+local nodeMap;nodeMap = util.switch()
     : case 'field'
     : call(function (source, results)
         local parent = source.parent
-        compiledMap[parent.type](parent, results)
+        nodeMap[parent.type](parent, results)
     end)
     : case 'getfield'
     : case 'setfield'
@@ -119,12 +122,19 @@ local function searchByGlobal(source, results)
     end
 end
 
+local function searchByID(source, results)
+    local idSources = localID.getSources(source)
+    if not idSources then
+        return
+    end
+end
+
 ---@param source  parser.object
 ---@param results parser.object[]
-local function searchByCompiled(source, results)
-    local compiled = compiledMap[source.type]
-    if compiled then
-        compiled(source, results)
+local function searchByNode(source, results)
+    local node = nodeMap[source.type]
+    if node then
+        node(source, results)
     end
 end
 
@@ -135,7 +145,8 @@ function vm.getDefs(source)
 
     searchBySimple(source, results)
     searchByGlobal(source, results)
-    searchByCompiled(source, results)
+    --searchByID(source, results)
+    searchByNode(source, results)
 
     return results
 end
diff --git a/script/vm/global-id.lua b/script/vm/global-id.lua
new file mode 100644
index 00000000..614e28c5
--- /dev/null
+++ b/script/vm/global-id.lua
@@ -0,0 +1,132 @@
+local util          = require 'utility'
+local guide         = require 'parser.guide'
+local globalBuilder = require 'vm.node.global'
+
+---@class parser.object
+---@field _globalID vm.node.global
+
+---@class vm.global-id
+local m = {}
+---@type table<string, vm.node.global>
+m.globals = util.defaultTable(globalBuilder)
+---@type table<uri, table<string, boolean>>
+m.globalSubs = util.defaultTable(function ()
+    return {}
+end)
+
+m.ID_SPLITE = '\x1F'
+
+local compilerGlobalMap = util.switch()
+    : case 'local'
+    : call(function (uri, source)
+        if source.tag ~= '_ENV' then
+            return
+        end
+        if source.ref then
+            for _, ref in ipairs(source.ref) do
+                m.compileNode(uri, ref)
+            end
+        end
+    end)
+    : case 'setglobal'
+    : call(function (uri, source)
+        local name = guide.getKeyName(source)
+        source._globalID = m.declareGlobal(name, uri, source)
+    end)
+    : case 'getglobal'
+    : call(function (uri, source)
+        local name   = guide.getKeyName(source)
+        local global = m.getGlobal(name)
+        global:addGet(uri, source)
+        source._globalID = global
+
+        local nxt = source.next
+        if nxt then
+            m.compileNode(uri, nxt)
+        end
+    end)
+    : case 'setfield'
+    ---@param uri    uri
+    ---@param source parser.object
+    : call(function (uri, source)
+        local parent = source.node._globalID
+        if not parent then
+            return
+        end
+        local name = parent:getName() .. m.ID_SPLITE .. guide.getKeyName(source)
+        source._globalID = m.declareGlobal(name, uri, source)
+    end)
+    : case 'getfield'
+    ---@param uri    uri
+    ---@param source parser.object
+    : call(function (uri, source)
+        local parent = source.node._globalID
+        if not parent then
+            return
+        end
+        local name = parent:getName() .. m.ID_SPLITE .. guide.getKeyName(source)
+        local global = m.getGlobal(name)
+        global:addGet(uri, source)
+        source._globalID = global
+
+        local nxt = source.next
+        if nxt then
+            m.compileNode(uri, nxt)
+        end
+    end)
+    : getMap()
+
+
+---@param name   string
+---@param uri    uri
+---@param source parser.object
+---@return vm.node.global
+function m.declareGlobal(name, uri, source)
+    m.globalSubs[uri][name] = true
+    local node = m.globals[name]
+    node:addSet(uri, source)
+    return node
+end
+
+---@param name string
+---@param uri? uri
+---@return vm.node.global
+function m.getGlobal(name, uri)
+    if uri then
+        m.globalSubs[uri][name] = true
+    end
+    return m.globals[name]
+end
+
+---@param source parser.object
+function m.compileNode(uri, source)
+    if source._globalID ~= nil then
+        return
+    end
+    source._globalID = false
+    local compiler = compilerGlobalMap[source.type]
+    if compiler then
+        compiler(uri, source)
+    end
+end
+
+---@param source parser.object
+function m.compileAst(source)
+    local uri = guide.getUri(source)
+    local env = guide.getENV(source)
+    m.compileNode(uri, env)
+end
+
+function m.getID(source)
+    return source._globalID
+end
+
+function m.dropUri(uri)
+    local globalSub = m.globalSubs[uri]
+    m.globalSubs[uri] = nil
+    for name in pairs(globalSub) do
+        m.globals[name]:dropUri(uri)
+    end
+end
+
+return m
diff --git a/script/vm/init.lua b/script/vm/init.lua
index f0bbb032..5a5de99e 100644
--- a/script/vm/init.lua
+++ b/script/vm/init.lua
@@ -1,5 +1,5 @@
 local vm = require 'vm.vm'
-require 'vm.node'
+require 'vm.state'
 require 'vm.getGlobals'
 require 'vm.getDocs'
 require 'vm.getLibrary'
diff --git a/script/vm/local-id.lua b/script/vm/local-id.lua
new file mode 100644
index 00000000..f19e85e6
--- /dev/null
+++ b/script/vm/local-id.lua
@@ -0,0 +1,21 @@
+local util = require 'utility'
+
+---@class parser.object
+---@field _localID string
+
+local compileMap = util.switch()
+    : getMap()
+
+local m = {}
+
+m.ID_SPLITE = '\x1F'
+
+function m.getID(source)
+    local compiler = compileMap[source.type]
+    if compiler then
+        return compiler(source)
+    end
+    return false
+end
+
+return m
diff --git a/script/vm/node/compiler.lua b/script/vm/node/compiler.lua
index 422f464e..296917d9 100644
--- a/script/vm/node/compiler.lua
+++ b/script/vm/node/compiler.lua
@@ -1,20 +1,18 @@
-local guide = require 'parser.guide'
-local util  = require 'utility'
-local state = require 'vm.state'
-local union = require 'vm.node.union'
+local guide    = require 'parser.guide'
+local util     = require 'utility'
+local state    = require 'vm.state'
+local union    = require 'vm.node.union'
+local globalID = require 'vm.global-id'
 
 ---@class parser.object
 ---@field _compiledNodes  boolean
 ---@field _node           vm.node
----@field _globalID       vm.node.global
 
 ---@class vm.node.compiler
 local m = {}
 
 ---@class vm.node.cross
 
-m.GLOBAL_SPLITE = '\x1F'
-
 ---@alias vm.node parser.object | vm.node.union | vm.node.cross
 
 ---@param  ... string
@@ -24,6 +22,9 @@ function m.getGlobalID(...)
 end
 
 function m.setNode(source, node)
+    if not node then
+        return
+    end
     local me = source._node
     if not me then
         source._node = node
@@ -31,23 +32,75 @@ function m.setNode(source, node)
     end
     if me.type == 'union'
     or me.type == 'cross' then
-        me:merge(source, node)
+        me:merge(node)
         return
     end
-    source._node = union(source, node)
+    source._node = union(me, node)
 end
 
-local function compileValue(uri, source, value)
-    if not value then
+function m.eachNode(node)
+    if node.type == 'union' then
+        return node:eachNode()
+    end
+    local first = true
+    return function ()
+        if first then
+            first = false
+            return node
+        end
+        return nil
+    end
+end
+
+local function getReturnOfFunction(func, index)
+    if not func._returns then
+        func._returns = util.defaultTable(function ()
+            return {
+                type   = 'function.return',
+                parent = func,
+                index  = index,
+            }
+        end)
+    end
+    return m.compileNode(guide.getUri(func), func._returns[index])
+end
+
+local function getReturn(func, index)
+    local node = m.compileNode(guide.getUri(func), func)
+    if not node then
         return
     end
-    if value.type == 'table'
-    or value.type == 'integer'
-    or value.type == 'number'
-    or value.type == 'string'
-    or value.type == 'function' then
+    for cnode in m.eachNode(node) do
+        if cnode.type == 'function' then
+            return getReturnOfFunction(cnode, index)
+        end
+    end
+end
+
+local valueMap = util.switch()
+    : case 'boolean'
+    : case 'table'
+    : case 'integer'
+    : case 'number'
+    : case 'string'
+    : case 'function'
+    : call(function (uri, source, value)
         state.declareLiteral(uri, value)
         m.setNode(source, value)
+    end)
+    : case 'call'
+    : call(function (uri, source, value)
+        m.setNode(source, getReturn(value.node, 1))
+    end)
+    : getMap()
+
+local function compileValue(uri, source, value)
+    if not value then
+        return
+    end
+    local f = valueMap[value.type]
+    if f then
+        f(uri, source, value)
     end
 end
 
@@ -67,6 +120,14 @@ local compilerMap = util.switch()
     : call(function (uri, source)
         m.setNode(source, m.compileNode(uri, source.node))
     end)
+    : case 'setfield'
+    : case 'getfield'
+    : call(function (uri, source)
+    end)
+    : case 'function.return'
+    : call(function (uri, source)
+        
+    end)
     : getMap()
 
 ---@param uri    uri
@@ -85,85 +146,4 @@ function m.compileNode(uri, source)
     return source._node
 end
 
-local compilerGlobalMap = util.switch()
-    : case 'local'
-    : call(function (uri, source)
-        if source.tag ~= '_ENV' then
-            return
-        end
-        if source.ref then
-            for _, ref in ipairs(source.ref) do
-                m.compileGlobalNode(uri, ref)
-            end
-        end
-    end)
-    : case 'setglobal'
-    : call(function (uri, source)
-        local name = guide.getKeyName(source)
-        source._globalID = state.declareGlobal(name, uri, source)
-    end)
-    : case 'getglobal'
-    : call(function (uri, source)
-        local name   = guide.getKeyName(source)
-        local global = state.getGlobal(name)
-        global:addGet(uri, source)
-        source._globalID = global
-
-        local nxt = source.next
-        if nxt then
-            m.compileGlobalNode(uri, nxt)
-        end
-    end)
-    : case 'setfield'
-    ---@param uri    uri
-    ---@param source parser.object
-    : call(function (uri, source)
-        local parent = source.node._globalID
-        if not parent then
-            return
-        end
-        local name = m.getGlobalID(parent:getName(), guide.getKeyName(source))
-        source._globalID = state.declareGlobal(name, uri, source)
-    end)
-    : case 'getfield'
-    ---@param uri    uri
-    ---@param source parser.object
-    : call(function (uri, source)
-        local parent = source.node._globalID
-        if not parent then
-            return
-        end
-        local name = m.getGlobalID(parent:getName(), guide.getKeyName(source))
-        local global = state.getGlobal(name)
-        global:addGet(uri, source)
-        source._globalID = global
-
-        local nxt = source.next
-        if nxt then
-            m.compileGlobalNode(uri, nxt)
-        end
-    end)
-    : getMap()
-
----@param uri    uri
----@param source parser.object
-function m.compileGlobalNode(uri, source)
-    if source._globalID ~= nil then
-        return
-    end
-    source._globalID = false
-    local compiler = compilerGlobalMap[source.type]
-    if compiler then
-        compiler(uri, source)
-    end
-end
-
----编译全局变量的node
----@param  root parser.object
-function m.compileGlobals(root)
-    local uri = guide.getUri(root)
-    local env = guide.getENV(root)
-    m.compileGlobalNode(uri, env)
-end
-
 return m
diff --git a/script/vm/node/init.lua b/script/vm/node/init.lua
deleted file mode 100644
index 524dab4e..00000000
--- a/script/vm/node/init.lua
+++ /dev/null
@@ -1,22 +0,0 @@
-local files    = require 'files'
-local compiler = require 'vm.node.compiler'
-local vmState  = require 'vm.state'
-
-for uri in files.eachFile() do
-    local state = files.getState(uri)
-    if state then
-        compiler.compileGlobals(state.ast)
-    end
-end
-
-files.watch(function (ev, uri)
-    if ev == 'update' then
-        local state = files.getState(uri)
-        if state then
-            compiler.compileGlobals(state.ast)
-        end
-    end
-    if ev == 'remove' then
-        vmState.dropUri(uri)
-    end
-end)
diff --git a/script/vm/node/union.lua b/script/vm/node/union.lua
index 928a9216..a8c917d9 100644
--- a/script/vm/node/union.lua
+++ b/script/vm/node/union.lua
@@ -5,9 +5,8 @@ local mt = {}
 mt.__index = mt
 mt.type = 'union'
 
----@param source parser.object
 ---@param node   vm.node
-function mt:merge(source, node)
+function mt:merge(node)
     if not node then
         return
     end
@@ -30,13 +29,21 @@ function mt:subscribeLiteral(source)
     end
 end
 
----@param source parser.object
+function mt:eachNode()
+    local i = 0
+    return function ()
+        i = i + 1
+        return self[i]
+    end
+end
+
+---@param me   parser.object
 ---@param node vm.node
 ---@return vm.node.union
-return function (source, node)
+return function (me, node)
     local union = setmetatable({
-        source,
+        [1] = me,
     }, mt)
-    union:merge(source, node)
+    union:merge(node)
     return union
 end
diff --git a/script/vm/state.lua b/script/vm/state.lua
index 15f4af70..1eddc9db 100644
--- a/script/vm/state.lua
+++ b/script/vm/state.lua
@@ -1,14 +1,9 @@
-local util   = require 'utility'
-local global = require 'vm.node.global'
+local util     = require 'utility'
+local files    = require 'files'
+local globalID = require 'vm.global-id'
 
 ---@class vm.state
 local m = {}
----@type table<string, vm.node.global>
-m.globals = util.defaultTable(global)
----@type table<uri, table<string, boolean>>
-m.globalSubs = util.defaultTable(function ()
-    return {}
-end)
 ---@type table<uri, parser.object[]>
 m.literals = util.multiTable(2)
 ---@type table<parser.object, table<parser.object, boolean>>
@@ -18,27 +13,6 @@ end)
 ---@type table<parser.object, boolean>
 m.allLiterals = {}
 
----@param name   string
----@param uri    uri
----@param source parser.object
----@return vm.node.global
-function m.declareGlobal(name, uri, source)
-    m.globalSubs[uri][name] = true
-    local node = m.globals[name]
-    node:addSet(uri, source)
-    return node
-end
-
----@param name string
----@param uri? uri
----@return vm.node.global
-function m.getGlobal(name, uri)
-    if uri then
-        m.globalSubs[uri][name] = true
-    end
-    return m.globals[name]
-end
-
 ---@param uri    uri
 ---@param source parser.object
 function m.declareLiteral(uri, source)
@@ -69,11 +43,6 @@ end
 
 ---@param uri uri
 function m.dropUri(uri)
-    local globalSub = m.globalSubs[uri]
-    m.globalSubs[uri] = nil
-    for name in pairs(globalSub) do
-        m.globals[name]:dropUri(uri)
-    end
     local literals = m.literals[uri]
     m.literals[uri] = nil
     for _, literal in ipairs(literals) do
@@ -86,4 +55,25 @@ function m.dropUri(uri)
     end
 end
 
+for uri in files.eachFile() do
+    local state = files.getState(uri)
+    if state then
+        globalID.compileAst(state.ast)
+    end
+end
+
+files.watch(function (ev, uri)
+    if ev == 'update' then
+        local state = files.getState(uri)
+        if state then
+            globalID.compileAst(state.ast)
+        end
+    end
+    if ev == 'remove' then
+        m.dropUri(uri)
+        globalID.dropUri(uri)
+    end
+end)
+
+
 return m
diff --git a/test/definition/table.lua b/test/definition/table.lua
index 61e8746d..0fb76191 100644
--- a/test/definition/table.lua
+++ b/test/definition/table.lua
@@ -134,6 +134,22 @@ local y = {
 t.<?insert?>()
 ]]
 
+
+TEST [[
+local x
+x.y.<!z!> = 1
+print(x.y.<?z?>)
+]]
+
+
+TEST [[
+local x
+x.y = {
+    <!z!> = 1
+}
+print(x.y.<?z?>)
+]]
+
 TEST [[
 local function f()
     local t = {}