Embed stdlib documentation to AST. (#2851)

* Embed stdlib documentation to AST.

* Extract documentation for attributes.

---------

Co-authored-by: Yong He <yhe@nvidia.com>

15 files changed, 89 insertions, 17 deletions

diff --git a/source/compiler-core/slang-doc-extractor.cpp b/source/compiler-core/slang-doc-extractor.cpp
index ca947c94f..ffbfc9904 100644
--- a/source/compiler-core/slang-doc-extractor.cpp
+++ b/source/compiler-core/slang-doc-extractor.cpp
@@ -638,6 +638,12 @@ SlangResult DocMarkupExtractor::_findMarkup(const FindInfo& info, SearchStyle se
         {
             return _findMarkup(info, Location::Before, out);
         }
+        case SearchStyle::Attribute:
+        {
+            FindInfo newInfo = info;
+            newInfo.tokenIndex -= 2;
+            return _findMarkup(newInfo, Location::Before, out);
+        }
         case SearchStyle::Variable:
         {
             Location locs[] = { Location::Before, Location::AfterSemicolon };
diff --git a/source/compiler-core/slang-doc-extractor.h b/source/compiler-core/slang-doc-extractor.h
index fc486250e..30346ce6f 100644
--- a/source/compiler-core/slang-doc-extractor.h
+++ b/source/compiler-core/slang-doc-extractor.h
@@ -99,6 +99,7 @@ public:
         Before,             ///< Only allows before
         Function,           ///< Function/method
         GenericParam,       ///< Generic parameter
+        Attribute,          ///< Attribute definition
     };
 
         /// An input search item
diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index 764e39cce..6a35f496b 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -112,12 +112,13 @@ interface __BuiltinSignedArithmeticType : __BuiltinArithmeticType {}
 interface __BuiltinIntegerType : __BuiltinArithmeticType
 {}
 
-/// Modifer to mark a function for forward-mode differentiation.
+/// Marks a function for forward-mode differentiation.
 /// i.e. the compiler will automatically generate a new function
 /// that computes the jacobian-vector product of the original.
 __attributeTarget(FunctionDeclBase)
 attribute_syntax [ForwardDifferentiable] : ForwardDifferentiableAttribute;
 
+/// Marks a function for backward-mode differentiation.
 __attributeTarget(FunctionDeclBase)
 attribute_syntax [BackwardDifferentiable] : BackwardDifferentiableAttribute;
 
diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang
index f8b36a3ac..f2f1d0cc3 100644
--- a/source/slang/diff.meta.slang
+++ b/source/slang/diff.meta.slang
@@ -25,6 +25,7 @@ attribute_syntax [DerivativeMember(memberName)] : DerivativeMemberAttribute;
 __attributeTarget(FunctionDeclBase)
 attribute_syntax [NoDiffThis] : NoDiffThisAttribute;
 
+/// Represents a GPU view of a tensor.
 __generic<T>
 __magic_type(TensorViewType)
 __intrinsic_type($(kIROp_TensorViewType))
@@ -231,6 +232,7 @@ extension TensorView<float>
     void InterlockedCompareExchange(vector<uint, N> index, float compare, float val);
 }
 
+/// Represents the handle of a Torch tensor object.
 __generic<T>
 __intrinsic_type($(kIROp_TorchTensorType))
 struct TorchTensor
@@ -294,10 +296,12 @@ struct TorchTensor
 __target_intrinsic(cpp, "AT_CUDA_CHECK(cudaStreamSynchronize(at::cuda::getCurrentCUDAStream()))")
 void syncTorchCudaStream();
 
+/// Constructs a `DifferentialPair` value from a primal value and a differential value.
 __generic<T: IDifferentiable>
 __intrinsic_op($(kIROp_MakeDifferentialPairUserCode))
 DifferentialPair<T> diffPair(T primal, T.Differential diff);
 
+/// Constructs a `DifferentialPair` value from a primal value and a zero differential value.
 __generic<T: IDifferentiable>
 [__unsafeForceInlineEarly]
 DifferentialPair<T> diffPair(T primal)
@@ -812,7 +816,7 @@ void __d_cross(inout DifferentialPair<vector<T, 3>> a, inout DifferentialPair<ve
     }
 #define SIMPLE_UNARY_DERIVATIVE_IMPL(NAME, DIFF_FUNC) UNARY_DERIVATIVE_IMPL(NAME, ReturnType.dmul(DIFF_FUNC, dpx.d), ReturnType.dmul(DIFF_FUNC, dOut))
 
-// Detach and set derivatives to zero
+/// Detach and set derivatives to zero.
 __generic<T : IDifferentiable>
 __intrinsic_op($(kIROp_DetachDerivative))
 T detach(T x);
diff --git a/source/slang/slang-ast-base.h b/source/slang/slang-ast-base.h
index 5e3773189..6e43b378a 100644
--- a/source/slang/slang-ast-base.h
+++ b/source/slang/slang-ast-base.h
@@ -407,6 +407,8 @@ public:
 
     NameLoc nameAndLoc;
 
+    RefPtr<MarkupEntry> markup;
+
     Name*     getName()       { return nameAndLoc.name; }
     SourceLoc getNameLoc()    { return nameAndLoc.loc ; }
     NameLoc   getNameAndLoc() { return nameAndLoc     ; }
diff --git a/source/slang/slang-ast-dump.cpp b/source/slang/slang-ast-dump.cpp
index 723608bb2..884f8b736 100644
--- a/source/slang/slang-ast-dump.cpp
+++ b/source/slang/slang-ast-dump.cpp
@@ -351,6 +351,10 @@ struct ASTDumpContext
     {
         m_writer->emit((int)kind);
     }
+    void dump(MarkupVisibility v)
+    {
+        m_writer->emit((int)v);
+    }
     void dump(const String& string)
     {
         dump(string.getUnownedSlice());
diff --git a/source/slang/slang-ast-modifier.h b/source/slang/slang-ast-modifier.h
index 6f79e900f..1083afae4 100644
--- a/source/slang/slang-ast-modifier.h
+++ b/source/slang/slang-ast-modifier.h
@@ -577,7 +577,9 @@ class AttributeTargetModifier : public Modifier
 class AttributeBase : public Modifier 
 {
     SLANG_AST_CLASS(AttributeBase)
- 
+    
+    AttributeDecl* attributeDecl = nullptr;
+
     List<Expr*> args;
 };
 
diff --git a/source/slang/slang-ast-print.cpp b/source/slang/slang-ast-print.cpp
index 23430cd13..da35a2818 100644
--- a/source/slang/slang-ast-print.cpp
+++ b/source/slang/slang-ast-print.cpp
@@ -366,6 +366,8 @@ void ASTPrinter::addDeclKindPrefix(Decl* decl)
                     continue;
                 if (as<SpecializedForTargetModifier>(modifier))
                     continue;
+                if (as<AttributeTargetModifier>(modifier))
+                    continue;
             }
             // Don't print out attributes.
             if (as<AttributeBase>(modifier))
@@ -430,6 +432,10 @@ void ASTPrinter::addDeclKindPrefix(Decl* decl)
     {
         m_builder << "associatedtype ";
     }
+    else if (const auto attribute = as<AttributeDecl>(decl))
+    {
+        m_builder << "attribute ";
+    }
 }
 
 void ASTPrinter::addDeclResultType(const DeclRef<Decl>& inDeclRef)
