[mlir:PDLL] Add better support for providing Constraint/Pattern/Rewrite documentation

[rriddle]

Summary

This commit enables providing long-form documentation more seamlessly to the LSP
by revamping decl documentation. For ODS imported constructs, we now also import
descriptions and attach them to decls when possible. For PDLL constructs, the LSP will
now try to provide documentation by parsing the comments directly above the decls
location within the source file. This commit also adds a new parser flag
enableDocumentation that gates the import and attachment of ODS documentation,
which is unnecessary in the normal build process (i.e. it should only be used/consumed
by tools).

[jpienaar]

Nice

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

151

OOC should we check for docstring?

[rriddle]

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

151

I think for simplicity it's good to just use whatever comments precede the decl, which should cover any form of //, ///, etc.

mlir/include/mlir/TableGen/Attribute.h

public:
   StringRef getAttrDefName() const;
 
   // Returns the code body for derived attribute. Aborts if this is not a
   // derived attribute.
   StringRef getDerivedCodeBody() const;
 
   // Returns the dialect for the attribute if defined.
   Dialect getDialect() const;
-
-  // Returns the description of the attribute.
-  StringRef getDescription() const;
 };
 
 // Wrapper class providing helper methods for accessing MLIR constant attribute
 // defined in TableGen. This class should closely reflect what is defined as
 // class `ConstantAttr` in TableGen.
 class ConstantAttr {
 public:
   explicit ConstantAttr(const llvm::DefInit *init);

mlir/include/mlir/TableGen/Constraint.h

public:
   Pred getPredicate() const;
 
   // Returns the condition template that can be used to check if a type or
   // attribute satisfies this constraint.  The template may contain "{0}" that
   // must be substituted with an expression returning an mlir::Type or
   // mlir::Attribute.
   std::string getConditionTemplate() const;
 
-  // Returns the user-readable description of this constraint. If the
-  // description is not provided, returns the TableGen def name.
+  // Returns the user-readable summary of this constraint. If the summary is not
+  // provided, returns the TableGen def name.
   StringRef getSummary() const;
 
+  // Returns the long-form description of this constraint. If the description is
+  // not provided, returns an empty string.
+  StringRef getDescription() const;
+
   /// Returns the name of the TablGen def of this constraint. In some cases
   /// where the current def is anonymous, the name of the base def is used (e.g.
   /// `Optional<>`/`Variadic<>` type constraints).
   StringRef getDefName() const;
 
   /// Returns a unique name for the TablGen def of this constraint. This is
   /// generally just the name of the def, but in some cases where the current
   /// def is anonymous, the name of the base def is attached (to provide more

mlir/include/mlir/TableGen/Type.h

public:
   std::string getCPPClassName() const;
 };
 
 // Wrapper class with helper methods for accessing Types defined in TableGen.
 class Type : public TypeConstraint {
 public:
   explicit Type(const llvm::Record *record);
 
-  // Returns the description of the type.
-  StringRef getDescription() const;
-
   // Returns the dialect for the type if defined.
   Dialect getDialect() const;
 };
 
 } // namespace tblgen
 } // namespace mlir
 
 #endif // MLIR_TABLEGEN_TYPE_H_

mlir/include/mlir/Tools/PDLL/AST/Nodes.h

 class Decl : public Node {
 public:
   /// Return the name of the decl, or nullptr if it doesn't have one.
   const Name *getName() const { return name; }
 
   /// Provide type casting support.
   static bool classof(const Node *node);
 
+  /// Set the documentation comment for this decl.
+  void setDocComment(Context &ctx, StringRef comment);
+
+  /// Return the documentation comment attached to this decl if it has been set.
+  /// Otherwise, returns None.
+  Optional<StringRef> getDocComment() const { return docComment; }
+
 protected:
   Decl(TypeID typeID, SMRange loc, const Name *name = nullptr)
       : Node(typeID, loc), name(name) {}
 
 private:
   /// The name of the decl. This is optional for some decls, such as
   /// PatternDecl.
   const Name *name;
+
+  /// The documentation comment attached to this decl. Defaults to None if
+  /// the comment is unset/unknown.
+  Optional<StringRef> docComment;
 };
 
 //===----------------------------------------------------------------------===//
 // ConstraintDecl
 //===----------------------------------------------------------------------===//
 
 /// This class represents the base of all AST Constraint decls. Constraints
 /// apply matcher conditions to, and define the type of PDLL variables.

mlir/include/mlir/Tools/PDLL/Parser/Parser.h

 namespace pdll {
 class CodeCompleteContext;
 
 namespace ast {
 class Context;
 class Module;
 } // namespace ast
 
-/// Parse an AST module from the main file of the given source manager. An
-/// optional code completion context may be provided to receive code completion
-/// suggestions. If a completion is hit, this method returns a failure.
+/// Parse an AST module from the main file of the given source manager.
+/// `enableDocumentation` is an optional flag that, when set, indicates that the
+/// parser should also include documentation when building the AST when
+/// possible. `codeCompleteContext` is an optional code completion context that
+/// may be provided to receive code completion suggestions. If a completion is
+/// hit, this method returns a failure.
 FailureOr<ast::Module *>
-parsePDLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
+parsePDLLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
+             bool enableDocumentation = false,
              CodeCompleteContext *codeCompleteContext = nullptr);
 } // namespace pdll
 } // namespace mlir
 
 #endif // MLIR_TOOLS_PDLL_PARSER_PARSER_H_

