Add extraInjectedClassDefinition/extraInjectedClassDeclaration to ODS Trait

mlir/include/mlir/IR/OpBase.td

class VariadicSuccessor<Successor successor>
   : Successor<successor.predicate, successor.summary>;
 
 
 //===----------------------------------------------------------------------===//
 // Trait definitions
 //===----------------------------------------------------------------------===//
 
 // Trait represents a trait regarding an attribute, operation, or type.
-class Trait;
+class Trait {
+  // Additional code that will be added to the public part of the generated
+  // C++ code of the entity (attribute, operation, or type) declaration.
+  code extraInjectedClassDeclaration = "";
+
+  // Additional code that will be added to the generated source file. The
+  // generated code is placed inside the entity's C++ namespace. `$cppClass` is
+  // replaced by the entity's C++ class name.
+  code extraInjectedClassDefinition = "";
+}
 
 // Define a Trait corresponding to a list of Traits, this allows for specifying
 // a list of traits as trait. Avoids needing to do `[Traits, ...] # ListOfTraits
 // # [Others, ...]` while still allowing providing convenient groupings.
 class TraitList<list<Trait> props> : Trait {
   list<Trait> traits = props;
 }
 

mlir/include/mlir/TableGen/Operator.h

public:
   bool hasSummary() const;
   StringRef getSummary() const;
 
   // Query functions for the assembly format of the operator.
   bool hasAssemblyFormat() const;
   StringRef getAssemblyFormat() const;
 
   // Returns this op's extra class declaration code.
-  StringRef getExtraClassDeclaration() const;
+  std::string getExtraClassDeclaration() const;
 
   // Returns this op's extra class definition code.
-  StringRef getExtraClassDefinition() const;
+  std::string getExtraClassDefinition() const;
 
   // Returns the Tablegen definition this operator was constructed from.
   // TODO: do not expose the TableGen record, this is a temporary solution to
   // OpEmitter requiring a Record because Operator does not provide enough
   // methods.
   const llvm::Record &getDef() const;
 
   // Returns the dialect of the op.

mlir/include/mlir/TableGen/Trait.h

public:
   // Returns an Trait corresponding to the init provided.
   static Trait create(const llvm::Init *init);
 
   Kind getKind() const { return kind; }
 
   // Returns the Tablegen definition this operator was constructed from.
   const llvm::Record &getDef() const { return *def; }
 
