fix: keep declaration docstrings when an attribute defines a constant

src/tests/VersoTests/LiterateHtml.lean

@@ -357,6 +357,23 @@ private def testDottedNameDocstrings (data : TestData) : Test := withTestDir dat
     | parts =>
       fail s!"Expected the docstring \"{docstring}\" once in the HTML, found it {parts.length - 1} times."
 
+/--
+A declaration keeps its docstring when one of its attributes defines a further constant, and a doc
+comment inside an attribute belongs to the constant defined there, including in nested attributes.
+-/
+private def testAttrDocstrings (data : TestData) : Test := withTestDir data fun jsonDir htmlDir _ _ => do
+  runLiterateHtml jsonDir htmlDir
+  let json ← IO.FS.readFile (jsonDir / "LitConfig" / "Core" / "AttrDocstrings.json")
+  for (decl, docstring) in [
+      ("first", "The docstring of the first declaration."),
+      ("second", "The docstring of the second declaration."),
+      ("secondAlias", "The docstring of the constant defined in the second attribute."),
+      ("third", "The docstring of the third declaration."),
+      ("thirdOuter", "The docstring of the outer constant."),
+      ("thirdInner", "The docstring of the inner constant.")] do
+    assertContains s!"\"AttrDocstrings.{decl}\",\{\"blocks\":[\{\"p\":[\{\"normal\":\"{docstring}\"" json
+      s!"The docstring \"{docstring}\" is not associated with `AttrDocstrings.{decl}`."
+
 /-- Excluded modules produce no HTML output and are absent from the navbar. -/
 private def testExclude (data : TestData) : Test := withTestDir data fun jsonDir htmlDir planFile tomlFile => do
   IO.FS.writeFile tomlFile "exclude = [\"LitConfig.NoDocstrings\"]\n"
@@ -1262,14 +1279,15 @@ private def htmlTests (data : TestData) (projectDir : System.FilePath) : List (S
   ("docstring code block messages", testDocstringCodeBlockMessages data),
   ("guard_msgs docstring", testGuardMsgsDocstring data),
   ("dotted name docstrings", testDottedNameDocstrings data),
+  ("attribute docstrings", testAttrDocstrings data),
   ("unknown extension fallback", testUnknownExtensionFallback)
 ]
 
 /-- The literate HTML generator produces the expected output for the single-root test project. -/
 @[test]
 def literateHtml : Test := do
   let projectDir := "test-projects/literate-config"
-  let modules := #["LitConfig", "LitConfig.Core", "LitConfig.Core.Basic", "LitConfig.NoDocstrings", "LitConfig.Builtins", "LitConfig.UserExt", "LitConfig.GuardMsgs", "LitConfig.DottedNames"]
+  let modules := #["LitConfig", "LitConfig.Core", "LitConfig.Core.Basic", "LitConfig.NoDocstrings", "LitConfig.Builtins", "LitConfig.UserExt", "LitConfig.GuardMsgs", "LitConfig.DottedNames", "LitConfig.Core.AttrDefinedConstant", "LitConfig.Core.AttrDocstrings"]
 
   -- First verify test project toolchain matches root toolchain
   let rootToolchain := (← IO.FS.readFile "lean-toolchain").trimAscii

src/verso-literate/VersoLiterateMain.lean

@@ -164,6 +164,11 @@ partial def findDocstringDefs (stx : Syntax) (t : InfoTree) : TermElabM Syntax :
     | _ => none
   let defSites ← defSites.filterM fun x => do
     return (← findInternalDocString? (← getEnv) x.1) |>.isSome
+  -- An attribute can define a further constant, as in `@[to_additive foo /-- ... -/] def bar`. The
+  -- declaration's modifiers hold `foo`'s definition site and `bar`'s doc comment, while `bar`'s
+  -- definition site lies outside them, so attributes are handled first and the definition sites
+  -- that they claim take no further part.
+  let (stx, defSites) ← (rewriteAttrComments stx).run defSites
   -- Now find the largest syntax object that contains just one of these definition sites and replace
   -- all doc comments with a placeholder token
   let stx ← stx.replaceM fun s => do
@@ -196,6 +201,21 @@ partial def findDocstringDefs (stx : Syntax) (t : InfoTree) : TermElabM Syntax :
     else pure none
   return stx
 where