mlir/lib/TableGen/Attribute.cpp

Dialect Attribute::getDialect() const {
   const llvm::RecordVal *record = def->getValue("dialect");
   if (record && record->getValue()) {
     if (DefInit *init = dyn_cast<DefInit>(record->getValue()))
       return Dialect(init->getDef());
   }
   return Dialect(nullptr);
 }
 
-StringRef Attribute::getDescription() const {
-  return def->getValueAsString("description");
-}
-
 ConstantAttr::ConstantAttr(const DefInit *init) : def(init->getDef()) {
   assert(def->isSubClassOf("ConstantAttr") &&
          "must be subclass of TableGen 'ConstantAttr' class");
 }
 
 Attribute ConstantAttr::getAttribute() const {
   return Attribute(def->getValueAsDef("attr"));
 }

mlir/lib/TableGen/Constraint.cpp

 }
 
 StringRef Constraint::getSummary() const {
   if (Optional<StringRef> summary = def->getValueAsOptionalString("summary"))
     return *summary;
   return def->getName();
 }
 
+StringRef Constraint::getDescription() const {
+  return def->getValueAsOptionalString("description").getValueOr("");
+}
+
 StringRef Constraint::getDefName() const {
   if (Optional<StringRef> baseDefName = getBaseDefName())
     return *baseDefName;
   return def->getName();
 }
 
 std::string Constraint::getUniqueDefName() const {
   std::string defName = def->getName().str();

mlir/lib/TableGen/Type.cpp

if (const llvm::RecordVal *value = def->getValue("dialect")) {
     Dialect dialect(cast<const llvm::DefInit>(value->getValue())->getDef());
     return (dialect.getCppNamespace() + "::" + className).str();
   }
   return className.str();
 }
 
 Type::Type(const llvm::Record *record) : TypeConstraint(record) {}
 
-StringRef Type::getDescription() const {
-  return def->getValueAsString("description");
-}
-
 Dialect Type::getDialect() const {
   return Dialect(def->getValueAsDef("dialect"));
 }

mlir/lib/Tools/PDLL/AST/Nodes.cpp

 //===----------------------------------------------------------------------===//
 
 TypeExpr *TypeExpr::create(Context &ctx, SMRange loc, StringRef value) {
   return new (ctx.getAllocator().Allocate<TypeExpr>())
       TypeExpr(ctx, loc, copyStringWithNull(ctx, value));
 }
 
 //===----------------------------------------------------------------------===//
+// Decl
+//===----------------------------------------------------------------------===//
+
+void Decl::setDocComment(Context &ctx, StringRef comment) {
+  docComment = comment.copy(ctx.getAllocator());
+}
+
+//===----------------------------------------------------------------------===//
 // AttrConstraintDecl
 //===----------------------------------------------------------------------===//
 
 AttrConstraintDecl *AttrConstraintDecl::create(Context &ctx, SMRange loc,
                                                Expr *typeExpr) {
   return new (ctx.getAllocator().Allocate<AttrConstraintDecl>())
       AttrConstraintDecl(loc, typeExpr);
 }

mlir/lib/Tools/PDLL/ODS/Operation.cpp

 
 //===----------------------------------------------------------------------===//
 // Operation
 //===----------------------------------------------------------------------===//
 
 Operation::Operation(StringRef name, StringRef summary, StringRef desc,
                      StringRef nativeClassName, bool supportsTypeInferrence,
                      llvm::SMLoc loc)
-    : name(name.str()), summary(summary.str()),
+    : name(name.str()), summary(summary.str()), description(desc.str()),
       nativeClassName(nativeClassName.str()),
       supportsTypeInferrence(supportsTypeInferrence),
-      location(loc, llvm::SMLoc::getFromPointer(loc.getPointer() + 1)) {
-  llvm::raw_string_ostream descOS(description);
-  raw_indented_ostream(descOS).printReindented(desc.rtrim(" \t"));
-}
+      location(loc, llvm::SMLoc::getFromPointer(loc.getPointer() + 1)) {}

mlir/lib/Tools/PDLL/Parser/Parser.cpp

 //===- Parser.cpp ---------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/PDLL/Parser/Parser.h"
 #include "Lexer.h"
