[TableGen] Add !setdagarg and !setdagname

- This patch proposes to add `!setdagarg` and `!setdagname` bang
  operators to produce a new DAG node after replacing the specified
  argument value/name from the given input DAG node. E.g.,
  `!setdagarg((foo 1, 2), 0, "x")` produces `(foo "x", 2)` and
  `!setdagname((foo 1:$a, 2:$b), 1, "c")` produces `(foo 1:$a, 2:$c)`.

Reviewed By: simon_tatham

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

4 files changed, 140 insertions, 41 deletions

diff --git a/llvm/lib/TableGen/Record.cpp b/llvm/lib/TableGen/Record.cpp
index ac25de1..039f5bc 100644
--- a/llvm/lib/TableGen/Record.cpp
+++ b/llvm/lib/TableGen/Record.cpp
@@ -1109,6 +1109,39 @@ std::optional<bool> BinOpInit::CompareInit(unsigned Opc, Init *LHS,
   return std::nullopt;
 }
 
+static std::optional<unsigned> getDagArgNoByKey(DagInit *Dag, Init *Key,
+                                                std::string &Error) {
+  // Accessor by index
+  if (IntInit *Idx = dyn_cast<IntInit>(Key)) {
+    int64_t Pos = Idx->getValue();
+    if (Pos < 0) {
+      // The index is negative.
+      Error =
+          (Twine("index ") + std::to_string(Pos) + Twine(" is negative")).str();
+      return std::nullopt;
+    }
+    if (Pos >= Dag->getNumArgs()) {
+      // The index is out-of-range.
+      Error = (Twine("index ") + std::to_string(Pos) +
+               " is out of range (dag has " +
+               std::to_string(Dag->getNumArgs()) + " arguments)")
+                  .str();
+      return std::nullopt;
+    }
+    return Pos;
+  }
+  assert(isa<StringInit>(Key));
+  // Accessor by name
+  StringInit *Name = dyn_cast<StringInit>(Key);
+  auto ArgNo = Dag->getArgNo(Name->getValue());
+  if (!ArgNo) {
+    // The key is not found.
+    Error = (Twine("key '") + Name->getValue() + Twine("' is not found")).str();
+    return std::nullopt;
+  }
+  return *ArgNo;
+}
+
 Init *BinOpInit::Fold(Record *CurRec) const {
   switch (getOpcode()) {
   case CONCAT: {
@@ -1279,51 +1312,19 @@ Init *BinOpInit::Fold(Record *CurRec) const {
   }
   case GETDAGARG: {
     DagInit *Dag = dyn_cast<DagInit>(LHS);
-    if (!Dag)
-      break;
+    if (Dag && isa<IntInit, StringInit>(RHS)) {
+      std::string Error;
+      auto ArgNo = getDagArgNoByKey(Dag, RHS, Error);
+      if (!ArgNo)
+        PrintFatalError(CurRec->getLoc(), "!getdagarg " + Error);
+
+      assert(*ArgNo < Dag->getNumArgs());
 
-    // Helper returning the specified argument.
-    auto getDagArgAsType = [](DagInit *Dag, unsigned Pos,
-                              RecTy *Type) -> Init * {
-      assert(Pos < Dag->getNumArgs());
-      Init *Arg = Dag->getArg(Pos);
+      Init *Arg = Dag->getArg(*ArgNo);
       if (auto *TI = dyn_cast<TypedInit>(Arg))
-        if (!TI->getType()->typeIsConvertibleTo(Type))
+        if (!TI->getType()->typeIsConvertibleTo(getType()))
           return UnsetInit::get(Dag->getRecordKeeper());
       return Arg;
-    };
-
-    // Accessor by index
-    if (IntInit *Idx = dyn_cast<IntInit>(RHS)) {
-      int64_t Pos = Idx->getValue();
-      if (Pos < 0) {
-        // The index is negative.
-        PrintFatalError(CurRec->getLoc(), Twine("!getdagarg index ") +
-                                              std::to_string(Pos) +
-                                              Twine(" is negative"));
-      }
-      if (Pos >= Dag->getNumArgs()) {
-        // The index is out-of-range.
-        PrintFatalError(CurRec->getLoc(),
-                        Twine("!getdagarg index ") + std::to_string(Pos) +
-                            " is out of range (dag has " +
-                            std::to_string(Dag->getNumArgs()) + " arguments)");
-      }
-      return getDagArgAsType(Dag, Pos, getType());
-    }
-    // Accessor by name
-    if (StringInit *Key = dyn_cast<StringInit>(RHS)) {
-      for (unsigned i = 0, e = Dag->getNumArgs(); i < e; ++i) {
-        StringInit *ArgName = Dag->getArgName(i);
-        if (!ArgName || ArgName->getValue() != Key->getValue())
-          continue;
-        // Found
-        return getDagArgAsType(Dag, i, getType());
-      }
-      // The key is not found.
-      PrintFatalError(CurRec->getLoc(), Twine("!getdagarg key '") +
-                                            Key->getValue() +
-                                            Twine("' is not found"));
     }
     break;
   }
@@ -1702,6 +1703,42 @@ Init *TernOpInit::Fold(Record *CurRec) const {
     }
     break;
   }
+
+  case SETDAGARG: {
+    DagInit *Dag = dyn_cast<DagInit>(LHS);
+    if (Dag && isa<IntInit, StringInit>(MHS)) {
+      std::string Error;
+      auto ArgNo = getDagArgNoByKey(Dag, MHS, Error);
+      if (!ArgNo)
+        PrintFatalError(CurRec->getLoc(), "!setdagarg " + Error);
+
+      assert(*ArgNo < Dag->getNumArgs());
+
+      SmallVector<Init *, 8> Args(Dag->getArgs());
+      SmallVector<StringInit *, 8> Names(Dag->getArgNames());
+      Args[*ArgNo] = RHS;
+      return DagInit::get(Dag->getOperator(), Dag->getName(), Args, Names);
+    }
+    break;
+  }
+
+  case SETDAGNAME: {
+    DagInit *Dag = dyn_cast<DagInit>(LHS);
+    if (Dag && isa<IntInit, StringInit>(MHS)) {
+      std::string Error;
+      auto ArgNo = getDagArgNoByKey(Dag, MHS, Error);
+      if (!ArgNo)
+        PrintFatalError(CurRec->getLoc(), "!setdagname " + Error);
+
+      assert(*ArgNo < Dag->getNumArgs());
+
+      SmallVector<Init *, 8> Args(Dag->getArgs());
+      SmallVector<StringInit *, 8> Names(Dag->getArgNames());
+      Names[*ArgNo] = dyn_cast<StringInit>(RHS);
+      return DagInit::get(Dag->getOperator(), Dag->getName(), Args, Names);
+    }
+    break;
+  }
   }
 
   return const_cast<TernOpInit *>(this);
@@ -1748,6 +1785,12 @@ std::string TernOpInit::getAsString() const {
   case SUBST: Result = "!subst"; break;
   case SUBSTR: Result = "!substr"; break;
   case FIND: Result = "!find"; break;
+  case SETDAGARG:
+    Result = "!setdagarg";
+    break;
+  case SETDAGNAME:
+    Result = "!setdagname";
+    break;
   }
   return (Result + "(" +
           (UnquotedLHS ? LHS->getAsUnquotedString() : LHS->getAsString()) +
@@ -2462,6 +2505,15 @@ Record *DagInit::getOperatorAsDef(ArrayRef<SMLoc> Loc) const {
   return nullptr;
 }
 
+std::optional<unsigned> DagInit::getArgNo(StringRef Name) const {
+  for (unsigned i = 0, e = getNumArgs(); i < e; ++i) {
+    StringInit *ArgName = getArgName(i);
+    if (ArgName && ArgName->getValue() == Name)
+      return i;
+  }
+  return std::nullopt;
+}
+
 Init *DagInit::resolveReferences(Resolver &R) const {
   SmallVector<Init*, 8> NewArgs;
   NewArgs.reserve(arg_size());
diff --git a/llvm/lib/TableGen/TGLexer.cpp b/llvm/lib/TableGen/TGLexer.cpp
index 84464dc..98f0e8c 100644
--- a/llvm/lib/TableGen/TGLexer.cpp
+++ b/llvm/lib/TableGen/TGLexer.cpp
@@ -594,6 +594,8 @@ tgtok::TokKind TGLexer::LexExclaim() {
           .Cases("getdagop", "getop", tgtok::XGetDagOp) // !getop is deprecated.
           .Case("getdagarg", tgtok::XGetDagArg)
           .Case("getdagname", tgtok::XGetDagName)
+          .Case("setdagarg", tgtok::XSetDagArg)
+          .Case("setdagname", tgtok::XSetDagName)
           .Case("exists", tgtok::XExists)
           .Case("tolower", tgtok::XToLower)
           .Case("toupper", tgtok::XToUpper)
diff --git a/llvm/lib/TableGen/TGLexer.h b/llvm/lib/TableGen/TGLexer.h
index dbd3853..c9bba98 100644
--- a/llvm/lib/TableGen/TGLexer.h
+++ b/llvm/lib/TableGen/TGLexer.h
@@ -129,6 +129,8 @@ enum TokKind {
   XRange,
   XGetDagArg,
   XGetDagName,
+  XSetDagArg,
+  XSetDagName,
 
   // Boolean literals.
   TrueVal,
diff --git a/llvm/lib/TableGen/TGParser.cpp b/llvm/lib/TableGen/TGParser.cpp
index e089222..067ce42 100644
--- a/llvm/lib/TableGen/TGParser.cpp
+++ b/llvm/lib/TableGen/TGParser.cpp
@@ -1769,6 +1769,8 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
     return ParseOperationForEachFilter(CurRec, ItemType);
   }
 
+  case tgtok::XSetDagArg:
+  case tgtok::XSetDagName:
   case tgtok::XDag:
   case tgtok::XIf:
   case tgtok::XSubst: { // Value ::= !ternop '(' Value ',' Value ',' Value ')'
@@ -1790,6 +1792,16 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
     case tgtok::XSubst:
       Code = TernOpInit::SUBST;
       break;
+    case tgtok::XSetDagArg:
+      Code = TernOpInit::SETDAGARG;
+      Type = DagRecTy::get(Records);
+      ItemType = nullptr;
+      break;
+    case tgtok::XSetDagName:
+      Code = TernOpInit::SETDAGNAME;
+      Type = DagRecTy::get(Records);
+      ItemType = nullptr;
+      break;
     }
     if (!consume(tgtok::l_paren)) {
       TokError("expected '(' after ternary operator");
@@ -1902,6 +1914,35 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
       Type = RHSt->getType();
       break;
     }
+    case tgtok::XSetDagArg: {
+      TypedInit *MHSt = dyn_cast<TypedInit>(MHS);
+      if (!MHSt || !isa<IntRecTy, StringRecTy>(MHSt->getType())) {
+        Error(MHSLoc, Twine("expected integer index or string name, got ") +
+                          (MHSt ? ("type '" + MHSt->getType()->getAsString())
+                                : ("'" + MHS->getAsString())) +
+                          "'");
+        return nullptr;
+      }
+      break;
+    }
+    case tgtok::XSetDagName: {
+      TypedInit *MHSt = dyn_cast<TypedInit>(MHS);
+      if (!MHSt || !isa<IntRecTy, StringRecTy>(MHSt->getType())) {
+        Error(MHSLoc, Twine("expected integer index or string name, got ") +
+                          (MHSt ? ("type '" + MHSt->getType()->getAsString())
+                                : ("'" + MHS->getAsString())) +
+                          "'");
+        return nullptr;
+      }
+      TypedInit *RHSt = dyn_cast<TypedInit>(RHS);
+      // The name could be a string or unset.
+      if (RHSt && !isa<StringRecTy>(RHSt->getType())) {
+        Error(RHSLoc, Twine("expected string or unset name, got type '") +
+                          RHSt->getType()->getAsString() + "'");
+        return nullptr;
+      }
+      break;
+    }
     }
     return (TernOpInit::get(Code, LHS, MHS, RHS, Type))->Fold(CurRec);
   }
@@ -2787,6 +2828,8 @@ Init *TGParser::ParseSimpleValue(Record *CurRec, RecTy *ItemType,
   case tgtok::XGetDagArg:
   case tgtok::XGetDagName:
   case tgtok::XSetDagOp: // Value ::= !binop '(' Value ',' Value ')'
+  case tgtok::XSetDagArg:
+  case tgtok::XSetDagName:
   case tgtok::XIf:
   case tgtok::XCond:
   case tgtok::XFoldl: