TableGen: Reimplement !foreach using the resolving mechanism

Summary:
This changes the syntax of !foreach so that the first "parameter" is
a new syntactic variable: !foreach(x, lst, expr) will define the
variable x within the scope of expr, and evaluation of the !foreach
will substitute elements of the given list (or dag) for x in expr.

Aside from leading to a nicer syntax, this allows more complex
expressions where x is deeply nested, or even constant expressions
in which x does not occur at all.

!foreach is currently not actually used anywhere in trunk, but I
plan to use it in the AMDGPU backend. If out-of-tree targets are
using it, they can adjust to the new syntax very easily.

Change-Id: Ib966694d8ab6542279d6bc358b6f4d767945a805

Reviewers: arsenm, craig.topper, tra, MartinO

Subscribers: wdng, llvm-commits, tpr

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

2 files changed, 171 insertions, 91 deletions

diff --git a/llvm/lib/TableGen/Record.cpp b/llvm/lib/TableGen/Record.cpp
index c08ae97dc57..990ef928f66 100644
--- a/llvm/lib/TableGen/Record.cpp
+++ b/llvm/lib/TableGen/Record.cpp
@@ -912,83 +912,57 @@ void TernOpInit::Profile(FoldingSetNodeID &ID) const {
   ProfileTernOpInit(ID, getOpcode(), getLHS(), getMHS(), getRHS(), getType());
 }
 
