3 files changed, 32 insertions, 17 deletions
diff --git a/server/src/core/diagnostics.lua b/server/src/core/diagnostics.lua
index d0fdbacd..089ff352 100644
--- a/server/src/core/diagnostics.lua
+++ b/server/src/core/diagnostics.lua
@@ -23,15 +23,18 @@ function mt:searchUnusedLocals(callback)
         if var.hide then
             goto NEXT_VAR
         end
-        for _, info in ipairs(var) do
+        local ok = self.vm:eachInfo(var, function (info)
             if info.type == 'get' then
-                goto NEXT_VAR
+                return true
             end
             if info.type == 'local' then
                 if info.source.start == 0 then
-                    goto NEXT_VAR
+                    return true
                 end
             end
+        end)
+        if ok then
+            goto NEXT_VAR
         end
         callback(var.source.start, var.source.finish, var.key)
         ::NEXT_VAR::
@@ -39,9 +42,8 @@ function mt:searchUnusedLocals(callback)
 end
 
 function mt:searchUndefinedGlobal(callback)
-    local results = self.results
-    local env = results.locals[1]
-    env.value:eachField(function (index, field)
+    local globalValue = self.vm.lsp.globalValue
+    globalValue:eachField(function (index, field)
         if field.value.lib then
             goto NEXT_VAR
         end
@@ -51,20 +53,22 @@ function mt:searchUndefinedGlobal(callback)
         if config.config.diagnostics.globals[index] then
             goto NEXT_VAR
         end
-        local lIndex = index:lower()
-        if lIndex == '' then
+        if index == '' then
             goto NEXT_VAR
         end
-        for _, info in ipairs(field) do
+        local ok = self.vm:eachInfo(field, function (info)
             if info.type == 'set' then
-                goto NEXT_VAR
+                return true
             end
+        end)
+        if ok then
+            goto NEXT_VAR
         end
-        for _, info in ipairs(field) do
+        self.vm:eachInfo(field, function (info)
             if info.type == 'get' then
                 callback(info.source.start, info.source.finish, tostring(index))
             end
-        end
+        end)
         ::NEXT_VAR::
     end)
 end
@@ -72,16 +76,19 @@ end
 function mt:searchUnusedLabel(callback)
     local results = self.results
     for _, label in ipairs(results.labels) do
-        for _, info in ipairs(label) do
+        local ok = self.vm:eachInfo(label, function (info)
             if info.type == 'goto' then
-                goto NEXT_LABEL
+                return true
             end
+        end)
+        if ok then
+            goto NEXT_LABEL
         end
-        for _, info in ipairs(label) do
+        self.vm:eachInfo(label, function (info)
             if info.type == 'set' then
                 callback(info.source.start, info.source.finish, label.key)
             end
-        end
+        end)
         ::NEXT_LABEL::
     end
 end
diff --git a/server/src/core/global.lua b/server/src/core/global.lua
index ac1170a7..a8f5092d 100644
--- a/server/src/core/global.lua
+++ b/server/src/core/global.lua
@@ -8,10 +8,16 @@ function mt:compileVM(uri)
 end
 
 function mt:markSet(uri)
+    if not uri then
+        return
+    end
     self.set[uri] = true
 end
 
 function mt:markGet(uri)
+    if not uri then
+        return
+    end
     self.get[uri] = true
 end
 
diff --git a/server/test/diagnostics/init.lua b/server/test/diagnostics/init.lua
index 8a2dcf61..acaf3249 100644
--- a/server/test/diagnostics/init.lua
+++ b/server/test/diagnostics/init.lua
@@ -1,5 +1,6 @@
 local core = require 'core'
 local parser  = require 'parser'
+local service = require 'service'
 
 rawset(_G, 'TEST', true)
 
@@ -38,10 +39,11 @@ end
 
 function TEST(script)
     local new_script, target = catch_target(script)
+    local lsp = service()
     local ast = parser:ast(new_script)
     assert(ast)
     local lines = parser:lines(new_script)
-    local vm = core.vm(ast)
+    local vm = core.vm(ast, lsp)
     assert(vm)
     local datas = core.diagnostics(vm, lines, 'test')
     local results = {}