Skip to content

Commit 42ede18

Browse files
hborlahborla
authored
Merge pull request #62713 from hborla/open-pack-element-type-repr
[ConstraintSystem] Enable pack element type reprs in expressions.
2 parents 28ee477 + 8956de4 commit 42ede18

20 files changed

+304
-94
lines changed

‎include/swift/AST/Expr.h

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3584,10 +3584,6 @@ class PackExpansionExpr final : public Expr,
35843584
: Expr(ExprKind::PackExpansion, implicit, type),
35853585
PatternExpr(patternExpr), DotsLoc(dotsLoc), Environment(environment) {
35863586
Bits.PackExpansionExpr.NumBindings = packElements.size();
3587-
3588-
assert(Bits.PackExpansionExpr.NumBindings > 0 &&
3589-
"PackExpansionExpr must have pack references");
3590-
35913587
std::uninitialized_copy(packElements.begin(), packElements.end(),
35923588
getTrailingObjects<PackElementExpr *>());
35933589
}
@@ -3623,9 +3619,7 @@ class PackExpansionExpr final : public Expr,
36233619
return {getTrailingObjects<PackElementExpr *>(), getNumBindings()};
36243620
}
36253621

3626-
void setBinding(unsigned i, PackElementExpr *e) {
3627-
getMutableBindings()[i] = e;
3628-
}
3622+
void getExpandedPacks(SmallVectorImpl<ASTNode> &packs);
36293623

36303624
GenericEnvironment *getGenericEnvironment() {
36313625
return Environment;

‎include/swift/AST/Types.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ class InOutType;
7272
class OpaqueTypeDecl;
7373
class OpenedArchetypeType;
7474
class PackType;
75+
class PackReferenceTypeRepr;
7576
class PlaceholderTypeRepr;
7677
enum class ReferenceCounting : uint8_t;
7778
enum class ResilienceExpansion : unsigned;

‎include/swift/Sema/ConstraintSystem.h

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -680,6 +680,13 @@ template <typename T = Decl> T *getAsDecl(ASTNode node) {
680680
return nullptr;
681681
}
682682

683+
template <typename T = TypeRepr>
684+
T *getAsTypeRepr(ASTNode node) {
685+
if (auto *type = node.dyn_cast<TypeRepr *>())
686+
return dyn_cast_or_null<T>(type);
687+
return nullptr;
688+
}
689+
683690
template <typename T = Stmt>
684691
T *getAsStmt(ASTNode node) {
685692
if (auto *S = node.dyn_cast<Stmt *>())
@@ -6108,6 +6115,46 @@ class HandlePlaceholderType {
61086115
}
61096116
};
61106117

6118+
/// A function object that opens a given pack type by generating a
6119+
/// \c PackElementOf constraint.
6120+
class OpenPackElementType {
6121+
ConstraintSystem &cs;
6122+
ConstraintLocator *locator;
6123+
GenericEnvironment *elementEnv;
6124+
6125+
public:
6126+
explicit OpenPackElementType(ConstraintSystem &cs,
6127+
const ConstraintLocatorBuilder &locator,
6128+
GenericEnvironment *elementEnv)
6129+
: cs(cs), elementEnv(elementEnv) {
6130+
this->locator = cs.getConstraintLocator(locator);
6131+
}
6132+
6133+
Type operator()(Type packType, PackReferenceTypeRepr *packRepr) const {
6134+
// Only assert we have an element environment when invoking the function
6135+
// object. In cases where pack elements are referenced outside of a
6136+
// pack expansion, type resolution will error before opening the pack
6137+
// element.
6138+
assert(elementEnv);
6139+
6140+
auto *elementType = cs.createTypeVariable(locator, TVO_CanBindToHole);
6141+
auto elementLoc = cs.getConstraintLocator(locator,
6142+
LocatorPathElt::OpenedPackElement(elementEnv));
6143+
6144+
// If we're opening a pack element from an explicit type repr,
6145+
// set the type repr types in the constraint system for generating
6146+
// ShapeOf constraints when visiting the PackExpansionExpr.
6147+
if (packRepr) {
6148+
cs.setType(packRepr->getPackType(), packType);
6149+
cs.setType(packRepr, elementType);
6150+
}
6151+
6152+
cs.addConstraint(ConstraintKind::PackElementOf, elementType,
6153+
packType, elementLoc);
6154+
return elementType;
6155+
}
6156+
};
6157+
61116158
/// Compute the shuffle required to map from a given tuple type to
61126159
/// another tuple type.
61136160
///

‎lib/AST/Expr.cpp

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1259,6 +1259,41 @@ PackExpansionExpr::create(ASTContext &ctx, Expr *patternExpr,
12591259
implicit, type);
12601260
}
12611261

