[TableGen] Add bang-operators !getop and !setop.

Summary:
These allow you to get and set the operator of a dag node, without
affecting its list of arguments.

`!getop` is slightly fiddly because in many contexts you need its
return value to have a static type more specific than 'any record'. It
works to say `!cast<BaseClass>(!getop(...))`, but it's cumbersome, so
I made `!getop` take an optional type suffix itself, so that can be
written as the shorter `!getop<BaseClass>(...)`.

Reviewers: hfinkel, nhaehnle

Reviewed By: nhaehnle

Subscribers: hiraditya, llvm-commits

Tags: #llvm

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

docs/TableGen/LangIntro.rst

@@ -183,6 +183,34 @@ supported include:
     Example: !dag(op, [a1, a2, ?], ["name1", "name2", "name3"]) results in
     (op a1:$name1, a2:$name2, ?:$name3).
 
+``!setop(dag, op)``
+    Return a DAG node with the same arguments as ``dag``, but with its
+    operator replaced with ``op``.
+
+    Example: ``!setop((foo 1, 2), bar)`` results in ``(bar 1, 2)``.
+
+``!getop(dag)``
+
+``!getop<type>(dag)``
+    Return the operator of the given DAG node.
+    Example: ``!getop((foo 1, 2))`` results in ``foo``.
+
+    The result of ``!getop`` can be used directly in a context where
+    any record value at all is acceptable (typically placing it into
+    another dag value). But in other contexts, it must be explicitly
+    cast to a particular class type. The ``!getop<type>`` syntax is
+    provided to make this easy.
+
+    For example, to assign the result to a class-typed value, you
+    could write either of these:
+      ``BaseClass b = !getop<BaseClass>(someDag);``
+
+      ``BaseClass b = !cast<BaseClass>(!getop(someDag));``
+
+    But to build a new dag node reusing the operator from another, no
+    cast is necessary:
+      ``dag d = !dag(!getop(someDag), args, names);``
+
 ``!listconcat(a, b, ...)``
     A list value that is the result of concatenating the 'a' and 'b' lists.
     The lists must have the same element type.

docs/TableGen/LangRef.rst

@@ -100,7 +100,8 @@ wide variety of meanings:
                :!or     !empty   !subst   !foreach   !strconcat
                :!cast   !listconcat       !size      !foldl
                :!isa    !dag     !le      !lt        !ge
-               :!gt     !ne      !mul     !listsplat
+               :!gt     !ne      !mul     !listsplat !setop
+               :!getop
 
 TableGen also has !cond operator that needs a slightly different
 syntax compared to other "bang operators":

include/llvm/TableGen/Record.h

@@ -751,7 +751,7 @@ public:
 ///
 class UnOpInit : public OpInit, public FoldingSetNode {
 public:
-  enum UnaryOp : uint8_t { CAST, HEAD, TAIL, SIZE, EMPTY };
+  enum UnaryOp : uint8_t { CAST, HEAD, TAIL, SIZE, EMPTY, GETOP };
 
 private:
   Init *LHS;
@@ -802,7 +802,7 @@ class BinOpInit : public OpInit, public FoldingSetNode {
 public:
   enum BinaryOp : uint8_t { ADD, MUL, AND, OR, SHL, SRA, SRL, LISTCONCAT,
                             LISTSPLAT, STRCONCAT, CONCAT, EQ, NE, LE, LT, GE,
-                            GT };
+                            GT, SETOP };
 
 private:
   Init *LHS, *RHS;

lib/TableGen/Record.cpp

@@ -788,6 +788,21 @@ Init *UnOpInit::Fold(Record *CurRec, bool IsFinal) const {
     if (StringInit *LHSs = dyn_cast<StringInit>(LHS))
       return IntInit::get(LHSs->getValue().empty());
     break;
+
+  case GETOP:
+    if (DagInit *Dag = dyn_cast<DagInit>(LHS)) {
+      DefInit *DI = DefInit::get(Dag->getOperatorAsDef({}));
+      if (!DI->getType()->typeIsA(getType())) {
+        PrintFatalError(CurRec->getLoc(),
+                        Twine("Expected type '") +
+                        getType()->getAsString() + "', got '" +
+                        DI->getType()->getAsString() + "' in: " +
+                        getAsString() + "\n");
+      } else {
+        return DI;
+      }
+    }
+    break;
   }
   return const_cast<UnOpInit *>(this);
 }
@@ -809,6 +824,7 @@ std::string UnOpInit::getAsString() const {
   case TAIL: Result = "!tail"; break;
   case SIZE: Result = "!size"; break;
   case EMPTY: Result = "!empty"; break;
+  case GETOP: Result = "!getop"; break;
   }
   return Result + "(" + LHS->getAsString() + ")";
 }
@@ -980,6 +996,20 @@ Init *BinOpInit::Fold(Record *CurRec) const {
 
     break;
   }
