[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

llvm/docs/TableGen/ProgRef.rst

@@ -225,10 +225,10 @@ TableGen provides "bang operators" that have a wide variety of uses:
                : !getdagname  !getdagop    !gt          !head        !if
                : !interleave  !isa         !le          !listconcat  !listremove
                : !listsplat   !logtwo      !lt          !mul         !ne
-               : !not         !or          !range       !setdagop    !shl
-               : !size        !sra         !srl         !strconcat   !sub
-               : !subst       !substr      !tail        !tolower     !toupper
-               : !xor
+               : !not         !or          !range       !setdagarg   !setdagname
+               : !setdagop    !shl         !size        !sra         !srl
+               : !strconcat   !sub         !subst       !substr      !tail
+               : !tolower     !toupper     !xor
 
 The ``!cond`` operator has a slightly different
 syntax compared to other bang operators, so it is defined separately:
@@ -1370,7 +1370,7 @@ DAG.
 
 The following bang operators are useful for working with DAGs:
 ``!con``, ``!dag``, ``!empty``, ``!foreach``, ``!getdagarg``, ``!getdagname``,
-``!getdagop``, ``!setdagop``, ``!size``.
+``!getdagop``, ``!setdagarg``, ``!setdagname``, ``!setdagop``, ``!size``.
 
 Defvar in a record body
 -----------------------
@@ -1821,6 +1821,16 @@ and non-0 as true.
 ``!range(``\ *list*\ ``)``
     Equivalent to ``!range(0, !size(list))``.
 
+``!setdagarg(``\ *dag*\ ``,``\ *key*\ ``,``\ *arg*\ ``)``
+    This operator produces a DAG node with the same operator and arguments as
+    *dag*, but replacing the value of the argument specified by the *key* with
+    *arg*. That *key* could be either an integer index or a string name.
+
+``!setdagname(``\ *dag*\ ``,``\ *key*\ ``,``\ *name*\ ``)``
+    This operator produces a DAG node with the same operator and arguments as
+    *dag*, but replacing the name of the argument specified by the *key* with
+    *name*. That *key* could be either an integer index or a string name.
+
 ``!setdagop(``\ *dag*\ ``,`` *op*\ ``)``
     This operator produces a DAG node with the same arguments as *dag*, but with its
     operator replaced with *op*.

llvm/include/llvm/TableGen/Record.h

@@ -918,7 +918,17 @@ class BinOpInit : public OpInit, public FoldingSetNode {
 /// !op (X, Y, Z) - Combine two inits.
 class TernOpInit : public OpInit, public FoldingSetNode {
 public:
-  enum TernaryOp : uint8_t { SUBST, FOREACH, FILTER, IF, DAG, SUBSTR, FIND };
+  enum TernaryOp : uint8_t {
+    SUBST,
+    FOREACH,
+    FILTER,
+    IF,
+    DAG,
+    SUBSTR,
+    FIND,
+    SETDAGARG,
+    SETDAGNAME,
+  };
 
 private:
   Init *LHS, *MHS, *RHS;
@@ -1407,6 +1417,10 @@ class DagInit final : public TypedInit, public FoldingSetNode,
     return getTrailingObjects<Init *>()[Num];
   }
 
+  /// This method looks up the specified argument name and returns its argument
+  /// number or std::nullopt if that argument name does not exist.
+  std::optional<unsigned> getArgNo(StringRef Name) const;
+
   StringInit *getArgName(unsigned Num) const {
     assert(Num < NumArgNames && "Arg number out of range!");
     return getTrailingObjects<StringInit *>()[Num];

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());

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)

llvm/lib/TableGen/TGLexer.h

@@ -129,6 +129,8 @@ enum TokKind {
   XRange,
   XGetDagArg,
   XGetDagName,
+  XSetDagArg,
+  XSetDagName,
 
   // Boolean literals.
   TrueVal,

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:

llvm/test/TableGen/getsetop.td

@@ -5,6 +5,9 @@
 // RUN: not llvm-tblgen -DERROR4 %s 2>&1 | FileCheck --check-prefix=ERROR4 %s
 // RUN: not llvm-tblgen -DERROR5 %s 2>&1 | FileCheck --check-prefix=ERROR5 %s
 // RUN: not llvm-tblgen -DERROR6 %s 2>&1 | FileCheck --check-prefix=ERROR6 %s
+// RUN: not llvm-tblgen -DERROR7 %s 2>&1 | FileCheck --check-prefix=ERROR7 %s
+// RUN: not llvm-tblgen -DERROR8 %s 2>&1 | FileCheck --check-prefix=ERROR8 %s
+// RUN: not llvm-tblgen -DERROR9 %s 2>&1 | FileCheck --check-prefix=ERROR9 %s
 
 // !setop and !getop are deprecated in favor of !setdagop and !getdagop.
 // Two tests retain the old names just to be sure they are still supported.
@@ -114,4 +117,34 @@ def test {
   dag in6 = (foo alice:$a, bob:$b);
   // CHECK: Base base = bob;
   Base base = !getdagarg<Base>(in6, 1);
+
+  // CHECK: dag orig_set_val = (foo 1, 2:$a, "val":$b);
+  dag orig_set_val = !setdagarg(orig, 2, "val");
+  // CHECK: dag orig_set_val_by_name = (foo 1, 2:$a, "aval":$b);
+  dag orig_set_val_by_name = !setdagarg(orig, "b", "aval");
+  // CHECK: dag orig_set_dag_val = (foo 1, 2:$a, (bar foo:$p, qux:$q):$b);
+  dag orig_set_dag_val = !setdagarg(orig, "b", (bar foo:$p, qux:$q));
+  // CHECK: dag orig_clr_val = (foo 1, ?:$a, ?:$b);
+  dag orig_clr_val = !setdagarg(orig, "a", ?);
+  // CHECK: dag orig_set_name = (foo 1:$c, 2:$a, ?:$b);
+  dag orig_set_name = !setdagname(orig, 0, "c");
+  // CHECK: dag orig_clr_name = (foo 1, 2, ?:$b);
+  dag orig_clr_name = !setdagname(orig, 1, ?);
+  // CHECK: dag orig_rename = (foo 1, 2:$x, ?:$y);
+  dag orig_rename = !setdagname(!setdagname(orig, "a", "x"), "b", "y");
+
+#ifdef ERROR7
+  // ERROR7: error: !setdagarg index -1 is negative
+  dag orig_negative = !setdagarg(orig, -1, "val");
+#endif
+
+#ifdef ERROR8
+  // ERROR8: error: !setdagarg index 3 is out of range (dag has 3 arguments)
+  dag orig_out_of_range = !setdagarg(orig, 3, "val");
+#endif
+
+#ifdef ERROR9
+  // ERROR9: error: expected integer index or string name, got type 'Base'
+  dag orig_out_of_range = !setdagarg(orig, foo, (foo qux:$a));
+#endif
 }