+#include "mlir/Support/IndentedOstream.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/TableGen/Argument.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Constraint.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/Tools/PDLL/AST/Context.h"
 #include "mlir/Tools/PDLL/AST/Diagnostic.h"

 //===----------------------------------------------------------------------===//
 // Parser
 //===----------------------------------------------------------------------===//
 
 namespace {
 class Parser {
 public:
   Parser(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
-         CodeCompleteContext *codeCompleteContext)
+         bool enableDocumentation, CodeCompleteContext *codeCompleteContext)
       : ctx(ctx), lexer(sourceMgr, ctx.getDiagEngine(), codeCompleteContext),
-        curToken(lexer.lexToken()), valueTy(ast::ValueType::get(ctx)),
+        curToken(lexer.lexToken()), enableDocumentation(enableDocumentation),
+        valueTy(ast::ValueType::get(ctx)),
         valueRangeTy(ast::ValueRangeType::get(ctx)),
         typeTy(ast::TypeType::get(ctx)),
         typeRangeTy(ast::TypeRangeType::get(ctx)),
         attrTy(ast::AttributeType::get(ctx)),
         codeCompleteContext(codeCompleteContext) {}
 
   /// Try to parse a new module. Returns nullptr in the case of failure.
   FailureOr<ast::Module *> parseModule();
private:
   ast::Expr *convertOpToValue(const ast::Expr *opExpr);
 
   /// Lookup ODS information for the given operation, returns nullptr if no
   /// information is found.
   const ods::Operation *lookupODSOperation(Optional<StringRef> opName) {
     return opName ? ctx.getODSContext().lookupOperation(*opName) : nullptr;
   }
 
+  /// Process the given documentation string, or return an empty string if
+  /// documentation isn't enabled.
+  StringRef processDoc(StringRef doc) {
+    return enableDocumentation ? doc : StringRef();
+  }
+
+  /// Process the given documentation string and format it, or return an empty
+  /// string if documentation isn't enabled.
+  std::string processAndFormatDoc(const Twine &doc) {
+    if (!enableDocumentation)
+      return "";
+    std::string docStr;
+    {
+      llvm::raw_string_ostream docOS(docStr);
+      std::string tmpDocStr = doc.str();
+      raw_indented_ostream(docOS).printReindented(
+          StringRef(tmpDocStr).rtrim(" \t"));
+    }
+    return docStr;
+  }
+
   //===--------------------------------------------------------------------===//
   // Directives
 
   LogicalResult parseDirective(SmallVectorImpl<ast::Decl *> &decls);
   LogicalResult parseInclude(SmallVectorImpl<ast::Decl *> &decls);
   LogicalResult parseTdInclude(StringRef filename, SMRange fileLoc,
                                SmallVectorImpl<ast::Decl *> &decls);
 
   /// Process the records of a parsed tablegen include file.
   void processTdIncludeRecords(llvm::RecordKeeper &tdRecords,
                                SmallVectorImpl<ast::Decl *> &decls);
 
   /// Create a user defined native constraint for a constraint imported from
   /// ODS.
   template <typename ConstraintT>
-  ast::Decl *createODSNativePDLLConstraintDecl(StringRef name,
-                                               StringRef codeBlock, SMRange loc,
-                                               ast::Type type,
-                                               StringRef nativeType);
+  ast::Decl *
+  createODSNativePDLLConstraintDecl(StringRef name, StringRef codeBlock,
+                                    SMRange loc, ast::Type type,
+                                    StringRef nativeType, StringRef docString);
   template <typename ConstraintT>
   ast::Decl *
   createODSNativePDLLConstraintDecl(const tblgen::Constraint &constraint,
                                     SMRange loc, ast::Type type,
                                     StringRef nativeType);
 
   //===--------------------------------------------------------------------===//
   // Decls
private:
   ast::Context &ctx;
 
   /// The lexer of this parser.
   Lexer lexer;
 
   /// The current token within the lexer.
   Token curToken;
 
+  /// A flag indicating if the parser should add documentation to AST nodes when
+  /// viable.
+  bool enableDocumentation;
+
   /// The most recently defined decl scope.
   ast::DeclScope *curDeclScope = nullptr;
   llvm::SpecificBumpPtrAllocator<ast::DeclScope> scopeAllocator;
 
   /// The current context of the parser.
   ParserContext parserContext = ParserContext::Global;
 
   /// Cached types to simplify verification and expression creation.
auto getLengthKind = [](const auto &value) {
     return value.isVariadic() ? ods::VariableLengthKind::Variadic
                               : ods::VariableLengthKind::Single;
   };
 
   // Insert a type constraint into the ODS context.
   ods::Context &odsContext = ctx.getODSContext();
   auto addTypeConstraint = [&](const tblgen::NamedTypeConstraint &cst)
       -> const ods::TypeConstraint & {
-    return odsContext.insertTypeConstraint(cst.constraint.getUniqueDefName(),
-                                           cst.constraint.getSummary(),
-                                           cst.constraint.getCPPClassName());
+    return odsContext.insertTypeConstraint(
+        cst.constraint.getUniqueDefName(),
+        processDoc(cst.constraint.getSummary()),
   };
   auto convertLocToRange = [&](llvm::SMLoc loc) -> llvm::SMRange {
     return {loc, llvm::SMLoc::getFromPointer(loc.getPointer() + 1)};
   };
 
   // Process the parsed tablegen records to build ODS information.
   /// Operations.
   for (llvm::Record *def : tdRecords.getAllDerivedDefinitions("Op")) {
     tblgen::Operator op(def);
 
     // Check to see if this operation is known to support type inferrence.
     bool supportsResultTypeInferrence =
         op.getTrait("::mlir::InferTypeOpInterface::Trait");
 
     bool inserted = false;
     ods::Operation *odsOp = nullptr;
     std::tie(odsOp, inserted) = odsContext.insertOperation(
-        op.getOperationName(), op.getSummary(), op.getDescription(),
-        op.getQualCppClassName(), supportsResultTypeInferrence,
-        op.getLoc().front());
+        op.getOperationName(), processDoc(op.getSummary()),
+        processAndFormatDoc(op.getDescription()), op.getQualCppClassName(),
+        supportsResultTypeInferrence, op.getLoc().front());
 
     // Ignore operations that have already been added.
     if (!inserted)
       continue;
 
     for (const tblgen::NamedAttribute &attr : op.getAttributes()) {
-      odsOp->appendAttribute(
-          attr.name, attr.attr.isOptional(),
-          odsContext.insertAttributeConstraint(attr.attr.getUniqueDefName(),
-                                               attr.attr.getSummary(),
-                                               attr.attr.getStorageType()));
+      odsOp->appendAttribute(attr.name, attr.attr.isOptional(),
+                             odsContext.insertAttributeConstraint(
+                                 attr.attr.getUniqueDefName(),
+                                 processDoc(attr.attr.getSummary()),
     }
     for (const tblgen::NamedTypeConstraint &operand : op.getOperands()) {
       odsOp->appendOperand(operand.name, getLengthKind(operand),
                            addTypeConstraint(operand));
     }
     for (const tblgen::NamedTypeConstraint &result : op.getResults()) {
       odsOp->appendResult(result.name, getLengthKind(result),
                           addTypeConstraint(result));
std::string cppClassName =
         llvm::formatv("{0}::{1}", def->getValueAsString("cppNamespace"),
                       def->getValueAsString("cppClassName"))
             .str();
     std::string codeBlock =
         llvm::formatv("return ::mlir::success(llvm::isa<{0}>(self));",
                       cppClassName)
             .str();
 
+    std::string desc =
+        processAndFormatDoc(def->getValueAsString("description"));
     if (def->isSubClassOf("OpInterface")) {
       decls.push_back(createODSNativePDLLConstraintDecl<ast::OpConstraintDecl>(
-          name, codeBlock, loc, opTy, cppClassName));
+          name, codeBlock, loc, opTy, cppClassName, desc));
     } else if (def->isSubClassOf("AttrInterface")) {
       decls.push_back(
           createODSNativePDLLConstraintDecl<ast::AttrConstraintDecl>(
-              name, codeBlock, loc, attrTy, cppClassName));
+              name, codeBlock, loc, attrTy, cppClassName, desc));
     } else if (def->isSubClassOf("TypeInterface")) {
       decls.push_back(
           createODSNativePDLLConstraintDecl<ast::TypeConstraintDecl>(
-              name, codeBlock, loc, typeTy, cppClassName));
+              name, codeBlock, loc, typeTy, cppClassName, desc));
     }
   }
 }
 
 template <typename ConstraintT>
