fix enum completion

2 files changed, 50 insertions, 29 deletions

diff --git a/script/core/completion.lua b/script/core/completion.lua
index 285adb38..f7f1e2c7 100644
--- a/script/core/completion.lua
+++ b/script/core/completion.lua
@@ -1140,39 +1140,29 @@ local function checkEqualEnum(ast, text, offset, results)
 end
 
 local function checkEqualEnumInString(ast, text, offset, results)
-    local list = {}
-    guide.eachSourceContain(ast.ast, offset, function (source)
-        if source.type == 'binary' then
-            if source.op.type == '=='
-            or source.op.type == '~=' then
-                list[#list+1] = source[1]
-            end
+    local source = findNearestSource(ast, offset)
+    local parent = source.parent
+    if parent.type == 'binary' then
+        if source ~= parent[2] then
+            return
         end
-        if not source.start then
+        if not parent.op then
             return
         end
-        if  source.start <= offset
-        and source.finish >= offset then
-            local parent = source.parent
-            if not parent then
-                return
-            end
-            if parent.type == 'local' then
-                list[#list+1] = parent
-            end
-            if parent.type == 'setlocal'
-            or parent.type == 'setglobal'
-            or parent.type == 'setfield'
-            or parent.type == 'setindex' then
-                list[#list+1] = parent.node
-            end
+        if parent.op.type ~= '==' and parent.op.type ~= '~=' then
+            return
         end
-    end)
-    table.sort(list, function (a, b)
-        return a.start > b.start
-    end)
-    local source = list[1]
-    checkEqualEnumLeft(ast, text, offset, source, results)
+        checkEqualEnumLeft(ast, text, offset, parent[1], results)
+    end
+    if parent.type == 'local' then
+        checkEqualEnumLeft(ast, text, offset, parent, results)
+    end
+    if parent.type == 'setlocal'
+    or parent.type == 'setglobal'
+    or parent.type == 'setfield'
+    or parent.type == 'setindex' then
+        checkEqualEnumLeft(ast, text, offset, parent.node, results)
+    end
 end
 
 local function isFuncArg(ast, offset)
diff --git a/test/completion/init.lua b/test/completion/init.lua
index 9c87ab33..ea1cdfdb 100644
--- a/test/completion/init.lua
+++ b/test/completion/init.lua
@@ -2665,3 +2665,34 @@ t.$
         },
     }
 }
+
+TEST [[
+---@alias enum '"aaa"'|'"bbb"'
+
+---@param x enum
+---@return enum
+local function f(x)
+end
+
+local r = f('$')
+]]
+{
+    {
+        label = "'aaa'",
+        kind  = define.CompletionItemKind.EnumMember,
+        textEdit = {
+            newText = "'aaa'",
+            start   = 103,
+            finish  = 104,
+        },
+    },
+    {
+        label = "'bbb'",
+        kind  = define.CompletionItemKind.EnumMember,
+        textEdit = {
+            newText = "'bbb'",
+            start   = 103,
+            finish  = 104,
+        },
+    },
+}