-// Evaluates operation RHSo after replacing all operands matching LHS with Arg.
-static Init *EvaluateOperation(OpInit *RHSo, Init *LHS, Init *Arg,
-                               Record *CurRec, MultiClass *CurMultiClass) {
-
-  SmallVector<Init *, 8> NewOperands;
-  NewOperands.reserve(RHSo->getNumOperands());
-  for (unsigned i = 0, e = RHSo->getNumOperands(); i < e; ++i) {
-    if (auto *RHSoo = dyn_cast<OpInit>(RHSo->getOperand(i))) {
-      if (Init *Result =
-              EvaluateOperation(RHSoo, LHS, Arg, CurRec, CurMultiClass))
-        NewOperands.push_back(Result);
-      else
-        NewOperands.push_back(RHSoo);
-    } else if (LHS->getAsString() == RHSo->getOperand(i)->getAsString()) {
-      NewOperands.push_back(Arg);
-    } else {
-      NewOperands.push_back(RHSo->getOperand(i));
-    }
+static Init *ForeachApply(Init *LHS, Init *MHSe, Init *RHS, Record *CurRec) {
+  MapResolver R(CurRec);
+  R.set(LHS, MHSe);
+  return RHS->resolveReferences(R);
+}
+
+static Init *ForeachDagApply(Init *LHS, DagInit *MHSd, Init *RHS,
+                             Record *CurRec) {
+  bool Change = false;
+  Init *Val = ForeachApply(LHS, MHSd->getOperator(), RHS, CurRec);
+  if (Val != MHSd->getOperator())
+    Change = true;
+
+  SmallVector<std::pair<Init *, StringInit *>, 8> NewArgs;
+  for (unsigned int i = 0; i < MHSd->getNumArgs(); ++i) {
+    Init *Arg = MHSd->getArg(i);
+    Init *NewArg;
+    StringInit *ArgName = MHSd->getArgName(i);
+
+    if (DagInit *Argd = dyn_cast<DagInit>(Arg))
+      NewArg = ForeachDagApply(LHS, Argd, RHS, CurRec);
+    else
+      NewArg = ForeachApply(LHS, Arg, RHS, CurRec);
+
+    NewArgs.push_back(std::make_pair(NewArg, ArgName));
+    if (Arg != NewArg)
+      Change = true;
   }
 
-  // Now run the operator and use its result as the new leaf
-  const OpInit *NewOp = RHSo->clone(NewOperands);
-  Init *NewVal = NewOp->Fold(CurRec, CurMultiClass);
-  return (NewVal != NewOp) ? NewVal : nullptr;
+  if (Change)
+    return DagInit::get(Val, nullptr, NewArgs);
+  return MHSd;
 }
 
 // Applies RHS to all elements of MHS, using LHS as a temp variable.
 static Init *ForeachHelper(Init *LHS, Init *MHS, Init *RHS, RecTy *Type,
-                           Record *CurRec, MultiClass *CurMultiClass) {
-  OpInit *RHSo = dyn_cast<OpInit>(RHS);
-
-  if (!RHSo)
-    PrintFatalError(CurRec->getLoc(), "!foreach requires an operator\n");
-
-  TypedInit *LHSt = dyn_cast<TypedInit>(LHS);
-
-  if (!LHSt)
-    PrintFatalError(CurRec->getLoc(), "!foreach requires typed variable\n");
-
-  DagInit *MHSd = dyn_cast<DagInit>(MHS);
-  if (MHSd) {
-    Init *Val = MHSd->getOperator();
-    if (Init *Result = EvaluateOperation(RHSo, LHS, Val, CurRec, CurMultiClass))
-      Val = Result;
-
-    SmallVector<std::pair<Init *, StringInit *>, 8> args;
-    for (unsigned int i = 0; i < MHSd->getNumArgs(); ++i) {
-      Init *Arg = MHSd->getArg(i);
-      StringInit *ArgName = MHSd->getArgName(i);
-      // If this is a dag, recurse
-      if (isa<DagInit>(Arg)) {
-        if (Init *Result =
-                ForeachHelper(LHS, Arg, RHSo, Type, CurRec, CurMultiClass))
-          Arg = Result;
-      } else if (Init *Result =
-                     EvaluateOperation(RHSo, LHS, Arg, CurRec, CurMultiClass)) {
-        Arg = Result;
-      }
+                           Record *CurRec) {
+  if (DagInit *MHSd = dyn_cast<DagInit>(MHS))
+    return ForeachDagApply(LHS, MHSd, RHS, CurRec);
 
-      // TODO: Process arg names
-      args.push_back(std::make_pair(Arg, ArgName));
-    }
-
-    return DagInit::get(Val, nullptr, args);
-  }
-
-  ListInit *MHSl = dyn_cast<ListInit>(MHS);
-  ListRecTy *ListType = dyn_cast<ListRecTy>(Type);
-  if (MHSl && ListType) {
+  if (ListInit *MHSl = dyn_cast<ListInit>(MHS)) {
     SmallVector<Init *, 8> NewList(MHSl->begin(), MHSl->end());
-    for (Init *&Arg : NewList) {
-      if (Init *Result =
-              EvaluateOperation(RHSo, LHS, Arg, CurRec, CurMultiClass))
-        Arg = Result;
+
+    for (Init *&Item : NewList) {
+      Init *NewItem = ForeachApply(LHS, Item, RHS, CurRec);
+      if (NewItem != Item)
+        Item = NewItem;
     }
-    return ListInit::get(NewList, ListType->getElementType());
+    return ListInit::get(NewList, cast<ListRecTy>(Type)->getElementType());
   }
+
   return nullptr;
 }
 
@@ -1038,8 +1012,7 @@ Init *TernOpInit::Fold(Record *CurRec, MultiClass *CurMultiClass) const {
   }
 
   case FOREACH: {
-    if (Init *Result =
-            ForeachHelper(LHS, MHS, RHS, getType(), CurRec, CurMultiClass))
+    if (Init *Result = ForeachHelper(LHS, MHS, RHS, getType(), CurRec))
       return Result;
     break;
   }
@@ -1081,7 +1054,15 @@ Init *TernOpInit::resolveReferences(Resolver &R) const {
   }
 
   Init *mhs = MHS->resolveReferences(R);
-  Init *rhs = RHS->resolveReferences(R);
+  Init *rhs;
+
+  if (getOpcode() == FOREACH) {
+    ShadowResolver SR(R);
+    SR.addShadow(lhs);
+    rhs = RHS->resolveReferences(SR);
+  } else {
+    rhs = RHS->resolveReferences(R);
+  }
 
   if (LHS != lhs || MHS != mhs || RHS != rhs)
     return (TernOpInit::get(getOpcode(), lhs, mhs, rhs, getType()))
@@ -1821,6 +1802,24 @@ Init *llvm::QualifyName(Record &CurRec, MultiClass *CurMultiClass,
   return NewName;
 }
 
+Init *MapResolver::resolve(Init *VarName) {
+  auto It = Map.find(VarName);
+  if (It == Map.end())
+    return nullptr;
+
+  Init *I = It->second.V;
+
+  if (!It->second.Resolved && Map.size() > 1) {
+    // Resolve mutual references among the mapped variables, but prevent
+    // infinite recursion.
+    Map.erase(It);
+    I = I->resolveReferences(*this);
+    Map[VarName] = {I, true};
+  }
+
+  return I;
+}
+
 Init *RecordResolver::resolve(Init *VarName) {
   Init *Val = Cache.lookup(VarName);
   if (Val)
diff --git a/llvm/lib/TableGen/TGParser.cpp b/llvm/lib/TableGen/TGParser.cpp
index 3ef65938282..b42580c9cda 100644
--- a/llvm/lib/TableGen/TGParser.cpp
+++ b/llvm/lib/TableGen/TGParser.cpp
@@ -985,8 +985,109 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
     return nullptr;
   }
 
+  case tgtok::XForEach: { // Value ::= !foreach '(' Id ',' Value ',' Value ')'
+    SMLoc OpLoc = Lex.getLoc();
+    Lex.Lex(); // eat the operation
+    if (Lex.getCode() != tgtok::l_paren) {
+      TokError("expected '(' after !foreach");
+      return nullptr;
+    }
+
+    if (Lex.Lex() != tgtok::Id) { // eat the '('
+      TokError("first argument of !foreach must be an identifier");
+      return nullptr;
+    }
+
+    Init *LHS = StringInit::get(Lex.getCurStrVal());
+
+    if (CurRec->getValue(LHS)) {
+      TokError((Twine("iteration variable '") + LHS->getAsString() +
+                "' already defined")
+                   .str());
+      return nullptr;
+    }
+
+    if (Lex.Lex() != tgtok::comma) { // eat the id
+      TokError("expected ',' in ternary operator");
+      return nullptr;
+    }
+    Lex.Lex();  // eat the ','
+
+    Init *MHS = ParseValue(CurRec);
+    if (!MHS)
+      return nullptr;
+
+    if (Lex.getCode() != tgtok::comma) {
+      TokError("expected ',' in ternary operator");
+      return nullptr;
+    }
+    Lex.Lex();  // eat the ','
+
+    TypedInit *MHSt = dyn_cast<TypedInit>(MHS);
+    if (!MHSt) {
+      TokError("could not get type of !foreach input");
+      return nullptr;
+    }
+
+    RecTy *InEltType = nullptr;
+    RecTy *OutEltType = nullptr;
+    bool IsDAG = false;
+
+    if (ListRecTy *InListTy = dyn_cast<ListRecTy>(MHSt->getType())) {
+      InEltType = InListTy->getElementType();
+      if (ItemType) {
+        if (ListRecTy *OutListTy = dyn_cast<ListRecTy>(ItemType)) {
+          OutEltType = OutListTy->getElementType();
+        } else {
+          Error(OpLoc,
+                "expected value of type '" + Twine(ItemType->getAsString()) +
+                "', but got !foreach of list type");
+          return nullptr;
+        }
+      }
+    } else if (DagRecTy *InDagTy = dyn_cast<DagRecTy>(MHSt->getType())) {
+      InEltType = InDagTy;
+      if (ItemType && !isa<DagRecTy>(ItemType)) {
+        Error(OpLoc,
+              "expected value of type '" + Twine(ItemType->getAsString()) +
+              "', but got !foreach of dag type");
+        return nullptr;
+      }
+      IsDAG = true;
+    } else {
+      TokError("!foreach must have list or dag input");
+      return nullptr;
+    }
+
+    CurRec->addValue(RecordVal(LHS, InEltType, false));
+    Init *RHS = ParseValue(CurRec, OutEltType);
+    CurRec->removeValue(LHS);
+    if (!RHS)
+      return nullptr;
+
+    if (Lex.getCode() != tgtok::r_paren) {
+      TokError("expected ')' in binary operator");
+      return nullptr;
+    }
+    Lex.Lex();  // eat the ')'
+
+    RecTy *OutType;
+    if (IsDAG) {
+      OutType = InEltType;
+    } else {
+      TypedInit *RHSt = dyn_cast<TypedInit>(RHS);
+      if (!RHSt) {
+        TokError("could not get type of !foreach result");
+        return nullptr;
+      }
+      OutType = RHSt->getType()->getListTy();
+    }
+
+    return (TernOpInit::get(TernOpInit::FOREACH, LHS, MHS, RHS, OutType))
+               ->Fold(CurRec, CurMultiClass);
+  }
+
   case tgtok::XIf:
-  case tgtok::XForEach:
   case tgtok::XSubst: {  // Value ::= !ternop '(' Value ',' Value ',' Value ')'
     TernOpInit::TernaryOp Code;
     RecTy *Type = nullptr;
@@ -998,9 +1099,6 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
     case tgtok::XIf:
       Code = TernOpInit::IF;
       break;
-    case tgtok::XForEach:
-      Code = TernOpInit::FOREACH;
-      break;
     case tgtok::XSubst:
       Code = TernOpInit::SUBST;
       break;
@@ -1081,23 +1179,6 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
       }
       break;
     }
-    case tgtok::XForEach: {
-      TypedInit *MHSt = dyn_cast<TypedInit>(MHS);
-      if (!MHSt) {
-        TokError("could not get type for !foreach");
-        return nullptr;
-      }
-      Type = MHSt->getType();
-      if (isa<ListRecTy>(Type)) {
-        TypedInit *RHSt = dyn_cast<TypedInit>(RHS);
-        if (!RHSt) {
-          TokError("could not get type of !foreach list elements");
-          return nullptr;
-        }
-        Type = RHSt->getType()->getListTy();
-      }
-      break;
-    }
     case tgtok::XSubst: {
       TypedInit *RHSt = dyn_cast<TypedInit>(RHS);
       if (!RHSt) {