6 files changed, 78 insertions, 30 deletions
diff --git a/script/core/diagnostics/undefined-doc-name.lua b/script/core/diagnostics/undefined-doc-name.lua
index c7318fcf..91d4b90e 100644
--- a/script/core/diagnostics/undefined-doc-name.lua
+++ b/script/core/diagnostics/undefined-doc-name.lua
@@ -1,7 +1,6 @@
 local files   = require 'files'
 local guide   = require 'parser.guide'
 local lang    = require 'language'
-local define  = require 'proto.define'
 local vm      = require 'vm'
 
 return function (uri, callback)
@@ -14,26 +13,6 @@ return function (uri, callback)
         return
     end
 
-    local classCache = {
-        ['any'] = true,
-        ['nil'] = true,
-    }
-    local function hasNameOfClassOrAlias(name)
-        if classCache[name] ~= nil then
-            return classCache[name]
-        end
-        local docs = vm.getDocDefines(name)
-        for _, otherDoc in ipairs(docs) do
-            if otherDoc.type == 'doc.class.name'
-            or otherDoc.type == 'doc.alias.name' then
-                classCache[name] = true
-                return true
-            end
-        end
-        classCache[name] = false
-        return false
-    end
-
     local function hasNameOfGeneric(name, source)
         if not source.typeGeneric then
             return false
@@ -56,7 +35,7 @@ return function (uri, callback)
         if name == '...' then
             return
         end
-        if hasNameOfClassOrAlias(name)
+        if vm.isDocDefined(name)
         or hasNameOfGeneric(name, source) then
             return
         end
diff --git a/script/proto/define.lua b/script/proto/define.lua
index a49a0ac1..b5e40c60 100644
--- a/script/proto/define.lua
+++ b/script/proto/define.lua
@@ -284,4 +284,19 @@ m.BuiltIn = {
     ['utf8']      = 'default',
 }
 
+m.BuiltinClass = {
+    ['unknown']       = true,
+    ['any']           = true,
+    ['nil']           = true,
+    ['boolean']       = true,
+    ['number']        = true,
+    ['integer']       = true,
+    ['thread']        = true,
+    ['table']         = true,
+    ['string']        = true,
+    ['userdata']      = true,
+    ['lightuserdata'] = true,
+    ['Function']      = true,
+}
+
 return m
diff --git a/script/vm/docs.lua b/script/vm/docs.lua
index 95cb4127..df87d46a 100644
--- a/script/vm/docs.lua
+++ b/script/vm/docs.lua
@@ -1,5 +1,4 @@
 local files = require 'files'
-local await = require 'await'
 local noder = require 'core.noder'
 
 local docsMap      = {}
@@ -25,7 +24,8 @@ local function pushDocs(uri)
     end
     local nodes = noder.compileNodes(state.ast)
     for id in pairs(nodes) do
-        if id:sub(1, 3) == 'dn:' then
+        if id:sub(1, 3) == 'dn:'
+        and noder.getFirstID(id) == id then
             if not docsMap[id] then
                 docsMap[id] = {}
             end
diff --git a/script/vm/getDocs.lua b/script/vm/getDocs.lua
index 0f12c763..bf99a247 100644
--- a/script/vm/getDocs.lua
+++ b/script/vm/getDocs.lua
@@ -3,7 +3,9 @@ local guide    = require 'parser.guide'
 ---@type vm
 local vm       = require 'vm.vm'
 local config   = require 'config'
-local searcher = require 'core.searcher'
+local docs     = require 'vm.docs'
+local define   = require 'proto.define'
+local noder    = require 'core.noder'
 
 local function getDocDefinesInAst(results, root, name)
     for _, doc in ipairs(root.docs) do
@@ -24,15 +26,61 @@ end
 ---@return parser.guide.object[]
 function vm.getDocDefines(name)
     local results = {}
-    for uri in files.eachFile() do
-        local state = files.getState(uri)
-        if state then
-            getDocDefinesInAst(results, state.ast, name)
+    if name then
+        local id = 'dn:' .. name
+        local uris = docs.getUrisByID(id)
+        if not uris then
+            return results
+        end
+        for uri in pairs(uris) do
+            local state = files.getState(uri)
+            if state then
+                getDocDefinesInAst(results, state.ast, name)
+            end
+        end
+    else
+        for uri in files.eachFile() do
+            local state = files.getState(uri)
+            if state then
+                getDocDefinesInAst(results, state.ast, name)
+            end
         end
     end
     return results
 end
 
+function vm.isDocDefined(name)
+    if define.BuiltinClass[name] then
+        return true
+    end
+    local cache = vm.getCache 'isDocDefined'
+    if cache[name] ~= nil then
+        return cache[name]
+    end
+    local id = 'dn:' .. name
+    local uris = docs.getUrisByID(id)
+    if not uris then
+        cache[name] = false
+        return
+    end
+    for uri in pairs(uris) do
+        local state = files.getState(uri)
+        local node = noder.getNodeByID(state.ast, id)
+        if node and node.sources then
+            for _, source in ipairs(node.sources) do
+                local doc = source.parent
+                if doc.type == 'doc.class'
+                or doc.type == 'doc.alias' then
+                    cache[name] = true
+                    return true
+                end
+            end
+        end
+    end
+    cache[name] = false
+    return false
+end
+
 function vm.getDocEnums(doc)
     if not doc then
         return nil
diff --git a/script/vm/globals.lua b/script/vm/globals.lua
index cb759551..3c74f8b0 100644
--- a/script/vm/globals.lua
+++ b/script/vm/globals.lua
@@ -1,5 +1,4 @@
 local files = require 'files'
-local await = require 'await'
 local noder = require 'core.noder'
 
 local globalsMap    = {}
diff --git a/script/vm/vm.lua b/script/vm/vm.lua
index ebd0102b..0e7f3176 100644
--- a/script/vm/vm.lua
+++ b/script/vm/vm.lua
@@ -59,6 +59,13 @@ function m.getArgInfo(source)
     return nil
 end
 
+function m.getSpecial(source)
+    if not source then
+        return nil
+    end
+    return source.special
+end
+
 function m.getKeyName(source)
     if not source then
         return nil