From c2018e05de0e9eebadbe094357d22883a608fdf5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9C=80=E8=90=8C=E5=B0=8F=E6=B1=90?= <sumneko@hotmail.com>
Date: Wed, 23 Aug 2023 15:13:51 +0800
Subject: [PATCH] support `---@class (exact)`

#1990
---
 changelog.md                             | 10 +++
 script/core/diagnostics/inject-field.lua | 81 ++++++++++++++++++++----
 script/parser/luadoc.lua                 |  1 +
 test/diagnostics/inject-field.lua        | 20 ++++++
 4 files changed, 99 insertions(+), 13 deletions(-)

diff --git a/changelog.md b/changelog.md
index a105d2769..474625a80 100644
--- a/changelog.md
+++ b/changelog.md
@@ -21,6 +21,16 @@
 
   assert(isAnimalType(animal, 'Cat'))
   ```
+* `NEW` `---@class` supports attribute `exact`
+  ```lua
+  ---@class (exact) Point
+  ---@field x number
+  ---@field y number
+  local m = {}
+  m.x = 1 -- OK
+  m.y = 2 -- OK
+  m.z = 3 -- Warning
+  ```
 
 * `FIX` wrong hover and signature for method with varargs and overloads
 * `FIX` [#2155]
diff --git a/script/core/diagnostics/inject-field.lua b/script/core/diagnostics/inject-field.lua
index 570ca270a..2866eef8f 100644
--- a/script/core/diagnostics/inject-field.lua
+++ b/script/core/diagnostics/inject-field.lua
@@ -37,14 +37,29 @@ return function (uri, callback)
             return
         end
 
+        local isExact
         local class = vm.getDefinedClass(uri, node)
         if class then
-            return
+            for _, doc in ipairs(class:getSets(uri)) do
+                if vm.docHasAttr(doc, 'exact') then
+                    isExact = true
+                    break
+                end
+            end
+            if not isExact then
+                return
+            end
+            if  src.type == 'setmethod'
+            and not guide.getSelfNode(node) then
+                return
+            end
         end
 
         for _, def in ipairs(vm.getDefs(src)) do
             local dnode = def.node
-            if dnode and vm.getDefinedClass(uri, dnode) then
+            if dnode
+            and not isExact
+            and vm.getDefinedClass(uri, dnode) then
                 return
             end
             if def.type == 'doc.type.field' then
@@ -55,16 +70,19 @@ return function (uri, callback)
             end
         end
 
-        local howToFix = lang.script('DIAG_INJECT_FIELD_FIX_CLASS', {
-            node = hname(node),
-            fix  = '---@class',
-        })
-        for _, ndef in ipairs(vm.getDefs(node)) do
-            if ndef.type == 'doc.type.table' then
-                howToFix = lang.script('DIAG_INJECT_FIELD_FIX_TABLE', {
-                    fix   = '[any]: any',
-                })
-                break
+        local howToFix = ''
+        if not isExact then
+            howToFix = lang.script('DIAG_INJECT_FIELD_FIX_CLASS', {
+                node = hname(node),
+                fix  = '---@class',
+            })
+            for _, ndef in ipairs(vm.getDefs(node)) do
+                if ndef.type == 'doc.type.table' then
+                    howToFix = lang.script('DIAG_INJECT_FIELD_FIX_TABLE', {
+                        fix   = '[any]: any',
+                    })
+                    break
+                end
             end
         end
 
@@ -79,7 +97,7 @@ return function (uri, callback)
                 finish  = src.field.finish,
                 message = message,
             }
-        elseif src.type == 'setfield' and src.method then
+        elseif src.type == 'setmethod' and src.method then
             callback {
                 start   = src.method.start,
                 finish  = src.method.finish,
@@ -89,4 +107,41 @@ return function (uri, callback)
     end
     guide.eachSourceType(ast.ast, 'setfield',  checkInjectField)
     guide.eachSourceType(ast.ast, 'setmethod', checkInjectField)
+
+    ---@async
+    local function checkExtraTableField(src)
+        await.delay()
+
+        if not src.bindSource then
+            return
+        end
+        if not vm.docHasAttr(src, 'exact') then
+            return
+        end
+        local value = src.bindSource.value
+        if not value or value.type ~= 'table' then
+            return
+        end
+        for _, field in ipairs(value) do
+            local defs = vm.getDefs(field)
+            for _, def in ipairs(defs) do
+                if def.type == 'doc.field' then
+                    goto nextField
+                end
+            end
+            local message = lang.script('DIAG_INJECT_FIELD', {
+                class = vm.getInfer(src):view(uri),
+                field = guide.getKeyName(src),
+                fix   = '',
+            })
+            callback {
+                start   = field.start,
+                finish  = field.finish,
+                message = message,
+            }
+            ::nextField::
+        end
+    end
+
+    guide.eachSourceType(ast.ast, 'doc.class', checkExtraTableField)
 end
diff --git a/script/parser/luadoc.lua b/script/parser/luadoc.lua
index 818696381..d7338918f 100644
--- a/script/parser/luadoc.lua
+++ b/script/parser/luadoc.lua
@@ -841,6 +841,7 @@ local docSwitch = util.switch()
             operators = {},
             calls     = {},
         }
+        result.docAttr = parseDocAttr(result)
         result.class = parseName('doc.class.name', result)
         if not result.class then
             pushWarning {
diff --git a/test/diagnostics/inject-field.lua b/test/diagnostics/inject-field.lua
index f4d847e9f..9bb0f8fcf 100644
--- a/test/diagnostics/inject-field.lua
+++ b/test/diagnostics/inject-field.lua
@@ -62,3 +62,23 @@ local t
 t.x = 1 -- OK
 t.y = 2 -- OK
 ]]
+
+
+TEST [[
+---@class (exact) Class
+---@field x number
+local m = {
+    x = 1, -- OK
+    <!y!> = 2, -- Warning
+}
+
+m.x = 1 -- OK
+m.<!y!> = 2 -- Warning
+
+function m:init() -- OK
+    self.x = 1 -- OK
+    self.<!y!> = 2 -- Warning
+    function self:<!xx!>() -- Warning
+    end
+end
+]]