update

3 files changed, 60 insertions, 27 deletions

diff --git a/script/core/highlight.lua b/script/core/highlight.lua
index d02f9004..d07290f6 100644
--- a/script/core/highlight.lua
+++ b/script/core/highlight.lua
@@ -7,9 +7,24 @@ local guide      = require 'parser.guide'
 
 ---@async
 local function eachRef(source, callback)
-    local results = vm.getRefs(source)
-    for i = 1, #results do
-        callback(results[i])
+    local uri = guide.getUri(source)
+    local mark = {}
+    local refs = vm.getRefs(source, function ()
+        return false
+    end)
+    for _, ref in ipairs(refs) do
+        if not mark[ref] then
+            mark[ref] = true
+            callback(ref)
+        end
+    end
+    local defs = vm.getDefs(source)
+    for _, def in ipairs(defs) do
+        if  not mark[def]
+        and guide.getUri(def) == uri then
+            mark[def] = true
+            callback(def)
+        end
     end
 end
 
diff --git a/script/vm/def.lua b/script/vm/def.lua
index 98d41fbd..1c9e1118 100644
--- a/script/vm/def.lua
+++ b/script/vm/def.lua
@@ -149,6 +149,12 @@ local nodeSwitch = util.switch()
             searchFieldSwitch(pn.type, pn, key, pushResult)
         end
     end)
+    : case 'tableindex'
+    : case 'tablefield'
+    : call(function (source, pushResult)
+        local tbl = source.parent
+        searchFieldSwitch(tbl.type, tbl, guide.getKeyName(source), pushResult)
+    end)
     : case 'doc.see.field'
     : call(function (source, pushResult)
         local parentNode = compiler.compileNode(source.parent.name)
diff --git a/script/vm/ref.lua b/script/vm/ref.lua
index 3f9262be..7849fda9 100644
--- a/script/vm/ref.lua
+++ b/script/vm/ref.lua
@@ -82,7 +82,23 @@ simpleSwitch = util.switch()
     end)
 
 ---@async
-local function searchField(source, pushResult)
+local function searchInAllFiles(suri, searcher, notify)
+    searcher(suri)
+
+    for uri in files.eachFile(suri) do
+        if  not vm.isMetaFile(uri)
+        and suri ~= uri then
+            local continue = notify(uri)
+            if continue == false then
+                break
+            end
+            searcher(uri)
+        end
+    end
+end
+
+---@async
+local function searchField(source, pushResult, fileNotify)
     local key = guide.getKeyName(source)
 
     ---@param src parser.object
@@ -133,15 +149,11 @@ local function searchField(source, pushResult)
         end)
     end
 
-    for uri in files.eachFile(guide.getUri(source)) do
-        if not vm.isMetaFile(uri) then
-            findWord(uri)
-        end
-    end
+    searchInAllFiles(guide.getUri(source), findWord, fileNotify)
 end
 
 ---@async
-local function searchFunction(source, pushResult)
+local function searchFunction(source, pushResult, fileNotify)
     ---@param src parser.object
     local function checkDef(src)
         for _, def in ipairs(vm.getDefs(src)) do
@@ -165,19 +177,15 @@ local function searchFunction(source, pushResult)
         end)
     end
 
-    for uri in files.eachFile(guide.getUri(source)) do
-        if not vm.isMetaFile(uri) then
-            findCall(uri)
-        end
-    end
+    searchInAllFiles(guide.getUri(source), findCall, fileNotify)
 end
 
 local searchByParentNode
 local nodeSwitch = util.switch()
     : case 'field'
     : case 'method'
-    : call(function (source, pushResult)
-        searchByParentNode(source.parent, pushResult)
+    : call(function (source, pushResult, fileNotify)
+        searchByParentNode(source.parent, pushResult, fileNotify)
     end)
     : case 'getfield'
     : case 'setfield'
@@ -186,7 +194,7 @@ local nodeSwitch = util.switch()
     : case 'getindex'
     : case 'setindex'
     ---@async
-    : call(function (source, pushResult)
+    : call(function (source, pushResult, fileNotify)
         local key = guide.getKeyName(source)
         if type(key) ~= 'string' then
             return
@@ -197,19 +205,19 @@ local nodeSwitch = util.switch()
             return
         end
 
-        searchField(source, pushResult)
+        searchField(source, pushResult, fileNotify)
     end)
     : case 'tablefield'
     : case 'tableindex'
     ---@async
-    : call(function (source, pushResult)
-        searchField(source, pushResult)
+    : call(function (source, pushResult, fileNotify)
+        searchField(source, pushResult, fileNotify)
     end)
     : case 'function'
     : case 'doc.type.function'
     ---@async
-    : call(function (source, pushResult)
-        searchFunction(source, pushResult)
+    : call(function (source, pushResult, fileNotify)
+        searchFunction(source, pushResult, fileNotify)
     end)
 
 ---@param source  parser.object
@@ -232,10 +240,12 @@ local function searchByLocalID(source, pushResult)
     end
 end
 
+---@async
 ---@param source  parser.object
 ---@param pushResult fun(src: parser.object)
-function searchByParentNode(source, pushResult)
-    nodeSwitch(source.type, source, pushResult)
+---@param fileNotify fun(uri: uri): boolean
+function searchByParentNode(source, pushResult, fileNotify)
+    nodeSwitch(source.type, source, pushResult, fileNotify)
 end
 
 local function searchByNode(source, pushResult)
@@ -253,7 +263,9 @@ local function searchByNode(source, pushResult)
 end
 
 ---@async
-function vm.getRefs(source)
+---@param source parser.object
+---@param fileNotify fun(uri: uri): boolean
+function vm.getRefs(source, fileNotify)
     local results = {}
     local mark    = {}
 
@@ -273,8 +285,8 @@ function vm.getRefs(source)
 
     searchBySimple(source, pushResult)
     searchByLocalID(source, pushResult)
-    searchByParentNode(source, pushResult)
     searchByNode(source, pushResult)
+    searchByParentNode(source, pushResult, fileNotify or await.delay)
 
     return results
 end