diff --git a/source/slang/slang-ast-support-types.h b/source/slang/slang-ast-support-types.h
index 443b6e7bd..3ab9d1a01 100644
--- a/source/slang/slang-ast-support-types.h
+++ b/source/slang/slang-ast-support-types.h
@@ -5,6 +5,7 @@
 
 #include "../compiler-core/slang-lexer.h"
 #include "../compiler-core/slang-name.h"
+#include "../compiler-core/slang-doc-extractor.h"
 
 #include "slang-profile.h"
 #include "slang-type-system-shared.h"
@@ -19,7 +20,6 @@
 #include "slang-ast-reflect.h"
 #include "slang-ref-object-reflect.h"
 
-
 #include <assert.h>
 
 namespace Slang
@@ -1555,6 +1555,16 @@ namespace Slang
         Backward
     };
 
+        /// Represents a markup (documentation) associated with a decl.
+    struct MarkupEntry : public SerialRefObject
+    {
+        SLANG_OBJ_CLASS(MarkupEntry)
+
+        NodeBase* m_node;                                           ///< The node this documentation is associated with
+        String m_markup;                                            ///< The raw contents of of markup associated with the decoration
+        MarkupVisibility m_visibility = MarkupVisibility::Public;   ///< How visible this decl is
+    };
+
     /// Get the inner most expr from an higher order expr chain, e.g. `__fwd_diff(__fwd_diff(f))`'s
     /// inner most expr is `f`.
     Expr* getInnerMostExprFromHigherOrderExpr(Expr* expr, FunctionDifferentiableLevel& outDiffLevel);
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index 78a9fd163..61f3f1196 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -780,6 +780,7 @@ namespace Slang
         attr->keywordName  = uncheckedAttr->keywordName;
         attr->args  = uncheckedAttr->args;
         attr->loc   = uncheckedAttr->loc;
