[ODS] Fix operation argument population to avoid crash

The `Operator` class keeps an `arguments` field, which contains pointers
to `operands` and `attributes` elements. Thus it must be populated after
`operands` and `attributes` are finalized so to have stable pointers.
SmallVector may re-allocate when still having new elements added, which
will invalidate pointers.

PiperOrigin-RevId: 280466896

2 files changed, 34 insertions, 6 deletions

diff --git a/mlir/lib/TableGen/Operator.cpp b/mlir/lib/TableGen/Operator.cpp
index 7d926d98b5c..8afffd03fcb 100644
--- a/mlir/lib/TableGen/Operator.cpp
+++ b/mlir/lib/TableGen/Operator.cpp
@@ -27,6 +27,8 @@
 #include "llvm/TableGen/Error.h"
 #include "llvm/TableGen/Record.h"
 
+#define DEBUG_TYPE "mlir-tblgen-operator"
+
 using namespace mlir;
 
 using llvm::DagInit;
@@ -205,12 +207,11 @@ void tblgen::Operator::populateOpStructure() {
   auto derivedAttrClass = recordKeeper.getClass("DerivedAttr");
   numNativeAttributes = 0;
 
-  // The argument ordering is operands, native attributes, derived
-  // attributes.
   DagInit *argumentValues = def.getValueAsDag("arguments");
-  unsigned i = 0;
+  unsigned numArgs = argumentValues->getNumArgs();
+
   // Handle operands and native attributes.
-  for (unsigned e = argumentValues->getNumArgs(); i != e; ++i) {
+  for (unsigned i = 0; i != numArgs; ++i) {
     auto arg = argumentValues->getArg(i);
     auto givenName = argumentValues->getArgNameStr(i);
     auto argDefInit = dyn_cast<DefInit>(arg);
@@ -222,7 +223,6 @@ void tblgen::Operator::populateOpStructure() {
     if (argDef->isSubClassOf(typeConstraintClass)) {
       operands.push_back(
           NamedTypeConstraint{givenName, TypeConstraint(argDefInit)});
-      arguments.emplace_back(&operands.back());
     } else if (argDef->isSubClassOf(attrClass)) {
       if (givenName.empty())
         PrintFatalError(argDef->getLoc(), "attributes must be named");
@@ -230,7 +230,6 @@ void tblgen::Operator::populateOpStructure() {
         PrintFatalError(argDef->getLoc(),
                         "derived attributes not allowed in argument list");
       attributes.push_back({givenName, Attribute(argDef)});
-      arguments.emplace_back(&attributes.back());
       ++numNativeAttributes;
     } else {
       PrintFatalError(def.getLoc(), "unexpected def type; only defs deriving "
@@ -258,6 +257,22 @@ void tblgen::Operator::populateOpStructure() {
     }
   }
 
+  // Populate `arguments`. This must happen after we've finalized `operands` and
+  // `attributes` because we will put their elements' pointers in `arguments`.
+  // SmallVector may perform re-allocation under the hood when adding new
+  // elements.
+  int operandIndex = 0, attrIndex = 0;
+  for (unsigned i = 0; i != numArgs; ++i) {
+    Record *argDef = dyn_cast<DefInit>(argumentValues->getArg(i))->getDef();
+
+    if (argDef->isSubClassOf(typeConstraintClass)) {
+      arguments.emplace_back(&operands[operandIndex++]);
+    } else {
+      assert(argDef->isSubClassOf(attrClass));
+      arguments.emplace_back(&attributes[attrIndex++]);
+    }
+  }
+
   auto *resultsDag = def.getValueAsDag("results");
   auto *outsOp = dyn_cast<DefInit>(resultsDag->getOperator());
   if (!outsOp || outsOp->getDef()->getName() != "outs") {
@@ -298,6 +313,8 @@ void tblgen::Operator::populateOpStructure() {
     }
     regions.push_back({name, Region(regionInit->getDef())});
   }
+
+  LLVM_DEBUG(print(llvm::dbgs()));
 }
 
 ArrayRef<llvm::SMLoc> tblgen::Operator::getLoc() const { return def.getLoc(); }
@@ -317,3 +334,13 @@ bool tblgen::Operator::hasSummary() const {
 StringRef tblgen::Operator::getSummary() const {
   return def.getValueAsString("summary");
 }
+
+void tblgen::Operator::print(llvm::raw_ostream &os) const {
+  os << "op '" << getOperationName() << "'\n";
+  for (Argument arg : arguments) {
+    if (auto *attr = arg.dyn_cast<NamedAttribute *>())
+      os << "[attribute] " << attr->name << '\n';
+    else
+      os << "[operand] " << arg.get<NamedTypeConstraint *>()->name << '\n';
+  }
+}
diff --git a/mlir/lib/TableGen/Pattern.cpp b/mlir/lib/TableGen/Pattern.cpp
index ddec0bacac4..d3c1dddd21e 100644
--- a/mlir/lib/TableGen/Pattern.cpp
+++ b/mlir/lib/TableGen/Pattern.cpp
@@ -211,6 +211,7 @@ int tblgen::SymbolInfoMap::SymbolInfo::getStaticValueCount() const {
 
 std::string
 tblgen::SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
+  LLVM_DEBUG(llvm::dbgs() << "getVarDecl for '" << name << "': ");
   switch (kind) {
   case Kind::Attr: {
     auto type =