[Mlir-commits] [mlir] [mlir][tblgen] Add custom parsing and printing within struct (PR #133939)

[River Riddle]

================
@@ -338,29 +387,56 @@ class DirectiveElementBase : public DirectiveElement {
   }
 };
 
+/// Base class for a directive that contains references to elements of type `T`
+/// in a vector.
+template <DirectiveElement::Kind DirectiveKind, typename T>
+class VectorDirectiveBase : public DirectiveElementBase<DirectiveKind> {
+public:
+  using Base = VectorDirectiveBase<DirectiveKind, T>;
+
+  VectorDirectiveBase(std::vector<T> &&elems) : elems(std::move(elems)) {}
+
+  /// Get the elements contained in this directive.
+  ArrayRef<T> getElements() const { return elems; }
+
+  /// Get the number of elements.
+  unsigned getNumElements() const { return elems.size(); }
+
+  /// Take all of the elements from this directive.
+  std::vector<T> takeElements() { return std::move(elems); }
+
+protected:
+  /// The elements captured by this directive.
+  std::vector<T> elems;
+};
+
 /// This class represents a custom format directive that is implemented by the
 /// user in C++. The directive accepts a list of arguments that is passed to the
 /// C++ function.
-class CustomDirective : public DirectiveElementBase<DirectiveElement::Custom> {
+class CustomDirective
+    : public VectorDirectiveBase<DirectiveElement::Custom, FormatElement *> {
 public:
+  using Base::Base;
   /// Create a custom directive with a name and list of arguments.
   CustomDirective(StringRef name, std::vector<FormatElement *> &&arguments)
-      : name(name), arguments(std::move(arguments)) {}
+      : Base(std::move(arguments)), name(name) {}
 
   /// Get the custom directive name.
   StringRef getName() const { return name; }
 
-  /// Get the arguments to the custom directive.
-  ArrayRef<FormatElement *> getArguments() const { return arguments; }
+  FailureOr<ParameterElement *> getFrontAsParam() const {
+    if (getNumElements() != 1)
+      return failure();
+    ParameterElement *param = dyn_cast<ParameterElement>(getElements()[0]);
+    if (param == nullptr)
+      return failure();
+    return param;
+  }
 
----------------
River707 wrote:

Can this be moved to AttrOrTypeFormatGen.cpp? We shouldn't expose logic specific to operations/attrtype formats here. Alternatively, this method could stay here, but it should be templated (to avoid specializing for ParameterElement):

```
template<typename T>
FailureOr<T *> getFrontAs() const {
```

https://github.com/llvm/llvm-project/pull/133939