1262+
void PackExpansionExpr::getExpandedPacks(SmallVectorImpl<ASTNode> &packs) {
1263+
struct PackCollector : public ASTWalker {
1264+
llvm::SmallVector<ASTNode, 2> packs;
1265+
1266+
virtual PreWalkResult<Expr *> walkToExprPre(Expr *E) override {
1267+
// Don't walk into nested pack expansions
1268+
if (isa<PackExpansionExpr>(E)) {
1269+
return Action::SkipChildren(E);
1270+
}
1271+
1272+
if (isa<PackElementExpr>(E)) {
1273+
packs.push_back(E);
1274+
}
1275+
1276+
return Action::Continue(E);
1277+
}
1278+
1279+
virtual PreWalkAction walkToTypeReprPre(TypeRepr *T) override {
1280+
// Don't walk into nested pack expansions
1281+
if (isa<PackExpansionTypeRepr>(T)) {
1282+
return Action::SkipChildren();
1283+
}
1284+
1285+
if (isa<PackReferenceTypeRepr>(T)) {
1286+
packs.push_back(T);
1287+
}
1288+
1289+
return Action::Continue();
1290+
}
1291+
} packCollector;
1292+
1293+
getPatternExpr()->walk(packCollector);
1294+
packs.append(packCollector.packs.begin(), packCollector.packs.end());
1295+
}
1296+
12621297
PackElementExpr *
12631298
PackElementExpr::create(ASTContext &ctx, SourceLoc eachLoc, Expr *packRefExpr,
12641299
bool implicit, Type type) {

‎lib/Sema/CSBindings.cpp

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1428,9 +1428,6 @@ void PotentialBindings::infer(Constraint *constraint) {
14281428
// Produce a potential binding to the opened element archetype corresponding
14291429
// to the pack type.
14301430
packType = packType->mapTypeOutOfContext();
1431-
if (auto *expansion = packType->getAs<PackExpansionType>())
1432-
packType = expansion->getPatternType();
1433-
14341431
auto *elementEnv = openedElement->getGenericEnvironment();
14351432
auto elementType = elementEnv->mapPackTypeIntoElementContext(packType);
14361433
addPotentialBinding({elementType, AllowedBindingKind::Exact, constraint});

‎lib/Sema/CSGen.cpp

Lines changed: 52 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include "swift/AST/ASTWalker.h"
2424
#include "swift/AST/Expr.h"
2525
#include "swift/AST/GenericSignature.h"
26+
#include "swift/AST/GenericEnvironment.h"
2627
#include "swift/AST/ParameterList.h"
2728
#include "swift/AST/PrettyStackTrace.h"
2829
#include "swift/AST/SubstitutionMap.h"
@@ -1416,11 +1417,23 @@ namespace {
14161417
Type
14171418
resolveTypeReferenceInExpression(TypeRepr *repr, TypeResolverContext resCtx,
14181419
const ConstraintLocatorBuilder &locator) {
1420+
TypeResolutionOptions options(resCtx);
1421+
14191422
// Introduce type variables for unbound generics.
14201423
const auto genericOpener = OpenUnboundGenericType(CS, locator);
14211424
const auto placeholderHandler = HandlePlaceholderType(CS, locator);
1425+
1426+
// Add a PackElementOf constraint for 'each T' type reprs.
1427+
GenericEnvironment *elementEnv = nullptr;
1428+
if (!PackElementEnvironments.empty()) {
1429+
options |= TypeResolutionFlags::AllowPackReferences;
1430+
elementEnv = PackElementEnvironments.back();
1431+
}
1432+
const auto packElementOpener = OpenPackElementType(CS, locator, elementEnv);
1433+
14221434
const auto result = TypeResolution::resolveContextualType(
1423-
repr, CS.DC, resCtx, genericOpener, placeholderHandler);
1435+
repr, CS.DC, options, genericOpener, placeholderHandler,
1436+
packElementOpener);
14241437
if (result->hasError()) {
14251438
CS.recordFix(
14261439
IgnoreInvalidASTNode::create(CS, CS.getConstraintLocator(locator)));
@@ -1659,14 +1672,20 @@ namespace {
16591672
ArrayRef<TypeRepr *> specializationArgs) {
16601673
// Resolve each type.
16611674
SmallVector<Type, 2> specializationArgTypes;
1662-
const auto options =
1675+
auto options =
16631676
TypeResolutionOptions(TypeResolverContext::InExpression);
16641677
for (auto specializationArg : specializationArgs) {
1678+
GenericEnvironment *elementEnv = nullptr;
1679+
if (!PackElementEnvironments.empty()) {
1680+
options |= TypeResolutionFlags::AllowPackReferences;
1681+
elementEnv = PackElementEnvironments.back();
1682+
}
16651683
const auto result = TypeResolution::resolveContextualType(
16661684
specializationArg, CurDC, options,
16671685
// Introduce type variables for unbound generics.
16681686
OpenUnboundGenericType(CS, locator),
1669-
HandlePlaceholderType(CS, locator));
1687+
HandlePlaceholderType(CS, locator),
1688+
OpenPackElementType(CS, locator, elementEnv));
16701689
if (result->hasError())
16711690
return true;
16721691

@@ -1723,14 +1742,21 @@ namespace {
17231742
// Bind the specified generic arguments to the type variables in the
17241743
// open type.
17251744
auto *const locator = CS.getConstraintLocator(expr);
1726-
const auto options =
1745+
auto options =
17271746
TypeResolutionOptions(TypeResolverContext::InExpression);
17281747
for (size_t i = 0, e = specializations.size(); i < e; ++i) {
1748+
GenericEnvironment *elementEnv = nullptr;
1749+
if (!PackElementEnvironments.empty()) {
1750+
options |= TypeResolutionFlags::AllowPackReferences;
1751+
elementEnv = PackElementEnvironments.back();
1752+
}
1753+
17291754
const auto result = TypeResolution::resolveContextualType(
17301755
specializations[i], CS.DC, options,
17311756
// Introduce type variables for unbound generics.
17321757
OpenUnboundGenericType(CS, locator),
1733-
HandlePlaceholderType(CS, locator));
1758+
HandlePlaceholderType(CS, locator),
1759+
OpenPackElementType(CS, locator, elementEnv));
17341760
if (result->hasError())
17351761
return Type();
17361762

@@ -2983,29 +3009,36 @@ namespace {
29833009
auto *shapeTypeVar = CS.createTypeVariable(shapeLoc,
29843010
TVO_CanBindToPack |
29853011
TVO_CanBindToHole);
2986-
auto packReference = expr->getPackElements().front()->getPackRefExpr();
2987-
auto packType = CS.simplifyType(CS.getType(packReference))
2988-
->castTo<PackExpansionType>()->getPatternType();
2989-
CS.addConstraint(ConstraintKind::ShapeOf, packType, shapeTypeVar,
2990-
CS.getConstraintLocator(expr));
3012+
3013+
// Generate ShapeOf constraints between all packs expanded by this
3014+
// pack expansion expression through the shape type variable.
3015+
SmallVector<ASTNode, 2> expandedPacks;
3016+
expr->getExpandedPacks(expandedPacks);
3017+
for (auto pack : expandedPacks) {
3018+
Type packType;
3019+
if (auto *elementExpr = getAsExpr<PackElementExpr>(pack)) {
3020+
packType = CS.getType(elementExpr->getPackRefExpr());
3021+
} else if (auto *elementType = getAsTypeRepr<PackReferenceTypeRepr>(pack)) {
3022+
packType = CS.getType(elementType->getPackType());
3023+
} else {
3024+
llvm_unreachable("unsupported pack reference ASTNode");
3025+
}
3026+
3027+
CS.addConstraint(ConstraintKind::ShapeOf, packType, shapeTypeVar,
3028+
CS.getConstraintLocator(expr));
3029+
}
29913030

29923031
return PackExpansionType::get(patternTy, shapeTypeVar);
29933032
}
29943033

29953034
Type visitPackElementExpr(PackElementExpr *expr) {
29963035
auto packType = CS.getType(expr->getPackRefExpr());
2997-
auto *elementType =
2998-
CS.createTypeVariable(CS.getConstraintLocator(expr),
2999-
TVO_CanBindToHole);
3000-
auto *elementEnv = PackElementEnvironments.back();
3001-
auto *elementLoc = CS.getConstraintLocator(
3002-
expr, LocatorPathElt::OpenedPackElement(elementEnv));
30033036

30043037
// The type of a PackElementExpr is the opened pack element archetype
30053038
// of the pack reference.
3006-
CS.addConstraint(ConstraintKind::PackElementOf, elementType, packType,
3007-
elementLoc);
3008-
return elementType;
3039+
OpenPackElementType openPackElement(CS, CS.getConstraintLocator(expr),
3040+
PackElementEnvironments.back());
3041+
return openPackElement(packType, /*packRepr*/ nullptr);
30093042
}
30103043

30113044
Type visitDynamicTypeExpr(DynamicTypeExpr *expr) {

‎lib/Sema/CSSimplify.cpp

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8612,8 +8612,6 @@ ConstraintSystem::simplifyPackElementOfConstraint(Type first, Type second,
86128612

86138613
// This constraint only exists to vend bindings.
86148614
auto *packEnv = DC->getGenericEnvironmentOfContext();
8615-
if (auto *expansion = packType->getAs<PackExpansionType>())
8616-
packType = expansion->getPatternType();
86178615
if (packType->isEqual(packEnv->mapElementTypeIntoPackContext
86188616
(elementType->mapTypeOutOfContext()))) {
86198617
return SolutionKind::Solved;

‎lib/Sema/CSSolver.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1190,7 +1190,8 @@ void ConstraintSystem::shrink(Expr *expr) {
11901190
// example:
11911191
// let foo: [Array<Float>] = [[0], [1], [2]] as [Array]
11921192
// let foo: [Array<Float>] = [[0], [1], [2]] as [Array<_>]
1193-
/*unboundTyOpener*/ nullptr, /*placeholderHandler*/ nullptr);
1193+
/*unboundTyOpener*/ nullptr, /*placeholderHandler*/ nullptr,
1194+
/*packElementOpener*/ nullptr);
11941195

11951196
// Looks like coercion type is invalid, let's skip this sub-tree.
11961197
if (coercionType->hasError())

‎lib/Sema/ConstraintSystem.cpp

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -780,7 +780,8 @@ Type ConstraintSystem::openUnboundGenericType(GenericTypeDecl *decl,
780780
TypeResolution::forInterface(
781781
DC, None,
782782
[](auto) -> Type { llvm_unreachable("should not be used"); },
783-
[](auto &, auto) -> Type { llvm_unreachable("should not be used"); })
783+
[](auto &, auto) -> Type { llvm_unreachable("should not be used"); },
784+
[](auto, auto) -> Type { llvm_unreachable("should not be used"); })
784785
.applyUnboundGenericArguments(decl, parentTy, SourceLoc(), arguments);
785786
if (!parentTy && !isTypeResolution) {
786787
result = DC->mapTypeIntoContext(result);
@@ -1281,6 +1282,10 @@ Type ConstraintSystem::getUnopenedTypeOfReference(
12811282
Type requestedType =
12821283
getType(value)->getWithoutSpecifierType()->getReferenceStorageReferent();
12831284

1285+
// Strip pack expansion types off of pack references.
1286+
if (auto *expansion = requestedType->getAs<PackExpansionType>())
1287+
requestedType = expansion->getPatternType();
1288+
12841289
// Adjust the type for concurrency if requested.
12851290
if (adjustForPreconcurrency)
12861291
requestedType = adjustVarTypeForConcurrency(
@@ -1596,7 +1601,8 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
15961601
auto type =
15971602
TypeResolution::forInterface(useDC, TypeResolverContext::InExpression,
15981603
/*unboundTyOpener*/ nullptr,
1599-
/*placeholderHandler*/ nullptr)
1604+
/*placeholderHandler*/ nullptr,
1605+
/*packElementOpener*/ nullptr)
16001606
.resolveTypeInContext(typeDecl, /*foundDC*/ nullptr,
16011607
/*isSpecialized=*/false);
16021608
type = useDC->mapTypeIntoContext(type);
@@ -2957,7 +2963,8 @@ FunctionType::ExtInfo ClosureEffectsRequest::evaluate(
29572963
castType = TypeResolution::resolveContextualType(
29582964
castTypeRepr, DC, TypeResolverContext::InExpression,
29592965
/*unboundTyOpener*/ nullptr,
2960-
/*placeholderHandler*/ nullptr);
2966+
/*placeholderHandler*/ nullptr,
2967+
/*packElementOpener*/ nullptr);
29612968
} else {
29622969
castType = isp->getCastType();
29632970
}

0 commit comments

Comments
 (0)