+  StringRef getExtraInjectedClassDeclaration() const;
+  StringRef getExtraInjectedClassDefinition() const;
+
 protected:
   // The TableGen definition of this trait.
   const llvm::Record *def;
   Kind kind;
 };
 
 // Trait corresponding to a native C++ Trait.
 class NativeTrait : public Trait {

mlir/lib/TableGen/Operator.cpp

 
 StringRef Operator::getCppNamespace() const { return cppNamespace; }
 
 int Operator::getNumResults() const {
   DagInit *results = def.getValueAsDag("results");
   return results->getNumArgs();
 }
 
-StringRef Operator::getExtraClassDeclaration() const {
+std::string Operator::getExtraClassDeclaration() const {
   constexpr auto attr = "extraClassDeclaration";
-  if (def.isValueUnset(attr))
-    return {};
-  return def.getValueAsString(attr);
+  std::string value;
+  if (!def.isValueUnset(attr))
+    value = def.getValueAsString(attr).str();
+  for (auto &trait : getTraits()) {
+    StringRef traitDeclaration = trait.getExtraInjectedClassDeclaration();
+    if (!traitDeclaration.empty()) {
+      value += "\n";
+      value += traitDeclaration.str();
+    }
+  }
+  return value;
 }
 
-StringRef Operator::getExtraClassDefinition() const {
+std::string Operator::getExtraClassDefinition() const {
   constexpr auto attr = "extraClassDefinition";
-  if (def.isValueUnset(attr))
-    return {};
-  return def.getValueAsString(attr);
+  std::string value;
+  if (!def.isValueUnset(attr))
+    value = def.getValueAsString(attr).str();
+  for (auto &trait : getTraits()) {
+    StringRef traitDefinition = trait.getExtraInjectedClassDefinition();
+    if (!traitDefinition.empty()) {
+      value += "\n";
+      value += traitDefinition.str();
+    }
+  }
+  return value;
 }
 
 const llvm::Record &Operator::getDef() const { return def; }
 
 bool Operator::skipDefaultBuilders() const {
   return def.getValueAsBit("skipDefaultBuilders");
 }
 

mlir/lib/TableGen/Trait.cpp

Trait Trait::create(const llvm::Init *init) {
   if (def->isSubClassOf("InterfaceTrait"))
     return Trait(Kind::Interface, def);
   assert(def->isSubClassOf("NativeTrait"));
   return Trait(Kind::Native, def);
 }
 
 Trait::Trait(Kind kind, const llvm::Record *def) : def(def), kind(kind) {}
 
+StringRef Trait::getExtraInjectedClassDeclaration() const {
+  constexpr auto attr = "extraInjectedClassDeclaration";
+  if (def->isValueUnset(attr))
+    return {};
+  return def->getValueAsString(attr);
+}
+
+StringRef Trait::getExtraInjectedClassDefinition() const {
+  constexpr auto attr = "extraInjectedClassDefinition";
+  if (def->isValueUnset(attr))
+    return {};
+  return def->getValueAsString(attr);
+}
+
 //===----------------------------------------------------------------------===//
 // NativeTrait
 //===----------------------------------------------------------------------===//
 
 std::string NativeTrait::getFullyQualifiedTraitName() const {
   llvm::StringRef trait = def->getValueAsString("trait");
   llvm::StringRef cppNamespace = def->getValueAsString("cppNamespace");
   return cppNamespace.empty() ? trait.str()

mlir/test/lib/Dialect/Test/TestOps.td

def FoldToCallOp : TEST_Op<"fold_to_call_op"> {
   let arguments = (ins FlatSymbolRefAttr:$callee);
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Test Traits
 //===----------------------------------------------------------------------===//
 
+def OpTraitWithInjectedCode : NativeOpTrait<"SameOperandsElementType"> {
+  code extraInjectedClassDeclaration = [{
+    bool testInjection();
+  }];
+  code extraInjectedClassDefinition = [{
+    bool $cppClass::testInjection() {
+      return true;
+    }
+  }];
+}
+
+def SameOperandElementTypeInjectedOp : TEST_Op<"same_operand_element_type_injected_code",
+    [OpTraitWithInjectedCode]> {
+  let arguments = (ins AnyType, AnyType);
+  let results = (outs AnyType);
+}
+
+
 def SameOperandElementTypeOp : TEST_Op<"same_operand_element_type",
     [SameOperandsElementType]> {
   let arguments = (ins AnyType, AnyType);
   let results = (outs AnyType);
 }
 
 def SameOperandAndResultElementTypeOp :
     TEST_Op<"same_operand_and_result_element_type",

mlir/tools/mlir-tblgen/OpClass.h

 public:
   /// Create an operation class with extra class declarations, whose default
   /// visibility is public. Also declares at the top of the class:
   ///
   /// - inheritance of constructors from `Op`
   /// - inheritance of `print`
   /// - a type alias for the associated adaptor class
   ///
-  OpClass(StringRef name, StringRef extraClassDeclaration,
+  OpClass(StringRef name, std::string extraClassDeclaration,
           std::string extraClassDefinition);
 
   /// Add an op trait.
   void addTrait(Twine trait) { parent.addTemplateParam(trait.str()); }
 
   /// The operation class is finalized by calling `Class::finalize` to delcare
   /// all pending private and public methods (ops don't have custom constructors
   /// or fields). Then, the extra class declarations are appended to the end of
   /// the class declaration.
   void finalize() override;
 
 private:
   /// Hand-written extra class declarations.
-  StringRef extraClassDeclaration;
+  std::string extraClassDeclaration;
   /// Hand-written extra class definitions.
   std::string extraClassDefinition;
   /// The parent class, which also contains the traits to be inherited.
   ParentClass &parent;
 };
 
 } // namespace tblgen
 } // namespace mlir
 
 #endif // MLIR_TOOLS_MLIRTBLGEN_OPCLASS_H_

mlir/tools/mlir-tblgen/OpClass.cpp

 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 //===----------------------------------------------------------------------===//
 // OpClass definitions
 //===----------------------------------------------------------------------===//
 
-OpClass::OpClass(StringRef name, StringRef extraClassDeclaration,
+OpClass::OpClass(StringRef name, std::string extraClassDeclaration,
                  std::string extraClassDefinition)
-    : Class(name.str()), extraClassDeclaration(extraClassDeclaration),
+    : Class(name.str()),
+      extraClassDeclaration(std::move(extraClassDeclaration)),
       extraClassDefinition(std::move(extraClassDefinition)),
       parent(addParent("::mlir::Op")) {
   parent.addTemplateParam(getClassName().str());
   declare<VisibilityDeclaration>(Visibility::Public);
   /// Inherit functions from Op.
   declare<UsingDeclaration>("Op::Op");
   declare<UsingDeclaration>("Op::print");
   /// Type alias for the adaptor class.
   declare<UsingDeclaration>("Adaptor", className + "Adaptor");
 }
 
 void OpClass::finalize() {
   Class::finalize();
   declare<VisibilityDeclaration>(Visibility::Public);
   declare<ExtraClassDeclaration>(extraClassDeclaration, extraClassDefinition);
 }