-ast::Decl *
-Parser::createODSNativePDLLConstraintDecl(StringRef name, StringRef codeBlock,
-                                          SMRange loc, ast::Type type,
-                                          StringRef nativeType) {
+ast::Decl *Parser::createODSNativePDLLConstraintDecl(
+    StringRef name, StringRef codeBlock, SMRange loc, ast::Type type,
+    StringRef nativeType, StringRef docString) {
   // Build the single input parameter.
   ast::DeclScope *argScope = pushDeclScope();
   auto *paramVar = ast::VariableDecl::create(
       ctx, ast::Name::create(ctx, "self", loc), type,
       /*initExpr=*/nullptr, ast::ConstraintRef(ConstraintT::create(ctx, loc)));
   argScope->add(paramVar);
   popDeclScope();
 
   // Build the native constraint.
   auto *constraintDecl = ast::UserConstraintDecl::createNative(
       ctx, ast::Name::create(ctx, name, loc), paramVar,
       /*results=*/llvm::None, codeBlock, ast::TupleType::get(ctx), nativeType);
+  constraintDecl->setDocComment(ctx, docString);
   curDeclScope->add(constraintDecl);
   return constraintDecl;
 }
 
 template <typename ConstraintT>
 ast::Decl *
 Parser::createODSNativePDLLConstraintDecl(const tblgen::Constraint &constraint,
                                           SMRange loc, ast::Type type,
                                           StringRef nativeType) {
   // Format the condition template.
   tblgen::FmtContext fmtContext;
   fmtContext.withSelf("self");
   std::string codeBlock = tblgen::tgfmt(
       "return ::mlir::success(" + constraint.getConditionTemplate() + ");",
       &fmtContext);
 
+  // If documentation was enabled, build the doc string for the generated
+  // constraint. It would be nice to do this lazily, but TableGen information is
+  // destroyed after we finish parsing the file.
+  std::string docString;
+  if (enableDocumentation) {
+    StringRef desc = constraint.getDescription();
+    docString = processAndFormatDoc(
+        constraint.getSummary() +
+        (desc.empty() ? "" : ("\n\n" + constraint.getDescription())));
+  }
+
   return createODSNativePDLLConstraintDecl<ConstraintT>(
-      constraint.getUniqueDefName(), codeBlock, loc, type, nativeType);
+      constraint.getUniqueDefName(), codeBlock, loc, type, nativeType,
+      docString);
 }
 
 //===----------------------------------------------------------------------===//
 // Decls
 
 FailureOr<ast::Decl *> Parser::parseTopLevelDecl() {
   FailureOr<ast::Decl *> decl;
   switch (curToken.getKind()) {
codeCompleteContext->codeCompleteOperationResultsSignature(opName,
                                                              currentNumResults);
 }
 
 //===----------------------------------------------------------------------===//
 // Parser
 //===----------------------------------------------------------------------===//
 
 FailureOr<ast::Module *>
-mlir::pdll::parsePDLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
+mlir::pdll::parsePDLLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
+                         bool enableDocumentation,
                          CodeCompleteContext *codeCompleteContext) {
-  Parser parser(ctx, sourceMgr, codeCompleteContext);
+  Parser parser(ctx, sourceMgr, enableDocumentation, codeCompleteContext);
   return parser.parseModule();
 }

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