+  -- Doc comments in an attribute instance belong to the only definition site in it that remains after
+  -- the nested instances have been handled, and the instance then claims that site. The state is the
+  -- definition sites not yet claimed.
+  rewriteAttrComments (stx : Syntax) : StateT (List (Name × Syntax)) TermElabM Syntax := do
+    let .node info kind args := stx | return stx
+    let new := Syntax.node info kind (← args.mapM rewriteAttrComments)
+    if kind != ``Lean.Parser.Term.attrInstance then return new
+    let some range := stx.getRange? | return new
+    let inside := (← get).filter fun (_, site) =>
+      site.getRange?.map (range.includes · true true) |>.getD false
+    let [(x, _)] := inside | return new
+    modify (·.filter (·.1 != x))
+    new.replaceM fun s => do
+      if s.isOfKind ``docComment then some <$> rewriteComment x s else pure none
+
   -- For `structure A.B.C`, the binder's syntax is `C` with the source range of `A.B.C`, so the
   -- identifiers are also compared by their source range.
   identMatch (stx1 stx2 : Syntax) : Bool :=

test-projects/literate-config/LitConfig.lean

@@ -5,6 +5,7 @@ import LitConfig.UserExt
 import LitConfig.Gallery
 import LitConfig.GuardMsgs
 import LitConfig.DottedNames
+import LitConfig.Core.AttrDocstrings
 import Verso
 
 /-!

test-projects/literate-config/LitConfig/Core/AttrDefinedConstant.lean

@@ -0,0 +1,44 @@
+import Lean.AddDecl
+import Lean.Attributes
+import Lean.DocString
+import Lean.Elab.Exception
+import Lean.Elab.InfoTree.Main
+import Lean.Parser.Term
+
+open Lean Elab
+
+set_option doc.verso false
+
+/-!
+# Attributes That Define Constants
+
+An attribute that defines a further constant whose name is written in the attribute, the way that
+Mathlib's `to_additive` does. The new constant takes its docstring from a doc comment in the attribute
+if there is one, and otherwise from the declaration. A nested attribute is applied to the new constant.
+-/
+
+/-- Defines a constant with the same type and value under the given name. -/
+syntax (name := alsoAs) "also_as " ident (" (" &"attr" " := " Parser.Term.attrInstance ")")? (docComment)? : attr
+
+initialize registerBuiltinAttribute {
+  name := `alsoAs
+  descr := "Defines a constant with the same type and value under the name in the attribute"
+  applicationTime := .afterCompilation
+  add := fun decl stx _ => do
+    let `(attr| also_as $id:ident $[(attr := $nested)]? $[$doc:docComment]?) := stx
+      | throwUnsupportedSyntax
+    let info ← getConstInfo decl
+    let name := id.getId
+    addDecl <| .defnDecl <|
+      mkDefinitionValEx name info.levelParams info.type
+        (.const decl (info.levelParams.map .param)) .abbrev .safe [decl]
+    if let some doc := doc then
+      addDocStringCore name (← getDocStringText doc)
+    else if let some doc ← findDocString? (← getEnv) decl then
+      addDocStringCore name doc
+    pushInfoLeaf <| .ofTermInfo {
+      elaborator := .anonymous, lctx := {}, expectedType? := none, stx := id,
+      expr := .const name [], isBinder := true }
+    if let some nested := nested then
+      Attribute.add name nested.raw[1].getKind nested.raw[1]
+}

test-projects/literate-config/LitConfig/Core/AttrDocstrings.lean

@@ -0,0 +1,23 @@
+import LitConfig.Core.AttrDefinedConstant
+
+set_option doc.verso false
+
+/-!
+# Docstrings and Attributes
+
+A declaration keeps its docstring when one of its attributes defines a further constant.
+-/
+
+/-- The docstring of the first declaration. -/
+@[also_as AttrDocstrings.firstAlias]
+def AttrDocstrings.first : Nat := 1
+
+/-- The docstring of the second declaration. -/
+@[also_as AttrDocstrings.secondAlias /-- The docstring of the constant defined in the second attribute. -/]
+def AttrDocstrings.second : Nat := 2
+
+/-- The docstring of the third declaration. -/
+@[also_as AttrDocstrings.thirdOuter
+    (attr := also_as AttrDocstrings.thirdInner /-- The docstring of the inner constant. -/)
+    /-- The docstring of the outer constant. -/]
+def AttrDocstrings.third : Nat := 3