+  case SETOP: {
+    DagInit *Dag = dyn_cast<DagInit>(LHS);
+    DefInit *Op = dyn_cast<DefInit>(RHS);
+    if (Dag && Op) {
+      SmallVector<Init*, 8> Args;
+      SmallVector<StringInit*, 8> ArgNames;
+      for (unsigned i = 0, e = Dag->getNumArgs(); i != e; ++i) {
+        Args.push_back(Dag->getArg(i));
+        ArgNames.push_back(Dag->getArgName(i));
+      }
+      return DagInit::get(Op, nullptr, Args, ArgNames);
+    }
+    break;
+  }
   case ADD:
   case MUL:
   case AND:
@@ -1042,6 +1072,7 @@ std::string BinOpInit::getAsString() const {
   case LISTCONCAT: Result = "!listconcat"; break;
   case LISTSPLAT: Result = "!listsplat"; break;
   case STRCONCAT: Result = "!strconcat"; break;
+  case SETOP: Result = "!setop"; break;
   }
   return Result + "(" + LHS->getAsString() + ", " + RHS->getAsString() + ")";
 }

lib/TableGen/TGLexer.cpp

@@ -559,6 +559,8 @@ tgtok::TokKind TGLexer::LexExclaim() {
     .Case("listconcat", tgtok::XListConcat)
     .Case("listsplat", tgtok::XListSplat)
     .Case("strconcat", tgtok::XStrConcat)
+    .Case("setop", tgtok::XSetOp)
+    .Case("getop", tgtok::XGetOp)
     .Default(tgtok::Error);
 
   return Kind != tgtok::Error ? Kind : ReturnError(Start-1, "Unknown operator");

lib/TableGen/TGLexer.h

@@ -51,7 +51,7 @@ namespace tgtok {
     // !keywords.
     XConcat, XADD, XMUL, XAND, XOR, XSRA, XSRL, XSHL, XListConcat, XListSplat,
     XStrConcat, XCast, XSubst, XForEach, XFoldl, XHead, XTail, XSize, XEmpty,
-    XIf, XCond, XEq, XIsA, XDag, XNe, XLe, XLt, XGe, XGt,
+    XIf, XCond, XEq, XIsA, XDag, XNe, XLe, XLt, XGe, XGt, XSetOp, XGetOp,
 
     // Integer value.
     IntVal,

lib/TableGen/TGParser.cpp

@@ -905,7 +905,8 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
   case tgtok::XTail:
   case tgtok::XSize:
   case tgtok::XEmpty:
-  case tgtok::XCast: {  // Value ::= !unop '(' Value ')'
+  case tgtok::XCast:
+  case tgtok::XGetOp: {  // Value ::= !unop '(' Value ')'
     UnOpInit::UnaryOp Code;
     RecTy *Type = nullptr;
 
@@ -941,6 +942,28 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
       Code = UnOpInit::EMPTY;
       Type = IntRecTy::get();
       break;
+    case tgtok::XGetOp:
+      Lex.Lex();  // eat the operation
+      if (Lex.getCode() == tgtok::less) {
+        // Parse an optional type suffix, so that you can say
+        // !getop<BaseClass>(someDag) as a shorthand for
+        // !cast<BaseClass>(!getop(someDag)).
+        Type = ParseOperatorType();
+
+        if (!Type) {
+          TokError("did not get type for unary operator");
+          return nullptr;
+        }
+
+        if (!isa<RecordRecTy>(Type)) {
+          TokError("type for !getop must be a record type");
+          // but keep parsing, to consume the operand
+        }
+      } else {
+        Type = RecordRecTy::get({});
+      }
+      Code = UnOpInit::GETOP;
+      break;
     }
     if (Lex.getCode() != tgtok::l_paren) {
       TokError("expected '(' after unary operator");
@@ -1055,7 +1078,8 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
   case tgtok::XGt:
   case tgtok::XListConcat:
   case tgtok::XListSplat:
-  case tgtok::XStrConcat: {  // Value ::= !binop '(' Value ',' Value ')'
+  case tgtok::XStrConcat:
+  case tgtok::XSetOp: {  // Value ::= !binop '(' Value ',' Value ')'
     tgtok::TokKind OpTok = Lex.getCode();
     SMLoc OpLoc = Lex.getLoc();
     Lex.Lex();  // eat the operation
@@ -1080,6 +1104,7 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
     case tgtok::XListConcat: Code = BinOpInit::LISTCONCAT; break;
     case tgtok::XListSplat: Code = BinOpInit::LISTSPLAT; break;
     case tgtok::XStrConcat: Code = BinOpInit::STRCONCAT; break;
+    case tgtok::XSetOp: Code = BinOpInit::SETOP; break;
     }
 
     RecTy *Type = nullptr;
@@ -1088,6 +1113,7 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
     default:
       llvm_unreachable("Unhandled code!");
     case tgtok::XConcat:
+    case tgtok::XSetOp:
       Type = DagRecTy::get();
       ArgType = DagRecTy::get();
       break;
@@ -1146,7 +1172,6 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
       InitList.push_back(ParseValue(CurRec, ArgType));
       if (!InitList.back()) return nullptr;
 
-      // All BinOps require their arguments to be of compatible types.
       RecTy *ListType = cast<TypedInit>(InitList.back())->getType();
       if (!ArgType) {
         ArgType = ListType;
@@ -1212,6 +1237,18 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
           ArgType = Resolved;
       }
 
+      // Deal with BinOps whose arguments have different types, by
+      // rewriting ArgType in between them.
+      switch (Code) {
+        case BinOpInit::SETOP:
+          // After parsing the first dag argument, switch to expecting
+          // a record, with no restriction on its superclasses.
+          ArgType = RecordRecTy::get({});
+          break;
+        default:
+          break;
+      }
+
       if (Lex.getCode() != tgtok::comma)
         break;
       Lex.Lex();  // eat the ','
@@ -2025,7 +2062,7 @@ Init *TGParser::ParseSimpleValue(Record *CurRec, RecTy *ItemType,
   case tgtok::l_paren: {         // Value ::= '(' IDValue DagArgList ')'
     Lex.Lex();   // eat the '('
     if (Lex.getCode() != tgtok::Id && Lex.getCode() != tgtok::XCast &&
-        Lex.getCode() != tgtok::question) {
+        Lex.getCode() != tgtok::question && Lex.getCode() != tgtok::XGetOp) {
       TokError("expected identifier in dag init");
       return nullptr;
     }
@@ -2063,7 +2100,8 @@ Init *TGParser::ParseSimpleValue(Record *CurRec, RecTy *ItemType,
   case tgtok::XTail:
   case tgtok::XSize:
   case tgtok::XEmpty:
-  case tgtok::XCast:  // Value ::= !unop '(' Value ')'
+  case tgtok::XCast:
+  case tgtok::XGetOp:  // Value ::= !unop '(' Value ')'
   case tgtok::XIsA:
   case tgtok::XConcat:
   case tgtok::XDag:
@@ -2082,7 +2120,8 @@ Init *TGParser::ParseSimpleValue(Record *CurRec, RecTy *ItemType,
   case tgtok::XGt:
   case tgtok::XListConcat:
   case tgtok::XListSplat:
-  case tgtok::XStrConcat:   // Value ::= !binop '(' Value ',' Value ')'
+  case tgtok::XStrConcat:
+  case tgtok::XSetOp:   // Value ::= !binop '(' Value ',' Value ')'
   case tgtok::XIf:
   case tgtok::XCond:
   case tgtok::XFoldl:

test/TableGen/getsetop.td

@@ -0,0 +1,61 @@
+// RUN: llvm-tblgen %s | FileCheck %s
+// RUN: not llvm-tblgen -DERROR1 %s 2>&1 | FileCheck --check-prefix=ERROR1 %s
+// RUN: not llvm-tblgen -DERROR2 %s 2>&1 | FileCheck --check-prefix=ERROR2 %s
+// RUN: not llvm-tblgen -DERROR3 %s 2>&1 | FileCheck --check-prefix=ERROR3 %s
+
+class Base;
+class OtherBase;
+
+def foo: Base;
+def bar: Base;
+def qux: OtherBase;
+
+def test {
+  dag orig = (foo 1, 2:$a, $b);
+  dag another = (qux "hello", $world);
+
+  // CHECK: dag replaceWithBar = (bar 1, 2:$a, ?:$b);
+  dag replaceWithBar = !setop(orig, bar);
+
+  // CHECK: dag replaceWithBaz = (qux 1, 2:$a, ?:$b);
+  dag replaceWithBaz = !setop(orig, qux);
+
+  // CHECK: Base getopWithCast = foo;
+  Base getopWithCast = !getop<Base>(orig);
+
+  // CHECK: dag getopToSetop = (foo "hello", ?:$world);
+  dag getopToSetop = !setop(another, !getop(orig));
+
+  // CHECK: dag getopToBangDag = (foo 1:$a, 2:$b, 3:$c);
+  dag getopToBangDag = !dag(!getop(orig), [1, 2, 3], ["a", "b", "c"]);
+
+  // CHECK: dag getopToDagInit = (foo "it worked");
+  dag getopToDagInit = (!getop(orig) "it worked");
+
+#ifdef ERROR1
+  // !getop(...) has a static type of 'any record at all, with no
+  // required superclasses'. That's too general to use in an
+  // assignment whose LHS demands an instance of Base, so we expect a
+  // static (parse-time) type-checking error.
+
+  // ERROR1: error: Value 'noCast' of type 'Base' is incompatible with initializer '!getop(orig)' of type '{}'
+  Base noCast = !getop(orig);
+#endif
+
+#ifdef ERROR2
+  // Here, we expect a _dynamic_ type error, when it turns out at
+  // evaluation time that the operator of 'another' is a record that
+  // isn't an instance of the specified base class.
+
+  // ERROR2: error: Expected type 'Base', got 'OtherBase' in: !getop((qux "hello", ?:$world))
+  Base badCast = !getop<Base>(another);
+#endif
+
+#ifdef ERROR3
+  // Obviously, you shouldn't be able to give any type to !getop that
+  // isn't a class type.
+
+  // ERROR3: error: type for !getop must be a record type
+  int ridiculousCast = !getop<int>(orig);
+#endif
+}