relatedDiags.emplace_back(
         note.getMessage().str());
   }
   if (!relatedDiags.empty())
     lspDiag.relatedInformation = std::move(relatedDiags);
 
   return lspDiag;
 }
 
+/// Get or extract the documentation for the given decl.
+static Optional<std::string> getDocumentationFor(llvm::SourceMgr &sourceMgr,
+                                                 const ast::Decl *decl) {
+  // If the decl already had documentation set, use it.
+  if (Optional<StringRef> doc = decl->getDocComment())
+    return doc->str();
+
+  // If the decl doesn't yet have documentation, try to extract it from the
+  // source file. This is a heuristic, and isn't intended to cover every case,
+  // but should cover the most common. We essentially look for a comment
+  // preceding the decl, and if we find one, use that as the documentation.
+  SMLoc startLoc = decl->getLoc().Start;
+  if (!startLoc.isValid())
+    return llvm::None;
+  int bufferId = sourceMgr.FindBufferContainingLoc(startLoc);
+  if (bufferId == 0)
+    return llvm::None;
+  const char *bufferStart =
+      sourceMgr.getMemoryBuffer(bufferId)->getBufferStart();
+  StringRef buffer(bufferStart, startLoc.getPointer() - bufferStart);
+
+  // Pop the last line from the buffer string.
+  auto popLastLine = [&]() -> Optional<StringRef> {
+    size_t newlineOffset = buffer.find_last_of("\n");
+    if (newlineOffset == StringRef::npos)
+      return llvm::None;
+    StringRef lastLine = buffer.drop_front(newlineOffset).trim();
+    buffer = buffer.take_front(newlineOffset);
+    return lastLine;
+  };
+
+  // Try to pop the current line, which contains the decl.
+  if (!popLastLine())
+    return llvm::None;
+
+  // Try to parse a comment string from the source file.
+  SmallVector<StringRef> commentLines;
+  while (Optional<StringRef> line = popLastLine()) {
+    // Check for a comment at the beginning of the line.
+    if (!line->startswith("//"))
+      break;
+
+    // Extract the document string from the comment.
+    commentLines.push_back(line->drop_while([](char c) { return c == '/'; }));
+  }
+
+  if (commentLines.empty())
+    return llvm::None;
+  return llvm::join(llvm::reverse(commentLines), "\n");
+}
+
 //===----------------------------------------------------------------------===//
 // PDLIndex
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct PDLIndexSymbol {
   explicit PDLIndexSymbol(const ast::Decl *definition)
       : definition(definition) {}
struct PDLDocument {
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  const lsp::Position &hoverPos);
   Optional<lsp::Hover> findHover(const ast::Decl *decl,
                                  const SMRange &hoverRange);
   lsp::Hover buildHoverForOpName(const ods::Operation *op,
                                  const SMRange &hoverRange);
   lsp::Hover buildHoverForVariable(const ast::VariableDecl *varDecl,
                                    const SMRange &hoverRange);
-  lsp::Hover buildHoverForPattern(const ast::PatternDecl *patternDecl,
+  lsp::Hover buildHoverForPattern(const ast::PatternDecl *decl,
                                   const SMRange &hoverRange);
   lsp::Hover buildHoverForCoreConstraint(const ast::CoreConstraintDecl *decl,
                                          const SMRange &hoverRange);
   template <typename T>
   lsp::Hover buildHoverForUserConstraintOrRewrite(StringRef typeName,
                                                   const T *decl,
                                                   const SMRange &hoverRange);
 
PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
 
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
 
   astContext.getDiagEngine().setHandlerFn([&](const ast::Diagnostic &diag) {
     if (auto lspDiag = getLspDiagnoticFromDiag(sourceMgr, diag, uri))
       diagnostics.push_back(std::move(*lspDiag));
   });
-  astModule = parsePDLAST(astContext, sourceMgr);
+  astModule = parsePDLLAST(astContext, sourceMgr, /*enableDocumentation=*/true);
 
   // Initialize the set of parsed includes.
   lsp::gatherIncludeFiles(sourceMgr, parsedIncludes);
 
   // If we failed to parse the module, there is nothing left to initialize.
   if (failed(astModule))
     return;
 
lsp::Hover PDLDocument::buildHoverForVariable(const ast::VariableDecl *varDecl,
   {
     llvm::raw_string_ostream hoverOS(hover.contents.value);
     hoverOS << "**Variable**: `" << varDecl->getName().getName() << "`\n***\n"
             << "Type: `" << varDecl->getType() << "`\n";
   }
   return hover;
 }
 
-lsp::Hover
-PDLDocument::buildHoverForPattern(const ast::PatternDecl *patternDecl,
+lsp::Hover PDLDocument::buildHoverForPattern(const ast::PatternDecl *decl,
                                              const SMRange &hoverRange) {
   lsp::Hover hover(lsp::Range(sourceMgr, hoverRange));
   {
     llvm::raw_string_ostream hoverOS(hover.contents.value);
     hoverOS << "**Pattern**";
-    if (const ast::Name *name = patternDecl->getName())
+    if (const ast::Name *name = decl->getName())
       hoverOS << ": `" << name->getName() << "`";
     hoverOS << "\n***\n";
-    if (Optional<uint16_t> benefit = patternDecl->getBenefit())
+    if (Optional<uint16_t> benefit = decl->getBenefit())
       hoverOS << "Benefit: " << *benefit << "\n";
-    if (patternDecl->hasBoundedRewriteRecursion())
+    if (decl->hasBoundedRewriteRecursion())
       hoverOS << "HasBoundedRewriteRecursion\n";
     hoverOS << "RootOp: `"
-            << patternDecl->getRootRewriteStmt()->getRootOpExpr()->getType()
-            << "`\n";
+            << decl->getRootRewriteStmt()->getRootOpExpr()->getType() << "`\n";
+
+    // Format the documentation for the decl.
+    if (Optional<std::string> doc = getDocumentationFor(sourceMgr, decl))
+      hoverOS << "\n" << *doc << "\n";
   }
   return hover;
 }
 
 lsp::Hover
 PDLDocument::buildHoverForCoreConstraint(const ast::CoreConstraintDecl *decl,
                                          const SMRange &hoverRange) {
   lsp::Hover hover(lsp::Range(sourceMgr, hoverRange));
if (!inputs.empty()) {
       hoverOS << "Parameters:\n";
       for (const ast::VariableDecl *input : inputs)
         hoverOS << "* " << input->getName().getName() << ": `"
                 << input->getType() << "`\n";
       hoverOS << "***\n";
     }
     ast::Type resultType = decl->getResultType();
     if (auto resultTupleTy = resultType.dyn_cast<ast::TupleType>()) {
-      if (resultTupleTy.empty())
-        return hover;
-
+      if (!resultTupleTy.empty()) {
         hoverOS << "Results:\n";
         for (auto it : llvm::zip(resultTupleTy.getElementNames(),
                                  resultTupleTy.getElementTypes())) {
           StringRef name = std::get<0>(it);
           hoverOS << "* " << (name.empty() ? "" : (name + ": ")) << "`"
                   << std::get<1>(it) << "`\n";
         }
+        hoverOS << "***\n";
+      }
     } else {
       hoverOS << "Results:\n* `" << resultType << "`\n";
-    }
       hoverOS << "***\n";
     }
+
+    // Format the documentation for the decl.
+    if (Optional<std::string> doc = getDocumentationFor(sourceMgr, decl))
+      hoverOS << "\n" << *doc << "\n";
+  }
   return hover;
 }
 
 //===----------------------------------------------------------------------===//
 // PDLDocument: Document Symbols
 //===----------------------------------------------------------------------===//
 
 void PDLDocument::findDocumentSymbols(

 
 //===----------------------------------------------------------------------===//
 // PDLDocument: Code Completion
 //===----------------------------------------------------------------------===//
 
 namespace {
 class LSPCodeCompleteContext : public CodeCompleteContext {
 public:
-  LSPCodeCompleteContext(SMLoc completeLoc, lsp::CompletionList &completionList,
+  LSPCodeCompleteContext(SMLoc completeLoc, llvm::SourceMgr &sourceMgr,
+                         lsp::CompletionList &completionList,
                          ods::Context &odsContext,
                          ArrayRef<std::string> includeDirs)
-      : CodeCompleteContext(completeLoc), completionList(completionList),
-        odsContext(odsContext), includeDirs(includeDirs) {}
+      : CodeCompleteContext(completeLoc), sourceMgr(sourceMgr),
+        completionList(completionList), odsContext(odsContext),
+        includeDirs(includeDirs) {}
 
   void codeCompleteTupleMemberAccess(ast::TupleType tupleType) final {
     ArrayRef<ast::Type> elementTypes = tupleType.getElementTypes();
     ArrayRef<StringRef> elementNames = tupleType.getElementNames();
     for (unsigned i = 0, e = tupleType.size(); i < e; ++i) {
       // Push back a completion item that uses the result index.
       lsp::CompletionItem item;
       item.label = llvm::formatv("{0} (field #{0})", i).str();
while (scope) {
             strOS << "(";
             llvm::interleaveComma(
                 cst->getInputs(), strOS, [&](const ast::VariableDecl *var) {
                   strOS << var->getName().getName() << ": " << var->getType();
                 });
             strOS << ") -> " << cst->getResultType();
           }
 
+          // Format the documentation for the constraint.
+          if (Optional<std::string> doc = getDocumentationFor(sourceMgr, cst)) {
+            item.documentation =
+                lsp::MarkupContent{lsp::MarkupKind::Markdown, std::move(*doc)};
+          }
+
           completionList.items.emplace_back(item);
         }
       }
 
       scope = scope->getParentScope();
     }
   }
 
void codeCompleteIncludeFilename(StringRef curPath) final {
     // the face of different iteration schemes for different platforms.
     llvm::sort(completionList.items, [](const lsp::CompletionItem &lhs,
                                         const lsp::CompletionItem &rhs) {
       return lhs.label < rhs.label;
     });
   }
 
 private:
+  llvm::SourceMgr &sourceMgr;
   lsp::CompletionList &completionList;
   ods::Context &odsContext;
   ArrayRef<std::string> includeDirs;
 };
 } // namespace
 
 lsp::CompletionList
 PDLDocument::getCodeCompletion(const lsp::URIForFile &uri,
                                const lsp::Position &completePos) {
   SMLoc posLoc = completePos.getAsSMLoc(sourceMgr);
   if (!posLoc.isValid())
     return lsp::CompletionList();
 
   // To perform code completion, we run another parse of the module with the
   // code completion context provided.
   ods::Context tmpODSContext;
   lsp::CompletionList completionList;
-  LSPCodeCompleteContext lspCompleteContext(
-      posLoc, completionList, tmpODSContext, sourceMgr.getIncludeDirs());
+  LSPCodeCompleteContext lspCompleteContext(posLoc, sourceMgr, completionList,
+                                            tmpODSContext,
+                                            sourceMgr.getIncludeDirs());
 
   ast::Context tmpContext(tmpODSContext);
-  (void)parsePDLAST(tmpContext, sourceMgr, &lspCompleteContext);
+  (void)parsePDLLAST(tmpContext, sourceMgr, /*enableDocumentation=*/true,
+                     &lspCompleteContext);
 
   return completionList;
 }
 
 //===----------------------------------------------------------------------===//
 // PDLDocument: Signature Help
 //===----------------------------------------------------------------------===//
 
 namespace {
 class LSPSignatureHelpContext : public CodeCompleteContext {
 public:
-  LSPSignatureHelpContext(SMLoc completeLoc, lsp::SignatureHelp &signatureHelp,
+  LSPSignatureHelpContext(SMLoc completeLoc, llvm::SourceMgr &sourceMgr,
+                          lsp::SignatureHelp &signatureHelp,
                           ods::Context &odsContext)
-      : CodeCompleteContext(completeLoc), signatureHelp(signatureHelp),
-        odsContext(odsContext) {}
+      : CodeCompleteContext(completeLoc), sourceMgr(sourceMgr),
+        signatureHelp(signatureHelp), odsContext(odsContext) {}
 
   void codeCompleteCallSignature(const ast::CallableDecl *callable,
                                  unsigned currentNumArgs) final {
     signatureHelp.activeParameter = currentNumArgs;
 
     lsp::SignatureInformation signatureInfo;
     {
       llvm::raw_string_ostream strOS(signatureInfo.label);
       strOS << callable->getName()->getName() << "(";
       auto formatParamFn = [&](const ast::VariableDecl *var) {
         unsigned paramStart = strOS.str().size();
         strOS << var->getName().getName() << ": " << var->getType();
         unsigned paramEnd = strOS.str().size();
         signatureInfo.parameters.emplace_back(lsp::ParameterInformation{
             StringRef(strOS.str()).slice(paramStart, paramEnd).str(),
             std::make_pair(paramStart, paramEnd), /*paramDoc*/ std::string()});
       };
       llvm::interleaveComma(callable->getInputs(), strOS, formatParamFn);
       strOS << ") -> " << callable->getResultType();
     }
+
+    // Format the documentation for the callable.
+    if (Optional<std::string> doc = getDocumentationFor(sourceMgr, callable))
+      signatureInfo.documentation = std::move(*doc);
+
     signatureHelp.signatures.emplace_back(std::move(signatureInfo));
   }
 
   void
   codeCompleteOperationOperandsSignature(Optional<StringRef> opName,
                                          unsigned currentNumOperands) final {
     const ods::Operation *odsOp =
         opName ? odsContext.lookupOperation(*opName) : nullptr;
if (currentValue == 0 && (!odsOp || !values.empty())) {
           StringRef(signatureInfo.label).drop_front().drop_back().str(),
           std::pair<unsigned, unsigned>(1, signatureInfo.label.size() - 1),
           ("All of the " + label + "s of the operation.").str()});
       signatureHelp.signatures.emplace_back(std::move(signatureInfo));
     }
   }
 
 private:
+  llvm::SourceMgr &sourceMgr;
   lsp::SignatureHelp &signatureHelp;
   ods::Context &odsContext;
 };
 } // namespace
 
 lsp::SignatureHelp PDLDocument::getSignatureHelp(const lsp::URIForFile &uri,
                                                  const lsp::Position &helpPos) {
   SMLoc posLoc = helpPos.getAsSMLoc(sourceMgr);
   if (!posLoc.isValid())
     return lsp::SignatureHelp();
 
   // To perform code completion, we run another parse of the module with the
   // code completion context provided.
   ods::Context tmpODSContext;
   lsp::SignatureHelp signatureHelp;
-  LSPSignatureHelpContext completeContext(posLoc, signatureHelp, tmpODSContext);
+  LSPSignatureHelpContext completeContext(posLoc, sourceMgr, signatureHelp,
+                                          tmpODSContext);
 
   ast::Context tmpContext(tmpODSContext);
-  (void)parsePDLAST(tmpContext, sourceMgr, &completeContext);
+  (void)parsePDLLAST(tmpContext, sourceMgr, /*enableDocumentation=*/true,
+                     &completeContext);
 
   return signatureHelp;
 }
 
 //===----------------------------------------------------------------------===//
 // PDLL ViewOutput
 //===----------------------------------------------------------------------===//
 