+        attr->attributeDecl = attrDecl;
 
         // We will start with checking steps that can be applied independent
         // of the concrete attribute type that was selected. These only need
diff --git a/source/slang/slang-doc-ast.cpp b/source/slang/slang-doc-ast.cpp
index 54155f06c..4f0d310bf 100644
--- a/source/slang/slang-doc-ast.cpp
+++ b/source/slang/slang-doc-ast.cpp
@@ -36,6 +36,10 @@ namespace Slang {
     {
         return SearchStyle::GenericParam;
     }
+    else if (as<AttributeDecl>(decl))
+    {
+        return SearchStyle::Attribute;
+    }
     else
     {
         // If can't determine just allow before
@@ -53,6 +57,8 @@ static void _addDeclRec(Decl* decl, List<Decl*>& outDecls)
     // If we don't have a loc, we have no way of locating documentation.
     if (decl->loc.isValid() || decl->nameAndLoc.loc.isValid())
     {
+        if (as<AttributeDecl>(decl))
+            printf("dd");
         outDecls.add(decl);
     }
 
diff --git a/source/slang/slang-doc-ast.h b/source/slang/slang-doc-ast.h
index df8f201f7..25c39d9d3 100644
--- a/source/slang/slang-doc-ast.h
+++ b/source/slang/slang-doc-ast.h
@@ -14,12 +14,8 @@ namespace Slang {
 class ASTMarkup : public RefObject
 {
 public:
-    struct Entry
-    {
-        NodeBase* m_node;                                           ///< The node this documentation is associated with
-        String m_markup;                                            ///< The raw contents of of markup associated with the decoration
-        MarkupVisibility m_visibility = MarkupVisibility::Public;   ///< How visible this decl is
-    };
+
+    typedef MarkupEntry Entry;
 
         /// Adds an entry, returns the reference to pre-existing node if there is one
     Entry& addEntry(NodeBase* base);
@@ -29,6 +25,9 @@ public:
         /// Get list of all of the entries in source order
     const List<Entry>& getEntries() const { return m_entries; }
 
+        /// Attaches the markup to the AST nodes.
+    void attachToAST();
+
 protected:
 
         /// Map from AST nodes to documentation entries
@@ -59,6 +58,17 @@ SLANG_INLINE ASTMarkup::Entry* ASTMarkup::getEntry(NodeBase* base)
     return (indexPtr) ? &m_entries[*indexPtr] : nullptr;
 }
 
+SLANG_INLINE void ASTMarkup::attachToAST()
+{
+    for (auto& entry : m_entries)
+    {
+        if (auto decl = as<Decl>(entry.m_node))
+        {
+            decl->markup = new MarkupEntry(entry);
+        }
+    }
+}
+
 /* Extracts documentation markup from source. 
 The comments are extracted and associated in declarations. The association
 is held in DocMarkup type. The comment style follows the doxygen style */
diff --git a/source/slang/slang-language-server-ast-lookup.cpp b/source/slang/slang-language-server-ast-lookup.cpp
index a9d7bf6d4..433047741 100644
--- a/source/slang/slang-language-server-ast-lookup.cpp
+++ b/source/slang/slang-language-server-ast-lookup.cpp
@@ -653,6 +653,12 @@ bool _findAstNodeImpl(ASTLookupContext& context, SyntaxNode* node)
                     context.results.add(result);
                     return true;
                 }
+                for (auto arg : attribute->args)
+                {
+                    ASTLookupExprVisitor exprVisitor(&context);
+                    if (exprVisitor.dispatchIfNotNull(arg))
+                        return true;
+                }
             }
         }
         if (auto container = as<ContainerDecl>(node))
