fix completion

3 files changed, 67 insertions, 37 deletions

diff --git a/script/core/completion.lua b/script/core/completion.lua
index 8ba01c24..60dbea7c 100644
--- a/script/core/completion.lua
+++ b/script/core/completion.lua
@@ -634,10 +634,19 @@ local function isInString(ast, offset)
     end)
 end
 
-local function checkKeyWord(ast, text, start, word, hasSpace, afterLocal, results)
+local function checkKeyWord(ast, text, start, offset, word, hasSpace, afterLocal, results)
     local snipType = config.config.completion.keywordSnippet
     local symbol = lookBackward.findSymbol(text, start - 1)
     local isExp = symbol == '(' or symbol == ',' or symbol == '='
+    local info = {
+        hasSpace = hasSpace,
+        isExp    = isExp,
+        text     = text,
+        start    = start,
+        uri      = guide.getUri(ast.ast),
+        offset   = offset,
+        ast      = ast,
+    }
     for _, data in ipairs(keyWordMap) do
         local key = data[1]
         local eq
@@ -665,7 +674,7 @@ local function checkKeyWord(ast, text, start, word, hasSpace, afterLocal, result
         if snipType == 'Both' or snipType == 'Replace' then
             local func = data[2]
             if func then
-                replaced = func(hasSpace, isExp, results, text, start, guide.getUri(ast.ast))
+                replaced = func(info, results)
                 extra = true
             end
         end
@@ -687,7 +696,7 @@ local function checkKeyWord(ast, text, start, word, hasSpace, afterLocal, result
         end
         local checkStop = data[3]
         if checkStop then
-            local stop = checkStop(ast, start)
+            local stop = checkStop(info)
             if stop then
                 return true
             end
@@ -1135,7 +1144,7 @@ local function tryWord(ast, text, offset, results)
             checkFunctionArgByDocParam(ast, word, start, results)
         else
             local afterLocal = isAfterLocal(text, start)
-            local stop = checkKeyWord(ast, text, start, word, hasSpace, afterLocal, results)
+            local stop = checkKeyWord(ast, text, start, offset, word, hasSpace, afterLocal, results)
             if stop then
                 return
             end
diff --git a/script/core/keyword.lua b/script/core/keyword.lua
index 3ce9e045..349f624c 100644
--- a/script/core/keyword.lua
+++ b/script/core/keyword.lua
@@ -3,8 +3,8 @@ local guide      = require 'core.guide'
 local files      = require 'files'
 
 local keyWordMap = {
-    {'do', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'do', function (info, results)
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'do .. end',
                 kind  = define.CompletionItemKind.Snippet,
@@ -23,13 +23,13 @@ end",
             }
         end
         return true
-    end, function (ast, start)
-        return guide.eachSourceContain(ast.ast, start, function (source)
+    end, function (info)
+        return guide.eachSourceContain(info.ast.ast, info.start, function (source)
             if source.type == 'while'
             or source.type == 'in'
             or source.type == 'loop' then
                 for i = 1, #source.keyword do
-                    if start == source.keyword[i] then
+                    if info.start == source.keyword[i] then
                         return true
                     end
                 end
@@ -39,8 +39,12 @@ end",
     {'and'},
     {'break'},
     {'else'},
-    {'elseif', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'elseif', function (info, results)
+        if info.text:find('^%s*then', info.offset + 1)
+        or info.text:find('^%s*do', info.offset + 1) then
+            return false
+        end
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'elseif .. then',
                 kind  = define.CompletionItemKind.Snippet,
@@ -59,8 +63,8 @@ end",
     end},
     {'end'},
     {'false'},
-    {'for', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'for', function (info, results)
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'for .. ipairs',
                 kind  = define.CompletionItemKind.Snippet,
@@ -119,13 +123,13 @@ end"
         end
         return true
     end},
-    {'function', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'function', function (info, results)
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'function ()',
                 kind  = define.CompletionItemKind.Snippet,
                 insertTextFormat = 2,
-                insertText = isExp and "\z
+                insertText = info.isExp and "\z
 ($1)\
 \t$0\
 end" or "\z
@@ -138,7 +142,7 @@ end"
                 label = 'function ()',
                 kind  = define.CompletionItemKind.Snippet,
                 insertTextFormat = 2,
-                insertText = isExp and "\z
+                insertText = info.isExp and "\z
 function ($1)\
 \t$0\
 end" or "\z
@@ -150,8 +154,12 @@ end"
         return true
     end},
     {'goto'},
-    {'if', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'if', function (info, results)
+        if info.text:find('^%s*then', info.offset + 1)
+        or info.text:find('^%s*do', info.offset + 1) then
+            return false
+        end
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'if .. then',
                 kind  = define.CompletionItemKind.Snippet,
@@ -174,8 +182,12 @@ end"
         end
         return true
     end},
-    {'in', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'in', function (info, results)
+        if info.text:find('^%s*then', info.offset + 1)
+        or info.text:find('^%s*do', info.offset + 1) then
+            return false
+        end
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'in ..',
                 kind  = define.CompletionItemKind.Snippet,
@@ -198,8 +210,8 @@ end"
         end
         return true
     end},
-    {'local', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'local', function (info, results)
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'local function',
                 kind  = define.CompletionItemKind.Snippet,
@@ -225,8 +237,8 @@ end"
     {'nil'},
     {'not'},
     {'or'},
-    {'repeat', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'repeat', function (info, results)
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'repeat .. until',
                 kind  = define.CompletionItemKind.Snippet,
@@ -246,8 +258,8 @@ until $1"
         end
         return true
     end},
-    {'return', function (hasSpace, isExp, results)
-        if not hasSpace then
+    {'return', function (info, results)
+        if not info.hasSpace then
             results[#results+1] = {
                 label = 'do return end',
                 kind  = define.CompletionItemKind.Snippet,
@@ -257,22 +269,21 @@ until $1"
         end
         return false
     end},
-    ---@param text string
-    {'then', function (hasSpace, isExp, results, text, start, uri)
-        local lines = files.getLines(uri)
-        local pos, first = text:match('%S+%s+()(%S+)', start)
+    {'then', function (info, results)
+        local lines = files.getLines(info.uri)
+        local pos, first = info.text:match('%S+%s+()(%S+)', info.start)
         if first == 'end'
         or first == 'else'
         or first == 'elseif' then
-            local startRow  = guide.positionOf(lines, start)
+            local startRow  = guide.positionOf(lines, info.start)
             local finishRow = guide.positionOf(lines, pos)
-            local startSp   = text:match('^%s*', lines[startRow].start)
-            local finishSp  = text:match('^%s*', lines[finishRow].start)
+            local startSp   = info.text:match('^%s*', lines[startRow].start)
+            local finishSp  = info.text:match('^%s*', lines[finishRow].start)
             if startSp == finishSp then
                 return false
             end
         end
-        if not hasSpace then
+        if not info.hasSpace then
             results[#results+1] = {
                 label = 'then .. end',
                 kind  = define.CompletionItemKind.Snippet,
@@ -287,8 +298,8 @@ end'
     end},
     {'true'},
     {'until'},
-    {'while', function (hasSpace, isExp, results)
-        if hasSpace then
+    {'while', function (info, results)
+        if info.hasSpace then
             results[#results+1] = {
                 label = 'while .. do',
                 kind  = define.CompletionItemKind.Snippet,
diff --git a/test/completion/init.lua b/test/completion/init.lua
index abd5b9a0..bf7a896f 100644
--- a/test/completion/init.lua
+++ b/test/completion/init.lua
@@ -2459,3 +2459,13 @@ local t = {
         kind  = define.CompletionItemKind.Property,
     },
 }
+
+TEST [[
+if $ then
+]]
+(nil)
+
+TEST [[
+elseif $ then
+]]
+(nil)