mlir/test/mlir-pdll-lsp-server/hover.test

 // RUN: mlir-pdll-lsp-server -pdll-extra-dir %S -pdll-extra-dir %S/../../include -lit-test < %s | FileCheck %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"pdll","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.pdll",
   "languageId":"pdll",
   "version":1,
-  "text":"Constraint FooCst();\nRewrite FooRewrite(op: Op) -> Op;\nPattern Foo {\nlet root: Op;\nerase root;\n}\n#include \"include/included.td\"\n#include \"include/included.pdll\""
+  "text":"Constraint FooCst();\n// This is documentation for the rewriter.\n/// And even more docs.\nRewrite FooRewrite(op: Op) -> Op;\nPattern Foo {\nlet root: Op;\nerase root;\n}\n#include \"include/included.td\"\n#include \"include/included.pdll\""
 }}}
 // -----
 // Hover on a variable.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":3,"character":6}
+  "position":{"line":5,"character":6}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "**Variable**: `root`\n***\nType: `Op`\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 8,
-// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:        "line": 5
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 4,
-// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:        "line": 5
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a pattern.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":2,"character":9}
+  "position":{"line":4,"character":9}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "**Pattern**: `Foo`\n***\nRootOp: `Op`\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 11,
-// CHECK-NEXT:        "line": 2
+// CHECK-NEXT:        "line": 4
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 8,
-// CHECK-NEXT:        "line": 2
+// CHECK-NEXT:        "line": 4
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a core constraint.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":3,"character":11}
+  "position":{"line":5,"character":11}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "**Constraint**: `Op`\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 12,
-// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:        "line": 5
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 10,
-// CHECK-NEXT:        "line": 3
+// CHECK-NEXT:        "line": 5
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a user constraint.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
   "position":{"line":0,"character":14}

 // CHECK-NEXT:        "line": 0
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on a user rewrite.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":1,"character":11}
+  "position":{"line":3,"character":11}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
-// CHECK-NEXT:      "value": "**Rewrite**: `FooRewrite`\n***\nParameters:\n* op: `Op`\n***\nResults:\n* `Op`\n***\n"
+// CHECK-NEXT:      "value": "**Rewrite**: `FooRewrite`\n***\nParameters:\n* op: `Op`\n***\nResults:\n* `Op`\n***\n\n This is documentation for the rewriter.\n And even more docs.\n"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 18,
-// CHECK-NEXT:        "line": 1
+// CHECK-NEXT:        "line": 3
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 8,
-// CHECK-NEXT:        "line": 1
+// CHECK-NEXT:        "line": 3
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on an include file.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":6,"character":15}
+  "position":{"line":8,"character":15}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "`included.td`\n***\n{{.*}}included.td"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 30,
-// CHECK-NEXT:        "line": 6
+// CHECK-NEXT:        "line": 8
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 9,
-// CHECK-NEXT:        "line": 6
+// CHECK-NEXT:        "line": 8
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 // Hover on an include file.
 {"jsonrpc":"2.0","id":1,"method":"textDocument/hover","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":7,"character":15}
