[mlir:PDLL] Add better support for providing Constraint/Pattern/Rewri…

…te documentation

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).

Differential Revision: https://reviews.llvm.org/D124881

mlir/include/mlir/TableGen/Attribute.h

@@ -113,9 +113,6 @@ class Attribute : public AttrConstraint {
 
   // 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

mlir/include/mlir/TableGen/Constraint.h

@@ -50,10 +50,14 @@ class Constraint {
   // 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).

mlir/include/mlir/TableGen/Type.h

@@ -65,9 +65,6 @@ 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;
 };

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

@@ -635,6 +635,13 @@ class Decl : public Node {
   /// 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) {}
@@ -643,6 +650,10 @@ class Decl : public Node {
   /// 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;
 };
 
 //===----------------------------------------------------------------------===//

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

@@ -26,12 +26,16 @@ 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,
-            CodeCompleteContext *codeCompleteContext = nullptr);
+parsePDLLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
+             bool enableDocumentation = false,
+             CodeCompleteContext *codeCompleteContext = nullptr);
 } // namespace pdll
 } // namespace mlir
 

mlir/lib/TableGen/Attribute.cpp

@@ -132,10 +132,6 @@ Dialect Attribute::getDialect() const {
   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");

mlir/lib/TableGen/Constraint.cpp

@@ -57,6 +57,10 @@ StringRef Constraint::getSummary() const {
   return def->getName();
 }
 
+StringRef Constraint::getDescription() const {
+  return def->getValueAsOptionalString("description").getValueOr("");
+}
+
 StringRef Constraint::getDefName() const {
   if (Optional<StringRef> baseDefName = getBaseDefName())
     return *baseDefName;

mlir/lib/TableGen/Type.cpp

@@ -76,10 +76,6 @@ std::string TypeConstraint::getCPPClassName() const {
 
 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

@@ -355,6 +355,14 @@ TypeExpr *TypeExpr::create(Context &ctx, SMRange loc, StringRef value) {
       TypeExpr(ctx, loc, copyStringWithNull(ctx, value));
 }
 
+//===----------------------------------------------------------------------===//
+// Decl
+//===----------------------------------------------------------------------===//
+
+void Decl::setDocComment(Context &ctx, StringRef comment) {
+  docComment = comment.copy(ctx.getAllocator());
+}
+
 //===----------------------------------------------------------------------===//
 // AttrConstraintDecl
 //===----------------------------------------------------------------------===//

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

@@ -20,10 +20,7 @@ using namespace mlir::pdll::ods;
 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

@@ -8,6 +8,7 @@
 
 #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"
@@ -43,9 +44,10 @@ 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)),
@@ -125,6 +127,27 @@ class Parser {
     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
 
@@ -140,10 +163,10 @@ class Parser {
   /// 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,
@@ -520,6 +543,10 @@ class Parser {
   /// 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;
@@ -801,9 +828,10 @@ void Parser::processTdIncludeRecords(llvm::RecordKeeper &tdRecords,
   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()),
+        cst.constraint.getCPPClassName());
   };
   auto convertLocToRange = [&](llvm::SMLoc loc) -> llvm::SMRange {
     return {loc, llvm::SMLoc::getFromPointer(loc.getPointer() + 1)};
@@ -821,20 +849,20 @@ void Parser::processTdIncludeRecords(llvm::RecordKeeper &tdRecords,
     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()),
+                                 attr.attr.getStorageType()));
     }
     for (const tblgen::NamedTypeConstraint &operand : op.getOperands()) {
       odsOp->appendOperand(operand.name, getLengthKind(operand),
@@ -883,26 +911,27 @@ void Parser::processTdIncludeRecords(llvm::RecordKeeper &tdRecords,
                       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(
@@ -915,6 +944,7 @@ Parser::createODSNativePDLLConstraintDecl(StringRef name, StringRef codeBlock,
   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;
 }
@@ -931,8 +961,20 @@ Parser::createODSNativePDLLConstraintDecl(const tblgen::Constraint &constraint,
       "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);
 }
 
 //===----------------------------------------------------------------------===//
@@ -3080,8 +3122,9 @@ void Parser::codeCompleteOperationResultsSignature(Optional<StringRef> opName,
 //===----------------------------------------------------------------------===//
 
 FailureOr<ast::Module *>
-mlir::pdll::parsePDLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
-                        CodeCompleteContext *codeCompleteContext) {
-  Parser parser(ctx, sourceMgr, codeCompleteContext);
+mlir::pdll::parsePDLLAST(ast::Context &ctx, llvm::SourceMgr &sourceMgr,
+                         bool enableDocumentation,
+                         CodeCompleteContext *codeCompleteContext) {
+  Parser parser(ctx, sourceMgr, enableDocumentation, codeCompleteContext);
   return parser.parseModule();
 }