7 files changed, 151 insertions, 16 deletions
diff --git a/server/src/core/highlight.lua b/server/src/core/highlight.lua
new file mode 100644
index 00000000..c71aab8b
--- /dev/null
+++ b/server/src/core/highlight.lua
@@ -0,0 +1,54 @@
+local findSource = require 'core.find_source'
+local parser = require 'parser'
+
+local DocumentHighlightKind = {
+    Text = 1,
+    Read = 2,
+    Write = 3,
+}
+
+local function parseResult(source)
+    local positions = {}
+    if source:bindLabel() then
+        source:bindLabel():eachInfo(function (info, src)
+            positions[#positions+1] = { src.start, src.finish, DocumentHighlightKind.Text }
+        end)
+        return positions
+    end
+    if source:bindLocal() then
+        local loc = source:bindLocal()
+        local mark = {}
+        loc:eachInfo(function (info, src)
+            if not mark[src] then
+                mark[src] = info
+                positions[#positions+1] = { src.start, src.finish, DocumentHighlightKind.Text }
+            end
+        end)
+        return positions
+    end
+    if source:bindValue() then
+        local parent = source:get 'parent'
+        local mark = {}
+        parent:eachInfo(function (info, src)
+            if not mark[src] then
+                mark[src] = info
+                if info.type == 'get child' or info.type == 'set child' then
+                    if info[1] == source[1] then
+                        positions[#positions+1] = {src.start, src.finish, DocumentHighlightKind.Text}
+                    end
+                end
+            end
+        end)
+        return positions
+    end
+    return nil
+end
+
+return function (vm, pos)
+    local source = findSource(vm, pos)
+    if not source then
+        return nil
+    end
+    local positions = parseResult(source)
+    return positions
+end
diff --git a/server/src/core/init.lua b/server/src/core/init.lua
index 1a556320..9e3995a7 100644
--- a/server/src/core/init.lua
+++ b/server/src/core/init.lua
@@ -11,6 +11,7 @@ local api = {
     signature      = require 'core.signature',
     documentSymbol = require 'core.document_symbol',
     global         = require 'core.global',
+    highlight      = require 'core.highlight',
 }
 
 return api
diff --git a/server/src/method/initialize.lua b/server/src/method/initialize.lua
index 538a8f9c..60e8ee1e 100644
--- a/server/src/method/initialize.lua
+++ b/server/src/method/initialize.lua
@@ -17,7 +17,7 @@ return function (lsp)
             referencesProvider = true,
             renameProvider = true,
             documentSymbolProvider = true,
-            --documentHighlightProvider = true,
+            documentHighlightProvider = true,
             signatureHelpProvider = {
                 triggerCharacters = { '(', ',' },
             },
diff --git a/server/src/method/textDocument/documentHighlight.lua b/server/src/method/textDocument/documentHighlight.lua
index 95ee0cfe..377ffcdf 100644
--- a/server/src/method/textDocument/documentHighlight.lua
+++ b/server/src/method/textDocument/documentHighlight.lua
@@ -2,24 +2,22 @@ local core = require 'core'
 
 return function (lsp, params)
     local uri = params.textDocument.uri
-    local newName = params.newName
     local vm, lines = lsp:loadVM(uri)
     if not vm then
-        return {}
+        return nil
     end
     local position = lines:positionAsChar(params.position.line + 1, params.position.character)
-    local positions = core.rename(vm, position, newName)
+    local positions = core.highlight(vm, position)
     if not positions then
-        return {}
+        return nil
     end
 
-    local TextEdit = {}
+    local result = {}
     for i, position in ipairs(positions) do
         local start, finish = position[1], position[2]
         local start_row,  start_col  = lines:rowcol(start)
         local finish_row, finish_col = lines:rowcol(finish)
-        TextEdit[i] = {
-            newText = newName,
+        result[i] = {
             range = {
                 start = {
                     line = start_row - 1,
@@ -30,15 +28,10 @@ return function (lsp, params)
                     -- 这里不用-1，因为前端期待的是匹配完成后的位置
                     character = finish_col,
                 },
-            }
+            },
+            kind = position[3],
         }
     end
 
-    local response = {
-        changes = {
-            [uri] = TextEdit,
-        },
-    }
-
-    return response
+    return result
 end
diff --git a/server/src/service.lua b/server/src/service.lua
index 41630d12..83f4aa0c 100644
--- a/server/src/service.lua
+++ b/server/src/service.lua
@@ -610,6 +610,7 @@ function mt:restartDueToMemoryLeak()
 end
 
 function mt:_testMemory()
+    collectgarbage()
     local cachedVM = 0
     local cachedSource = 0
     local cachedFunction = 0
diff --git a/server/test/highlight/init.lua b/server/test/highlight/init.lua
new file mode 100644
index 00000000..4e68de0c
--- /dev/null
+++ b/server/test/highlight/init.lua
@@ -0,0 +1,85 @@
+local core = require 'core'
+local parser  = require 'parser'
+local buildVM = require 'vm'
+
+local function catch_target(script)
+    local list = {}
+    local cur = 1
+    while true do
+        local start, finish  = script:find('<[!?].-[!?]>', cur)
+        if not start then
+            break
+        end
+        list[#list+1] = { start + 2, finish - 2 }
+        cur = finish + 1
+    end
+    return list
+end
+
+local function founded(targets, results)
+    if #targets ~= #results then
+        return false
+    end
+    for _, target in ipairs(targets) do
+        for _, result in ipairs(results) do
+            if target[1] == result[1] and target[2] == result[2] then
+                goto NEXT
+            end
+        end
+        do return false end
+        ::NEXT::
+    end
+    return true
+end
+
+function TEST(newName)
+    return function (script)
+        local target = catch_target(script)
+        local start  = script:find('<?', 1, true)
+        local finish = script:find('?>', 1, true)
+        local pos = (start + finish) // 2 + 1
+        local new_script = script:gsub('<[!?]', '  '):gsub('[!?]>', '  ')
+        local ast = parser:ast(new_script)
+        assert(ast)
+        local vm = buildVM(ast)
+        assert(vm)
+
+        local positions = core.highlight(vm, pos)
+        if positions then
+            assert(founded(target, positions))
+        else
+            assert(#target == 0)
+        end
+    end
+end
+
+TEST 'b' [[
+local <?a?> = 1
+]]
+
+TEST 'b' [[
+local <?a?> = 1
+<!a!> = 2
+<!a!> = <!a!>
+]]
+
+TEST 'b' [[
+t.<?a?> = 1
+a = t.<!a!>
+]]
+
+TEST 'b' [[
+t[<!'a'!>] = 1
+a = t.<?a?>
+]]
+
+TEST 'b' [[
+:: <?a?> ::
+goto <!a!>
+]]
+
+TEST 'b' [[
+local function f(<!a!>)
+    return <?a?>
+end
+]]
diff --git a/server/test/main.lua b/server/test/main.lua
index 07ea9ebb..10d5e155 100644
--- a/server/test/main.lua
+++ b/server/test/main.lua
@@ -27,6 +27,7 @@ local function main()
     test 'core'
     test 'definition'
     test 'rename'
+    test 'highlight'
     test 'references'
     test 'diagnostics'
     test 'type_inference'