diff --git a/source/slang/slang-language-server.cpp b/source/slang/slang-language-server.cpp
index b826d7cfe..dbe031fd3 100644
--- a/source/slang/slang-language-server.cpp
+++ b/source/slang/slang-language-server.cpp
@@ -230,10 +230,6 @@ String getDeclKindString(DeclRef<Decl> declRef)
     {
         return "(generic value parameter) ";
     }
-    else if (declRef.as<AttributeDecl>())
-    {
-        return "(attribute) ";
-    }
     else if (auto varDecl = declRef.as<VarDeclBase>())
     {
         auto parent = declRef.getParent();
@@ -422,8 +418,16 @@ static void _tryGetDocumentation(StringBuilder& sb, WorkspaceVersion* workspace,
     auto definingModule = getModuleDecl(decl);
     if (definingModule)
     {
-        auto markupAST = workspace->getOrCreateMarkupAST(definingModule);
-        auto markupEntry = markupAST->getEntry(decl);
+        MarkupEntry* markupEntry = nullptr;
+        if (decl->markup)
+        {
+            markupEntry = decl->markup;
+        }
+        else
+        {
+            auto markupAST = workspace->getOrCreateMarkupAST(definingModule);
+            markupEntry = markupAST->getEntry(decl);
+        }
         if (markupEntry)
         {
             sb << "\n";
@@ -705,6 +709,10 @@ SlangResult LanguageServer::hover(
     {
         fillDeclRefHoverInfo(DeclRef<Decl>(decl, nullptr));
     }
+    else if (auto attr = as<Attribute>(leafNode))
+    {
+        fillDeclRefHoverInfo(DeclRef<Decl>(attr->attributeDecl, nullptr));
+    }
     if (sb.getLength() == 0)
     {
         m_connection->sendResult(NullResponse::get(), responseId);
@@ -951,7 +959,7 @@ SlangResult LanguageServer::completion(
         return SLANG_OK;
     }
 
-    // Don't general completion suggestions after typing '['.
+    // Don't generate completion suggestions after typing '['.
     if (args.context.triggerKind ==
         LanguageServerProtocol::kCompletionTriggerKindTriggerCharacter &&
         args.context.triggerCharacter == "[")
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 1dc61b937..8aea1bc37 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -4543,6 +4543,11 @@ void Session::addBuiltinSource(
     auto module = compileRequest->translationUnits[translationUnitIndex]->getModule();
     auto moduleDecl = module->getModuleDecl();
 
+    // Extact documentation markup.
+    ASTMarkup markup;
+    ASTMarkupUtil::extract(moduleDecl, sourceManager, &sink, &markup);
+    markup.attachToAST();
+
     // Put in the loaded module map
     linkage->mapNameToLoadedModules.add(moduleName, module);