Skip to content

Commit d3bc4f5

Browse files
hborlahborla
authored
Merge pull request #62582 from hborla/pack-element-expr
2 parents cdf59ad + 472bd1f commit d3bc4f5

21 files changed

+296
-141
lines changed

‎include/swift/AST/Expr.h‎

Lines changed: 56 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -3515,6 +3515,50 @@ class VarargExpansionExpr : public Expr {
35153515
}
35163516
};
35173517

3518+
/// A pack element expression spelled with the contextual \c each
3519+
/// keyword applied to a pack reference expression.
3520+
///
3521+
/// \code
3522+
/// func zip<T..., U...>(t: (each T)..., u: (each U)...) {
3523+
/// let zipped = (each t, each u)...
3524+
/// }
3525+
/// \endcode
3526+
///
3527+
/// Pack elements can only appear in the pattern expression of a
3528+
/// \c PackExpansionExpr.
3529+
class PackElementExpr final : public Expr {
3530+
SourceLoc EachLoc;
3531+
Expr *PackRefExpr;
3532+
3533+
PackElementExpr(SourceLoc eachLoc, Expr *packRefExpr,
3534+
bool implicit = false, Type type = Type())
3535+
: Expr(ExprKind::PackElement, implicit, type),
3536+
EachLoc(eachLoc), PackRefExpr(packRefExpr) {}
3537+
3538+
public:
3539+
static PackElementExpr *create(ASTContext &ctx, SourceLoc eachLoc,
3540+
Expr *packRefExpr, bool implicit = false,
3541+
Type type = Type());
3542+
3543+
Expr *getPackRefExpr() const { return PackRefExpr; }
3544+
3545+
void setPackRefExpr(Expr *packRefExpr) {
3546+
PackRefExpr = packRefExpr;
3547+
}
3548+
3549+
SourceLoc getStartLoc() const {
3550+
return EachLoc;
3551+
}
3552+
3553+
SourceLoc getEndLoc() const {
3554+
return PackRefExpr->getEndLoc();
3555+
}
3556+
3557+
static bool classof(const Expr *E) {
3558+
return E->getKind() == ExprKind::PackElement;
3559+
}
3560+
};
3561+
35183562
/// A pack expansion expression is a pattern expression followed by
35193563
/// the expansion operator '...'. The pattern expression contains
35203564
/// references to parameter packs of length N, and the expansion
@@ -3525,51 +3569,41 @@ class VarargExpansionExpr : public Expr {
35253569
/// call argument lists, the elements of a tuple value, and the source
35263570
/// of a for-in loop.
35273571
class PackExpansionExpr final : public Expr,
3528-
private llvm::TrailingObjects<PackExpansionExpr,
3529-
OpaqueValueExpr *, Expr *> {
3572+
private llvm::TrailingObjects<PackExpansionExpr, PackElementExpr *> {
35303573
friend TrailingObjects;
35313574

35323575
Expr *PatternExpr;
35333576
SourceLoc DotsLoc;
35343577
GenericEnvironment *Environment;
35353578

35363579
PackExpansionExpr(Expr *patternExpr,
3537-
ArrayRef<OpaqueValueExpr *> opaqueValues,
3538-
ArrayRef<Expr *> bindings,
3580+
ArrayRef<PackElementExpr *> packElements,
35393581
SourceLoc dotsLoc,
35403582
GenericEnvironment *environment,
35413583
bool implicit, Type type)
35423584
: Expr(ExprKind::PackExpansion, implicit, type),
35433585
PatternExpr(patternExpr), DotsLoc(dotsLoc), Environment(environment) {
3544-
assert(opaqueValues.size() == bindings.size());
3545-
Bits.PackExpansionExpr.NumBindings = opaqueValues.size();
3586+
Bits.PackExpansionExpr.NumBindings = packElements.size();
35463587

35473588
assert(Bits.PackExpansionExpr.NumBindings > 0 &&
35483589
"PackExpansionExpr must have pack references");
35493590

3550-
std::uninitialized_copy(opaqueValues.begin(), opaqueValues.end(),
3551-
getTrailingObjects<OpaqueValueExpr *>());
3552-
std::uninitialized_copy(bindings.begin(), bindings.end(),
3553-
getTrailingObjects<Expr *>());
3554-
}
3555-
3556-
size_t numTrailingObjects(OverloadToken<OpaqueValueExpr *>) const {
3557-
return getNumBindings();
3591+
std::uninitialized_copy(packElements.begin(), packElements.end(),
3592+
getTrailingObjects<PackElementExpr *>());
35583593
}
35593594

3560-
size_t numTrailingObjects(OverloadToken<Expr *>) const {
3595+
size_t numTrailingObjects(OverloadToken<PackElementExpr *>) const {
35613596
return getNumBindings();
35623597
}
35633598

3564-
MutableArrayRef<Expr *> getMutableBindings() {
3565-
return {getTrailingObjects<Expr *>(), getNumBindings()};
3599+
MutableArrayRef<PackElementExpr *> getMutableBindings() {
3600+
return {getTrailingObjects<PackElementExpr *>(), getNumBindings()};
35663601
}
35673602

35683603
public:
35693604
static PackExpansionExpr *create(ASTContext &ctx,
35703605
Expr *patternExpr,
3571-
ArrayRef<OpaqueValueExpr *> opaqueValues,
3572-
ArrayRef<Expr *> bindings,
3606+
ArrayRef<PackElementExpr *> packElements,
35733607
SourceLoc dotsLoc,
35743608
GenericEnvironment *environment,
35753609
bool implicit = false,
@@ -3585,15 +3619,11 @@ class PackExpansionExpr final : public Expr,
35853619
return Bits.PackExpansionExpr.NumBindings;
35863620
}
35873621

3588-
ArrayRef<OpaqueValueExpr *> getOpaqueValues() {
3589-
return {getTrailingObjects<OpaqueValueExpr *>(), getNumBindings()};
3590-
}
3591-
3592-
ArrayRef<Expr *> getBindings() {
3593-
return {getTrailingObjects<Expr *>(), getNumBindings()};
3622+
ArrayRef<PackElementExpr *> getPackElements() {
3623+
return {getTrailingObjects<PackElementExpr *>(), getNumBindings()};
35943624
}
35953625

3596-
void setBinding(unsigned i, Expr *e) {
3626+
void setBinding(unsigned i, PackElementExpr *e) {
35973627
getMutableBindings()[i] = e;
35983628
}
35993629

‎include/swift/AST/ExprNodes.def‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,6 +130,7 @@ ABSTRACT_EXPR(AbstractClosure, Expr)
130130
EXPR(InOut, Expr)
131131
EXPR(VarargExpansion, Expr)
132132
EXPR(PackExpansion, Expr)
133+
EXPR(PackElement, Expr)
133134
EXPR(DynamicType, Expr)
134135
EXPR(RebindSelfInConstructor, Expr)
135136
EXPR(OpaqueValue, Expr)

‎include/swift/Sema/ConstraintLocator.h‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -687,6 +687,21 @@ class LocatorPathElt::PackElement final : public StoredIntegerElement<1> {
687687
}
688688
};
689689

690+
class LocatorPathElt::OpenedPackElement final
691+
: public StoredPointerElement<GenericEnvironment> {
692+
public:
693+
OpenedPackElement(GenericEnvironment *env)
694+
: StoredPointerElement(PathElementKind::OpenedPackElement, env) {}
695+
696+
GenericEnvironment *getGenericEnvironment() const {
697+
return getStoredPointer();
698+
}
699+
700+
static bool classof(const LocatorPathElt *elt) {
701+
return elt->getKind() == PathElementKind::OpenedPackElement;
702+
}
703+
};
704+
690705
class LocatorPathElt::KeyPathComponent final : public StoredIntegerElement<1> {
691706
public:
692707
KeyPathComponent(unsigned index)

‎include/swift/Sema/ConstraintLocatorPathElts.def‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,9 @@ CUSTOM_LOCATOR_PATH_ELT(PackType)
195195
/// An element of a pack type - the T in <T, U, V, ...>
196196
CUSTOM_LOCATOR_PATH_ELT(PackElement)
197197

198+
/// Stores the generic environment for an opened pack element.
199+
CUSTOM_LOCATOR_PATH_ELT(OpenedPackElement)
200+
198201
/// The shape of a parameter pack.
199202
SIMPLE_LOCATOR_PATH_ELT(PackShape)
200203

‎lib/AST/ASTDumper.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2504,6 +2504,12 @@ class PrintExpr : public ExprVisitor<PrintExpr> {
25042504
PrintWithColorRAII(OS, ParenthesisColor) << ')';
25052505
}
25062506

2507+
void visitPackElementExpr(PackElementExpr *E) {
2508+
printCommon(E, "pack_element_expr") << "\n";
2509+
printRec(E->getPackRefExpr());
2510+
PrintWithColorRAII(OS, ParenthesisColor) << ')';
2511+
}
2512+
25072513
void visitForceTryExpr(ForceTryExpr *E) {
25082514
printCommon(E, "force_try_expr");
25092515
OS << '\n';

‎lib/AST/ASTPrinter.cpp‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4881,6 +4881,10 @@ void PrintAST::visitPackExpansionExpr(PackExpansionExpr *expr) {
48814881
visit(expr->getPatternExpr());
48824882
}
48834883

4884+
void PrintAST::visitPackElementExpr(PackElementExpr *expr) {
4885+
visit(expr->getPackRefExpr());
4886+
}
4887+
48844888
void PrintAST::visitArchetypeToSuperExpr(ArchetypeToSuperExpr *expr) {
48854889
}
48864890

‎lib/AST/ASTVerifier.cpp‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -836,12 +836,6 @@ class Verifier : public ASTWalker {
836836
return false;
837837

838838
Generics.push_back(expr->getGenericEnvironment());
839-
840-
for (auto *placeholder : expr->getOpaqueValues()) {
841-
assert(!OpaqueValues.count(placeholder));
842-
OpaqueValues[placeholder] = 0;
843-
}
844-
845839
return true;
846840
}
847841

@@ -854,13 +848,6 @@ class Verifier : public ASTWalker {
854848
verifyCheckedAlwaysBase(E);
855849
}
856850

857-
void cleanup(PackExpansionExpr *expr) {
858-
for (auto *placeholder : expr->getOpaqueValues()) {
859-
assert(OpaqueValues.count(placeholder));
860-
OpaqueValues.erase(placeholder);
861-
}
862-
}
863-
864851
bool shouldVerify(MakeTemporarilyEscapableExpr *expr) {
865852
if (!shouldVerify(cast<Expr>(expr)))
866853
return false;

‎lib/AST/ASTWalker.cpp‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -859,6 +859,14 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
859859
return nullptr;
860860
}
861861

862+
Expr *visitPackElementExpr(PackElementExpr *E) {
863+
if (Expr *pattern = doIt(E->getPackRefExpr())) {
864+
E->setPackRefExpr(pattern);
865+
return E;
866+
}
867+
return nullptr;
868+
}
869+
862870
Expr *visitSequenceExpr(SequenceExpr *E) {
863871
for (unsigned i = 0, e = E->getNumElements(); i != e; ++i)
864872
if (Expr *Elt = doIt(E->getElement(i)))

‎lib/AST/Expr.cpp‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -394,6 +394,7 @@ ConcreteDeclRef Expr::getReferencedDecl(bool stopAtParenExpr) const {
394394

395395
NO_REFERENCE(VarargExpansion);
396396
NO_REFERENCE(PackExpansion);
397+
NO_REFERENCE(PackElement);
397398
NO_REFERENCE(DynamicType);
398399

399400
PASS_THROUGH_REFERENCE(RebindSelfInConstructor, getSubExpr);
@@ -750,6 +751,7 @@ bool Expr::canAppendPostfixExpression(bool appendingPostfixOperator) const {
750751
case ExprKind::MakeTemporarilyEscapable:
751752
case ExprKind::VarargExpansion:
752753
case ExprKind::PackExpansion:
754+
case ExprKind::PackElement:
753755
return false;
754756

755757
case ExprKind::Call:
@@ -927,6 +929,7 @@ bool Expr::isValidParentOfTypeExpr(Expr *typeExpr) const {
927929
case ExprKind::InOut:
928930
case ExprKind::VarargExpansion:
929931
case ExprKind::PackExpansion:
932+
case ExprKind::PackElement:
930933
case ExprKind::DynamicType:
931934
case ExprKind::RebindSelfInConstructor:
932935
case ExprKind::OpaqueValue:
@@ -1245,19 +1248,23 @@ VarargExpansionExpr *VarargExpansionExpr::createArrayExpansion(ASTContext &ctx,
12451248

12461249
PackExpansionExpr *
12471250
PackExpansionExpr::create(ASTContext &ctx, Expr *patternExpr,
1248-
ArrayRef<OpaqueValueExpr *> opaqueValues,
1249-
ArrayRef<Expr *> bindings, SourceLoc dotsLoc,
1250-
GenericEnvironment *environment,
1251+
ArrayRef<PackElementExpr *> packElements,
1252+
SourceLoc dotsLoc, GenericEnvironment *environment,
12511253
bool implicit, Type type) {
12521254
size_t size =
1253-
totalSizeToAlloc<OpaqueValueExpr *, Expr *>(opaqueValues.size(),
1254-
bindings.size());
1255+
totalSizeToAlloc<PackElementExpr *>(packElements.size());
12551256
void *mem = ctx.Allocate(size, alignof(PackExpansionExpr));
1256-
return ::new (mem) PackExpansionExpr(patternExpr, opaqueValues,
1257-
bindings, dotsLoc, environment,
1257+
return ::new (mem) PackExpansionExpr(patternExpr, packElements,
1258+
dotsLoc, environment,
12581259
implicit, type);
12591260
}
12601261

1262+
PackElementExpr *
1263+
PackElementExpr::create(ASTContext &ctx, SourceLoc eachLoc, Expr *packRefExpr,
1264+
bool implicit, Type type) {
1265+
return new (ctx) PackElementExpr(eachLoc, packRefExpr, implicit, type);
1266+
}
1267+
12611268
SequenceExpr *SequenceExpr::create(ASTContext &ctx, ArrayRef<Expr*> elements) {
12621269
assert(elements.size() & 1 && "even number of elements in sequence");
12631270
size_t bytes = totalSizeToAlloc<Expr *>(elements.size());

‎lib/Parse/ParseExpr.cpp‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1612,6 +1612,19 @@ ParserResult<Expr> Parser::parseExprPrimary(Diag<> ID, bool isExprBasic) {
16121612
return makeParserResult(typeExpr);
16131613
}
16141614

1615+
// 'each' followed by another identifier is a pack element expr.
1616+
if (Tok.isContextualKeyword("each") &&
1617+
peekToken().is(tok::identifier) &&
1618+
!peekToken().isAtStartOfLine()) {
1619+
SourceLoc loc = consumeToken();
1620+
ParserResult<Expr> ref = parseExpr(ID);
1621+
if (ref.isNull())
1622+
return ref;
1623+
1624+
auto *packRef = PackElementExpr::create(Context, loc, ref.get());
1625+
return makeParserResult(packRef);
1626+
}
1627+
16151628
LLVM_FALLTHROUGH;
16161629
case tok::kw_Self: // Self
16171630
return parseExprIdentifier();

0 commit comments

Comments
 (0)