supports `---@enum (key)`

11 files changed, 271 insertions, 38 deletions

diff --git a/changelog.md b/changelog.md
index 38e53e95..1c5b53f2 100644
--- a/changelog.md
+++ b/changelog.md
@@ -4,6 +4,24 @@
 * `NEW` support `---@type` and `--[[@as]]` for return statement
 * `NEW` commandline parameter `--force-accept-workspace`: allowing the use of the root directory or home directory as the workspace
 * `NEW` diagnostic: `inject-field`
+* `NEW` `---@enum` supports attribute `key`
+  ```lua
+  ---@enum (key) AnimalType
+  local enum = {
+    Cat = 1,
+    Dog = 2,
+  }
+  
+  ---@param animal userdata
+  ---@param atp AnimalType
+  ---@return boolean
+  local function isAnimalType(animal, atp)
+    return API.isAnimalType(animal, enum[atp])
+  end
+
+  assert(isAnimalType(animal, 'Cat'))
+  ```
+
 * `FIX` wrong hover and signature for method with varargs and overloads
 * `FIX` [#2155]
 * `FIX` [#2224]
diff --git a/script/core/completion/completion.lua b/script/core/completion/completion.lua
index 15bc0fcb..5a61919c 100644
--- a/script/core/completion/completion.lua
+++ b/script/core/completion/completion.lua
@@ -1251,6 +1251,46 @@ local function insertDocEnum(state, pos, doc, enums)
     return enums
 end
 
+---@param state     parser.state
+---@param pos       integer
+---@param doc       vm.node.object
+---@param enums     table[]
+---@return table[]?
+local function insertDocEnumKey(state, pos, doc, enums)
+    local tbl = doc.bindSource
+    if not tbl then
+        return nil
+    end
+    local keyEnums = {}
+    for _, field in ipairs(tbl) do
+        if field.type == 'tablefield'
+        or field.type == 'tableindex' then
+            if not field.value then
+                goto CONTINUE
+            end
+            local key = guide.getKeyName(field)
+            if not key then
+                goto CONTINUE
+            end
+            enums[#enums+1] = {
+                label  = ('%q'):format(key),
+                kind   = define.CompletionItemKind.EnumMember,
+                id     = stack(field, function (newField) ---@async
+                    return {
+                        detail      = buildDetail(newField),
+                        description = buildDesc(newField),
+                    }
+                end),
+            }
+            ::CONTINUE::
+        end
+    end
+    for _, enum in ipairs(keyEnums) do
+        enums[#enums+1] = enum
+    end
+    return enums
+end
+
 local function buildInsertDocFunction(doc)
     local args = {}
     for i, arg in ipairs(doc.args) do
@@ -1316,7 +1356,11 @@ local function insertEnum(state, pos, src, enums, isInArray, mark)
     elseif src.type == 'global' and src.cate == 'type' then
         for _, set in ipairs(src:getSets(state.uri)) do
             if set.type == 'doc.enum' then
-                insertDocEnum(state, pos, set, enums)
+                if vm.docHasAttr(set, 'key') then
+                    insertDocEnumKey(state, pos, set, enums)
+                else
+                    insertDocEnum(state, pos, set, enums)
+                end
             end
         end
     end
diff --git a/script/core/hover/description.lua b/script/core/hover/description.lua
index f5890b21..58fb9fbe 100644
--- a/script/core/hover/description.lua
+++ b/script/core/hover/description.lua
@@ -477,34 +477,55 @@ local function tryDocEnum(source)
     if not tbl then
         return
     end
-    local md = markdown()
-    md:add('lua', '{')
-    for _, field in ipairs(tbl) do
-        if field.type == 'tablefield'
-        or field.type == 'tableindex' then
-            if not field.value then
-                goto CONTINUE
-            end
-            local key = guide.getKeyName(field)
-            if not key then
-                goto CONTINUE
-            end
-            if field.value.type == 'integer'
-            or field.value.type == 'string' then
-                md:add('lua', ('    %s: %s = %s,'):format(key, field.value.type, field.value[1]))
+    if vm.docHasAttr(source, 'key') then
+        local md = markdown()
+        local keys = {}
+        for _, field in ipairs(tbl) do
+            if field.type == 'tablefield'
+            or field.type == 'tableindex' then
+                if not field.value then
+                    goto CONTINUE
+                end
+                local key = guide.getKeyName(field)
+                if not key then
+                    goto CONTINUE
+                end
+                keys[#keys+1] = ('%q'):format(key)
+                ::CONTINUE::
             end
-            if field.value.type == 'binary'
-            or field.value.type == 'unary' then
-                local number = vm.getNumber(field.value)
-                if number then
-                    md:add('lua', ('    %s: %s = %s,'):format(key, math.tointeger(number) and 'integer' or 'number', number))
+        end
+        md:add('lua', table.concat(keys, ' | '))
+        return md:string()
+    else
+        local md = markdown()
+        md:add('lua', '{')
+        for _, field in ipairs(tbl) do
+            if field.type == 'tablefield'
+            or field.type == 'tableindex' then
+                if not field.value then
+                    goto CONTINUE
+                end
+                local key = guide.getKeyName(field)
+                if not key then
+                    goto CONTINUE
+                end
+                if field.value.type == 'integer'
+                or field.value.type == 'string' then
+                    md:add('lua', ('    %s: %s = %s,'):format(key, field.value.type, field.value[1]))
+                end
+                if field.value.type == 'binary'
+                or field.value.type == 'unary' then
+                    local number = vm.getNumber(field.value)
+                    if number then
+                        md:add('lua', ('    %s: %s = %s,'):format(key, math.tointeger(number) and 'integer' or 'number', number))
+                    end
                 end
+                ::CONTINUE::
             end
-            ::CONTINUE::
         end
+        md:add('lua', '}')
+        return md:string()
     end
-    md:add('lua', '}')
-    return md:string()
 end
 
 ---@async
diff --git a/script/core/semantic-tokens.lua b/script/core/semantic-tokens.lua
index cd19e2ee..4e1d8e00 100644
--- a/script/core/semantic-tokens.lua
+++ b/script/core/semantic-tokens.lua
@@ -711,6 +711,14 @@ local Care = util.switch()
             type       = define.TokenTypes.namespace,
         }
     end)
+    : case 'doc.attr'
+    : call(function (source, options, results)
+        results[#results+1] = {
+            start      = source.start,
+            finish     = source.finish,
+            type       = define.TokenTypes.decorator,
+        }
+    end)
 
 ---@param state table
 ---@param results table
diff --git a/script/parser/guide.lua b/script/parser/guide.lua
index dc46fecf..94816be8 100644
--- a/script/parser/guide.lua
+++ b/script/parser/guide.lua
@@ -159,7 +159,7 @@ local childMap = {
     ['doc.class']          = {'class', '#extends', '#signs', 'comment'},
     ['doc.type']           = {'#types', 'name', 'comment'},
     ['doc.alias']          = {'alias', 'extends', 'comment'},
-    ['doc.enum']           = {'enum', 'extends', 'comment'},
+    ['doc.enum']           = {'enum', 'extends', 'comment', 'docAttr'},
     ['doc.param']          = {'param', 'extends', 'comment'},
     ['doc.return']         = {'#returns', 'comment'},
     ['doc.field']          = {'field', 'extends', 'comment'},
@@ -182,6 +182,7 @@ local childMap = {
     ['doc.cast.block']     = {'extends'},
     ['doc.operator']       = {'op', 'exp', 'extends'},
     ['doc.meta']           = {'name'},
+    ['doc.attr']           = {'#names'},
 }
 
 ---@type table<string, fun(obj: parser.object, list: parser.object[])>
diff --git a/script/parser/luadoc.lua b/script/parser/luadoc.lua
index 3454b241..1365adfc 100644
--- a/script/parser/luadoc.lua
+++ b/script/parser/luadoc.lua
@@ -156,6 +156,7 @@ Symbol              <-  ({} {
 ---@field calls?            parser.object[]
 ---@field generics?         parser.object[]
 ---@field generic?          parser.object
+---@field docAttr?          parser.object
 
 local function parseTokens(text, offset)
     Ci = 0
@@ -252,6 +253,40 @@ local function nextSymbolOrError(symbol)
     return false
 end
 
+local function parseDocAttr(parent)
+    if not checkToken('symbol', '(', 1) then
+        return nil
+    end
+    nextToken()
+
+    local attrs = {
+        type   = 'doc.attr',
+        parent = parent,
+        start  = getStart(),
+        finish = getStart(),
+        names  = {},
+    }
+
+    while true do
+        if checkToken('symbol', ',', 1) then
+            nextToken()
+            goto continue
+        end
+        local name = parseName('doc.attr.name', attrs)
+        if not name then
+            break
+        end
+        attrs.names[#attrs.names+1] = name
+        attrs.finish = name.finish
+        ::continue::
+    end
+
+    nextSymbolOrError(')')
+    attrs.finish = getFinish()
+
+    return attrs
+end
+
 local function parseIndexField(parent)
     if not checkToken('symbol', '[', 1) then
         return nil
@@ -1428,17 +1463,22 @@ local docSwitch = util.switch()
     end)
     : case 'enum'
     : call(function ()
+        local attr = parseDocAttr()
         local name = parseName('doc.enum.name')
         if not name then
             return nil
         end
         local result = {
-            type   = 'doc.enum',
-            start  = name.start,
-            finish = name.finish,
-            enum   = name,
+            type    = 'doc.enum',
+            start   = name.start,
+            finish  = name.finish,
+            enum    = name,
+            docAttr = attr,
         }
         name.parent = result
+        if attr then
+            attr.parent = result
+        end
         return result
     end)
     : case 'private'
diff --git a/script/vm/doc.lua b/script/vm/doc.lua
index 5cb039fe..6ac39910 100644
--- a/script/vm/doc.lua
+++ b/script/vm/doc.lua
@@ -471,3 +471,18 @@ function vm.getCastTargetHead(doc)
     end
     return nil
 end
+
+---@param doc parser.object
+---@param key string
+---@return boolean
+function vm.docHasAttr(doc, key)
+    if not doc.docAttr then
+        return false
+    end
+    for _, name in ipairs(doc.docAttr.names) do
+        if name[1] == key then
+            return true
+        end
+    end
+    return false
+end
diff --git a/script/vm/global.lua b/script/vm/global.lua
index 4e2d0617..e830f6d8 100644
--- a/script/vm/global.lua
+++ b/script/vm/global.lua
@@ -372,16 +372,36 @@ local compilerGlobalSwitch = util.switch()
             return
         end
         source._enums = {}
-        for _, field in ipairs(tbl) do
-            if     field.type == 'tablefield' then
-                source._enums[#source._enums+1] = field
-                local subType = vm.declareGlobal('type', name .. '.' .. field.field[1], uri)
-                subType:addSet(uri, field)
-            elseif field.type == 'tableindex' then
-                source._enums[#source._enums+1] = field
-                if field.index.type == 'string' then
-                    local subType = vm.declareGlobal('type', name .. '.' .. field.index[1], uri)
+        if vm.docHasAttr(source, 'key') then
+            for _, field in ipairs(tbl) do
+                if     field.type == 'tablefield' then
+                    source._enums[#source._enums+1] = {
+                        type   = 'doc.type.string',
+                        start  = field.field.start,
+                        finish = field.field.finish,
+                        [1]    = field.field[1],
+                    }
+                elseif field.type == 'tableindex' then
+                    source._enums[#source._enums+1] = {
+                        type   = 'doc.type.string',
+                        start  = field.index.start,
+                        finish = field.index.finish,
+                        [1]    = field.index[1],
+                    }
+                end
+            end
+        else
+            for _, field in ipairs(tbl) do
+                if     field.type == 'tablefield' then
+                    source._enums[#source._enums+1] = field
+                    local subType = vm.declareGlobal('type', name .. '.' .. field.field[1], uri)
                     subType:addSet(uri, field)
+                elseif field.type == 'tableindex' then
+                    source._enums[#source._enums+1] = field
+                    if field.index.type == 'string' then
+                        local subType = vm.declareGlobal('type', name .. '.' .. field.index[1], uri)
+                        subType:addSet(uri, field)
+                    end
                 end
             end
         end
diff --git a/test/completion/common.lua b/test/completion/common.lua
index e33b2a1b..bd317259 100644
--- a/test/completion/common.lua
+++ b/test/completion/common.lua
@@ -3816,6 +3816,32 @@ f(<??>)
 }
 
 TEST [[
+local x = 1
+local y = 2
+
+---@enum(key) Enum
+local t = {
+    x = x,
+    y = y,
+}
+
+---@param p Enum
+local function f(p) end
+
+f(<??>)
+]]
+{
+    {
+        label    = '"x"',
+        kind     = define.CompletionItemKind.EnumMember,
+    },
+    {
+        label    = '"y"',
+        kind     = define.CompletionItemKind.EnumMember,
+    },
+}
+
+TEST [[
 --
 <??>
 ]]
diff --git a/test/crossfile/hover.lua b/test/crossfile/hover.lua
index 617b6f7e..a18c714b 100644
--- a/test/crossfile/hover.lua
+++ b/test/crossfile/hover.lua
@@ -1637,6 +1637,30 @@ TEST {
     {
         path = 'a.lua',
         content = [[
+            ---@enum(key) <?A?>
+            local t = {
+                x = 1 << 0,
+                y = 1 << 1,
+                z = 1 << 2,
+            }
+        ]]
+    },
+    hover = [[
+```lua
+(enum) A
+```
+
+---
+
+```lua
+"x" | "y" | "z"
+```]]
+}
+
+TEST {
+    {
+        path = 'a.lua',
+        content = [[
             ---@alias someType
             ---| "#" # description
             
diff --git a/test/diagnostics/param-type-mismatch.lua b/test/diagnostics/param-type-mismatch.lua
index 2a4bfcb1..e31e9933 100644
--- a/test/diagnostics/param-type-mismatch.lua
+++ b/test/diagnostics/param-type-mismatch.lua
@@ -91,6 +91,22 @@ f(<!{ h = 1 }!>)
 ]]
 
 TEST [[
+---@enum(key) A
+local t = {
+    x = 1,
+    ['y'] = 2,
+}
+
+---@param x A
+local function f(x)
+end
+
+f('x')
+f('y')
+f(<!'z'!>)
+]]
+
+TEST [[
 ---@generic T: string | boolean | table
 ---@param x T
 ---@return T