Skip to content

Commit 66c8dce

Browse files
jurahuljurahul
authored
[TableGen] Add a !listflatten operator to TableGen (#109346)
Add a !listflatten operator that will transform an input list of type `list<list<X>>` to `list<X>` by concatenating elements of the constituent lists of the input argument.
1 parent 3ec5e74 commit 66c8dce

File tree

8 files changed

+110
-12
lines changed

8 files changed

+110
-12
lines changed

‎llvm/docs/TableGen/ProgRef.rst‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -223,12 +223,12 @@ TableGen provides "bang operators" that have a wide variety of uses:
223223
: !div !empty !eq !exists !filter
224224
: !find !foldl !foreach !ge !getdagarg
225225
: !getdagname !getdagop !gt !head !if
226-
: !interleave !isa !le !listconcat !listremove
227-
: !listsplat !logtwo !lt !mul !ne
228-
: !not !or !range !repr !setdagarg
229-
: !setdagname !setdagop !shl !size !sra
230-
: !srl !strconcat !sub !subst !substr
231-
: !tail !tolower !toupper !xor
226+
: !interleave !isa !le !listconcat !listflatten
227+
: !listremove !listsplat !logtwo !lt !mul
228+
: !ne !not !or !range !repr
229+
: !setdagarg !setdagname !setdagop !shl !size
230+
: !sra !srl !strconcat !sub !subst
231+
: !substr !tail !tolower !toupper !xor
232232

233233
The ``!cond`` operator has a slightly different
234234
syntax compared to other bang operators, so it is defined separately:
@@ -1832,6 +1832,12 @@ and non-0 as true.
18321832
This operator concatenates the list arguments *list1*, *list2*, etc., and
18331833
produces the resulting list. The lists must have the same element type.
18341834

1835+
``!listflatten(``\ *list*\ ``)``
1836+
This operator flattens a list of lists *list* and produces a list with all
1837+
elements of the constituent lists concatenated. If *list* is of type
1838+
``list<list<X>>`` the resulting list is of type ``list<X>``. If *list*'s
1839+
element type is not a list, the result is *list* itself.
1840+
18351841
``!listremove(``\ *list1*\ ``,`` *list2*\ ``)``
18361842
This operator returns a copy of *list1* removing all elements that also occur in
18371843
*list2*. The lists must have the same element type.

‎llvm/include/llvm/TableGen/Record.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -847,7 +847,8 @@ class UnOpInit : public OpInit, public FoldingSetNode {
847847
EMPTY,
848848
GETDAGOP,
849849
LOG2,
850-
REPR
850+
REPR,
851+
LISTFLATTEN,
851852
};
852853

853854
private:

‎llvm/lib/TableGen/Record.cpp‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -986,6 +986,32 @@ Init *UnOpInit::Fold(Record *CurRec, bool IsFinal) const {
986986
}
987987
}
988988
break;
989+
990+
case LISTFLATTEN:
991+
if (ListInit *LHSList = dyn_cast<ListInit>(LHS)) {
992+
ListRecTy *InnerListTy = dyn_cast<ListRecTy>(LHSList->getElementType());
993+
// list of non-lists, !listflatten() is a NOP.
994+
if (!InnerListTy)
995+
return LHS;
996+
997+
auto Flatten = [](ListInit *List) -> std::optional<std::vector<Init *>> {
998+
std::vector<Init *> Flattened;
999+
// Concatenate elements of all the inner lists.
1000+
for (Init *InnerInit : List->getValues()) {
1001+
ListInit *InnerList = dyn_cast<ListInit>(InnerInit);
1002+
if (!InnerList)
1003+
return std::nullopt;
1004+
for (Init *InnerElem : InnerList->getValues())
1005+
Flattened.push_back(InnerElem);
1006+
};
1007+
return Flattened;
1008+
};
1009+
1010+
auto Flattened = Flatten(LHSList);
1011+
if (Flattened)
1012+
return ListInit::get(*Flattened, InnerListTy->getElementType());
1013+
}
1014+
break;
9891015
}
9901016
return const_cast<UnOpInit *>(this);
9911017
}
@@ -1010,6 +1036,9 @@ std::string UnOpInit::getAsString() const {
10101036
case EMPTY: Result = "!empty"; break;
10111037
case GETDAGOP: Result = "!getdagop"; break;
10121038
case LOG2 : Result = "!logtwo"; break;
1039+
case LISTFLATTEN:
1040+
Result = "!listflatten";
1041+
break;
10131042
case REPR:
10141043
Result = "!repr";
10151044
break;

‎llvm/lib/TableGen/TGLexer.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -628,6 +628,7 @@ tgtok::TokKind TGLexer::LexExclaim() {
628628
.Case("foreach", tgtok::XForEach)
629629
.Case("filter", tgtok::XFilter)
630630
.Case("listconcat", tgtok::XListConcat)
631+
.Case("listflatten", tgtok::XListFlatten)
631632
.Case("listsplat", tgtok::XListSplat)
632633
.Case("listremove", tgtok::XListRemove)
633634
.Case("range", tgtok::XRange)

‎llvm/lib/TableGen/TGLexer.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ enum TokKind {
122122
XSRL,
123123
XSHL,
124124
XListConcat,
125+
XListFlatten,
125126
XListSplat,
126127
XStrConcat,
127128
XInterleave,

‎llvm/lib/TableGen/TGParser.cpp‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1190,6 +1190,7 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
11901190
case tgtok::XNOT:
11911191
case tgtok::XToLower:
11921192
case tgtok::XToUpper:
1193+
case tgtok::XListFlatten:
11931194
case tgtok::XLOG2:
11941195
case tgtok::XHead:
11951196
case tgtok::XTail:
@@ -1235,6 +1236,11 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
12351236
Code = UnOpInit::NOT;
12361237
Type = IntRecTy::get(Records);
12371238
break;
1239+
case tgtok::XListFlatten:
1240+
Lex.Lex(); // eat the operation.
1241+
Code = UnOpInit::LISTFLATTEN;
1242+
Type = IntRecTy::get(Records); // Bogus type used here.
1243+
break;
12381244
case tgtok::XLOG2:
12391245
Lex.Lex(); // eat the operation
12401246
Code = UnOpInit::LOG2;
@@ -1309,7 +1315,8 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
13091315
}
13101316
}
13111317

1312-
if (Code == UnOpInit::HEAD || Code == UnOpInit::TAIL) {
1318+
if (Code == UnOpInit::HEAD || Code == UnOpInit::TAIL ||
1319+
Code == UnOpInit::LISTFLATTEN) {
13131320
ListInit *LHSl = dyn_cast<ListInit>(LHS);
13141321
TypedInit *LHSt = dyn_cast<TypedInit>(LHS);
13151322
if (!LHSl && !LHSt) {
@@ -1328,19 +1335,34 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
13281335
TokError("empty list argument in unary operator");
13291336
return nullptr;
13301337
}
1338+
bool UseElementType =
1339+
Code == UnOpInit::HEAD || Code == UnOpInit::LISTFLATTEN;
13311340
if (LHSl) {
13321341
Init *Item = LHSl->getElement(0);
13331342
TypedInit *Itemt = dyn_cast<TypedInit>(Item);
13341343
if (!Itemt) {
13351344
TokError("untyped list element in unary operator");
13361345
return nullptr;
13371346
}
1338-
Type = (Code == UnOpInit::HEAD) ? Itemt->getType()
1339-
: ListRecTy::get(Itemt->getType());
1347+
Type = UseElementType ? Itemt->getType()
1348+
: ListRecTy::get(Itemt->getType());
13401349
} else {
13411350
assert(LHSt && "expected list type argument in unary operator");
13421351
ListRecTy *LType = dyn_cast<ListRecTy>(LHSt->getType());
1343-
Type = (Code == UnOpInit::HEAD) ? LType->getElementType() : LType;
1352+
Type = UseElementType ? LType->getElementType() : LType;
1353+
}
1354+
1355+
// for !listflatten, we expect a list of lists, but also support a list of
1356+
// non-lists, where !listflatten will be a NOP.
1357+
if (Code == UnOpInit::LISTFLATTEN) {
1358+
ListRecTy *InnerListTy = dyn_cast<ListRecTy>(Type);
1359+
if (InnerListTy) {
1360+
// listflatten will convert list<list<X>> to list<X>.
1361+
Type = ListRecTy::get(InnerListTy->getElementType());
1362+
} else {
1363+
// If its a list of non-lists, !listflatten will be a NOP.
1364+
Type = ListRecTy::get(Type);
1365+
}
13441366
}
13451367
}
13461368

@@ -1378,7 +1400,7 @@ Init *TGParser::ParseOperation(Record *CurRec, RecTy *ItemType) {
13781400

13791401
case tgtok::XExists: {
13801402
// Value ::= !exists '<' Type '>' '(' Value ')'
1381-
Lex.Lex(); // eat the operation
1403+
Lex.Lex(); // eat the operation.
13821404

13831405
RecTy *Type = ParseOperatorType();
13841406
if (!Type)

‎llvm/test/TableGen/listflatten-error.td‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
// RUN: not llvm-tblgen %s 2>&1 | FileCheck %s -DFILE=%s
2+
3+
// CHECK: [[FILE]]:[[@LINE+2]]:33: error: expected list type argument in unary operator
4+
class Flatten<int A> {
5+
list<int> F = !listflatten(A);
6+
}

‎llvm/test/TableGen/listflatten.td‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// RUN: llvm-tblgen %s | FileCheck %s
2+
3+
class Flatten<list<int> A, list<int> B> {
4+
list<int> Flat1 = !listflatten([A, B, [6], [7, 8]]);
5+
6+
list<list<int>> X = [A, B];
7+
list<int> Flat2 = !listflatten(!listconcat(X, [[7]]));
8+
9+
// Generate a nested list of integers.
10+
list<int> Y0 = [1, 2, 3, 4];
11+
list<list<int>> Y1 = !foreach(elem, Y0, [elem]);
12+
list<list<list<int>>> Y2 = !foreach(elem, Y1, [elem]);
13+
list<list<list<list<int>>>> Y3 = !foreach(elem, Y2, [elem]);
14+
15+
// Flatten it completely.
16+
list<int> Flat3=!listflatten(!listflatten(!listflatten(Y3)));
17+
18+
// Flatten it partially.
19+
list<list<list<int>>> Flat4 = !listflatten(Y3);
20+
list<list<int>> Flat5 = !listflatten(!listflatten(Y3));
21+
22+
// Test NOP flattening.
23+
list<string> Flat6 = !listflatten(["a", "b"]);
24+
}
25+
26+
// CHECK: list<int> Flat1 = [1, 2, 3, 4, 5, 6, 7, 8];
27+
// CHECK: list<int> Flat2 = [1, 2, 3, 4, 5, 7];
28+
// CHECK: list<int> Flat3 = [1, 2, 3, 4];
29+
// CHECK{LITERAL}: list<list<list<int>>> Flat4 = [[[1]], [[2]], [[3]], [[4]]];
30+
// CHECK: list<string> Flat6 = ["a", "b"];
31+
def F : Flatten<[1,2], [3,4,5]>;
32+

0 commit comments

Comments
 (0)