+  "position":{"line":9,"character":15}
 }}
 //      CHECK:  "id": 1,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "contents": {
 // CHECK-NEXT:      "kind": "markdown",
 // CHECK-NEXT:      "value": "`included.pdll`\n***\n{{.*}}included.pdll"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "range": {
 // CHECK-NEXT:      "end": {
 // CHECK-NEXT:        "character": 32,
-// CHECK-NEXT:        "line": 7
+// CHECK-NEXT:        "line": 9
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "start": {
 // CHECK-NEXT:        "character": 9,
-// CHECK-NEXT:        "line": 7
+// CHECK-NEXT:        "line": 9
 // CHECK-NEXT:      }
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 {"jsonrpc":"2.0","id":7,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/tools/mlir-pdll/mlir-pdll.cpp

 static LogicalResult
 processBuffer(raw_ostream &os, std::unique_ptr<llvm::MemoryBuffer> chunkBuffer,
               OutputType outputType, std::vector<std::string> &includeDirs,
               bool dumpODS, std::set<std::string> *includedFiles) {
   llvm::SourceMgr sourceMgr;
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(chunkBuffer), SMLoc());
 
+  // If we are dumping ODS information, also enable documentation to ensure the
+  // summary and description information is imported as well.
+  bool enableDocumentation = dumpODS;
+
   ods::Context odsContext;
   ast::Context astContext(odsContext);
-  FailureOr<ast::Module *> module = parsePDLAST(astContext, sourceMgr);
+  FailureOr<ast::Module *> module =
+      parsePDLLAST(astContext, sourceMgr, enableDocumentation);
   if (failed(module))
     return failure();
 
   // Add the files that were included to the set.
   if (includedFiles) {
     for (unsigned i = 1, e = sourceMgr.getNumBuffers(); i < e; ++i) {
       includedFiles->insert(
           sourceMgr.getMemoryBuffer(i + 1)->getBufferIdentifier().str());