Skip to content

Commit 53d8ba8

Browse files
slavapestovslavapestov
committed
Sema: Record pack element environment captures
1 parent 9dcf2a2 commit 53d8ba8

File tree

5 files changed

+147
-24
lines changed

5 files changed

+147
-24
lines changed

‎include/swift/AST/CaptureInfo.h

Lines changed: 35 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ class ValueDecl;
4343
class FuncDecl;
4444
class OpaqueValueExpr;
4545
class VarDecl;
46+
class GenericEnvironment;
4647

4748
/// CapturedValue includes both the declaration being captured, along with flags
4849
/// that indicate how it is captured.
@@ -51,7 +52,8 @@ class CapturedValue {
5152

5253
public:
5354
using Storage =
54-
llvm::PointerIntPair<llvm::PointerUnion<ValueDecl*, OpaqueValueExpr*>, 2,
55+
llvm::PointerIntPair<llvm::PointerUnion<ValueDecl *,
56+
OpaqueValueExpr *>, 2,
5557
unsigned>;
5658

5759
private:
@@ -140,19 +142,27 @@ class DynamicSelfType;
140142
/// Stores information about captured variables.
141143
class CaptureInfo {
142144
class CaptureInfoStorage final
143-
: public llvm::TrailingObjects<CaptureInfoStorage, CapturedValue> {
145+
: public llvm::TrailingObjects<CaptureInfoStorage,
146+
CapturedValue,
147+
GenericEnvironment *> {
144148

145149
DynamicSelfType *DynamicSelf;
146150
OpaqueValueExpr *OpaqueValue;
147-
unsigned Count;
151+
unsigned NumCapturedValues;
152+
unsigned NumGenericEnvironments;
153+
148154
public:
149-
explicit CaptureInfoStorage(unsigned count, DynamicSelfType *dynamicSelf,
150-
OpaqueValueExpr *opaqueValue)
151-
: DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue), Count(count) { }
155+
explicit CaptureInfoStorage(DynamicSelfType *dynamicSelf,
156+
OpaqueValueExpr *opaqueValue,
157+
unsigned numCapturedValues,
158+
unsigned numGenericEnvironments)
159+
: DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue),
160+
NumCapturedValues(numCapturedValues),
161+
NumGenericEnvironments(numGenericEnvironments) { }
152162

153-
ArrayRef<CapturedValue> getCaptures() const {
154-
return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), Count);
155-
}
163+
ArrayRef<CapturedValue> getCaptures() const;
164+
165+
ArrayRef<GenericEnvironment *> getGenericEnvironments() const;
156166

157167
DynamicSelfType *getDynamicSelfType() const {
158168
return DynamicSelf;
@@ -161,6 +171,10 @@ class CaptureInfo {
161171
OpaqueValueExpr *getOpaqueValue() const {
162172
return OpaqueValue;
163173
}
174+
175+
unsigned numTrailingObjects(OverloadToken<CapturedValue>) const {
176+
return NumCapturedValues;
177+
}
164178
};
165179

166180
enum class Flags : unsigned {
@@ -173,7 +187,9 @@ class CaptureInfo {
173187
public:
174188
/// The default-constructed CaptureInfo is "not yet computed".
175189
CaptureInfo() = default;
176-
CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
190+
CaptureInfo(ASTContext &ctx,
191+
ArrayRef<CapturedValue> captures,
192+
ArrayRef<GenericEnvironment *> genericEnv,
177193
DynamicSelfType *dynamicSelf, OpaqueValueExpr *opaqueValue,
178194
bool genericParamCaptures);
179195

@@ -190,12 +206,20 @@ class CaptureInfo {
190206
!hasDynamicSelfCapture() && !hasOpaqueValueCapture();
191207
}
192208

209+
/// Returns all captured values and opaque expressions.
193210
ArrayRef<CapturedValue> getCaptures() const {
194211
assert(hasBeenComputed());
195212
return StorageAndFlags.getPointer()->getCaptures();
196213
}
197214

198-
/// \returns true if the function captures any generic type parameters.
215+
/// Returns all captured pack element environments.
216+
ArrayRef<GenericEnvironment *> getGenericEnvironments() const {
217+
assert(hasBeenComputed());
218+
return StorageAndFlags.getPointer()->getGenericEnvironments();
219+
}
220+
221+
/// \returns true if the function captures the primary generic environment
222+
/// from its innermost declaration context.
199223
bool hasGenericParamCaptures() const {
200224
assert(hasBeenComputed());
201225
return StorageAndFlags.getInt().contains(Flags::HasGenericParamCaptures);

‎include/swift/Sema/ConstraintSystem.h

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2229,6 +2229,10 @@ class ConstraintSystem {
22292229
/// from declared parameters/result and body.
22302230
llvm::MapVector<const ClosureExpr *, FunctionType *> ClosureTypes;
22312231

2232+
/// Maps closures and local functions to the pack expansion expressions they
2233+
/// capture.
2234+
llvm::MapVector<AnyFunctionRef, SmallVector<PackExpansionExpr *, 1>> CapturedExpansions;
2235+
22322236
/// Maps expressions for implied results (e.g implicit 'then' statements,
22332237
/// implicit 'return' statements in single expression body closures) to their
22342238
/// result kind.
@@ -3164,6 +3168,19 @@ class ConstraintSystem {
31643168
return nullptr;
31653169
}
31663170

3171+
SmallVector<PackExpansionExpr *, 1> getCapturedExpansions(AnyFunctionRef func) const {
3172+
auto result = CapturedExpansions.find(func);
3173+
if (result == CapturedExpansions.end())
3174+
return {};
3175+
3176+
return result->second;
3177+
}
3178+
3179+
void setCapturedExpansions(AnyFunctionRef func, SmallVector<PackExpansionExpr *, 1> exprs) {
3180+
assert(CapturedExpansions.count(func) == 0 && "Cannot reset captured expansions");
3181+
CapturedExpansions.insert({func, exprs});
3182+
}
3183+
31673184
TypeVariableType *getKeyPathValueType(const KeyPathExpr *keyPath) const {
31683185
auto result = getKeyPathValueTypeIfAvailable(keyPath);
31693186
assert(result);

‎lib/AST/CaptureInfo.cpp

Lines changed: 41 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,15 +13,28 @@
1313
#include "swift/AST/CaptureInfo.h"
1414
#include "swift/AST/ASTContext.h"
1515
#include "swift/AST/Decl.h"
16+
#include "swift/AST/GenericEnvironment.h"
1617
#include "llvm/Support/raw_ostream.h"
1718

1819
using namespace swift;
1920

21+
ArrayRef<CapturedValue>
22+
CaptureInfo::CaptureInfoStorage::getCaptures() const {
23+
return llvm::ArrayRef(this->getTrailingObjects<CapturedValue>(), NumCapturedValues);
24+
}
25+
26+
ArrayRef<GenericEnvironment *>
27+
CaptureInfo::CaptureInfoStorage::getGenericEnvironments() const {
28+
return llvm::ArrayRef(this->getTrailingObjects<GenericEnvironment *>(), NumGenericEnvironments);
29+
}
30+
2031
//===----------------------------------------------------------------------===//
2132
// MARK: CaptureInfo
2233
//===----------------------------------------------------------------------===//
2334

24-
CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
35+
CaptureInfo::CaptureInfo(ASTContext &ctx,
36+
ArrayRef<CapturedValue> captures,
37+
ArrayRef<GenericEnvironment *> genericEnv,
2538
DynamicSelfType *dynamicSelf,
2639
OpaqueValueExpr *opaqueValue,
2740
bool genericParamCaptures) {
@@ -30,30 +43,43 @@ CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
3043
static_assert(IsTriviallyDestructible<CaptureInfo::CaptureInfoStorage>::value,
3144
"Capture info is alloc'd on the ASTContext and not destroyed");
3245

46+
// This is the only kind of local generic environment we can capture right now.
47+
#ifndef NDEBUG
48+
for (auto *env : genericEnv) {
49+
assert(env->getKind() == GenericEnvironment::Kind::OpenedElement);
50+
}
51+
#endif
52+
3353
OptionSet<Flags> flags;
3454
if (genericParamCaptures)
3555
flags |= Flags::HasGenericParamCaptures;
3656

37-
if (captures.empty() && !dynamicSelf && !opaqueValue) {
57+
if (captures.empty() && genericEnv.empty() && !dynamicSelf && !opaqueValue) {
3858
*this = CaptureInfo::empty();
3959
StorageAndFlags.setInt(flags);
4060
return;
4161
}
4262

4363
size_t storageToAlloc =
44-
CaptureInfoStorage::totalSizeToAlloc<CapturedValue>(captures.size());
64+
CaptureInfoStorage::totalSizeToAlloc<CapturedValue,
65+
GenericEnvironment *>(captures.size(),
66+
genericEnv.size());
4567
void *storageBuf = ctx.Allocate(storageToAlloc, alignof(CaptureInfoStorage));
46-
auto *storage = new (storageBuf) CaptureInfoStorage(captures.size(),
47-
dynamicSelf,
48-
opaqueValue);
68+
auto *storage = new (storageBuf) CaptureInfoStorage(dynamicSelf,
69+
opaqueValue,
70+
captures.size(),
71+
genericEnv.size());
4972
StorageAndFlags.setPointerAndInt(storage, flags);
5073
std::uninitialized_copy(captures.begin(), captures.end(),
5174
storage->getTrailingObjects<CapturedValue>());
75+
std::uninitialized_copy(genericEnv.begin(), genericEnv.end(),
76+
storage->getTrailingObjects<GenericEnvironment *>());
5277
}
5378

5479
CaptureInfo CaptureInfo::empty() {
55-
static const CaptureInfoStorage empty{0, /*dynamicSelf*/nullptr,
56-
/*opaqueValue*/nullptr};
80+
static const CaptureInfoStorage empty{/*dynamicSelf*/nullptr,
81+
/*opaqueValue*/nullptr,
82+
0, 0};
5783
CaptureInfo result;
5884
result.StorageAndFlags.setPointer(&empty);
5985
return result;
@@ -109,6 +135,13 @@ void CaptureInfo::print(raw_ostream &OS) const {
109135
OS << "<noescape>";
110136
},
111137
[&] { OS << ", "; });
138+
139+
interleave(getGenericEnvironments(),
140+
[&](GenericEnvironment *genericEnv) {
141+
OS << " shape_class=";
142+
OS << genericEnv->getOpenedElementShapeClass();
143+
},
144+
[&] { OS << ","; });
112145
OS << ')';
113146
}
114147

‎lib/SIL/IR/TypeLowering.cpp

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4150,6 +4150,9 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
41504150
// that IRGen can pass dynamic 'Self' metadata.
41514151
std::optional<CapturedValue> selfCapture;
41524152

4153+
// Captured pack element environments.
4154+
llvm::SetVector<GenericEnvironment *> genericEnv;
4155+
41534156
bool capturesGenericParams = false;
41544157
DynamicSelfType *capturesDynamicSelf = nullptr;
41554158
OpaqueValueExpr *capturesOpaqueValue = nullptr;
@@ -4176,6 +4179,13 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
41764179
if (captureInfo.hasOpaqueValueCapture())
41774180
capturesOpaqueValue = captureInfo.getOpaqueValue();
41784181

4182+
// Any set of mutually-recursive local functions will capture the same
4183+
// element environments, because we can't "cross" a pack expansion
4184+
// expression.
4185+
for (auto *env : captureInfo.getGenericEnvironments()) {
4186+
genericEnv.insert(env);
4187+
}
4188+
41794189
for (auto capture : captureInfo.getCaptures()) {
41804190
if (!capture.isLocalCapture())
41814191
continue;
@@ -4399,8 +4409,9 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
43994409
}
44004410

44014411
// Cache the uniqued set of transitive captures.
4402-
CaptureInfo info{Context, resultingCaptures, capturesDynamicSelf,
4403-
capturesOpaqueValue, capturesGenericParams};
4412+
CaptureInfo info(Context, resultingCaptures, genericEnv.getArrayRef(),
4413+
capturesDynamicSelf, capturesOpaqueValue,
4414+
capturesGenericParams);
44044415
auto inserted = LoweredCaptures.insert({fn, info});
44054416
assert(inserted.second && "already in map?!");
44064417
(void)inserted;

‎lib/Sema/TypeCheckCaptures.cpp

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,16 @@ class FindCapturedVars : public ASTWalker {
3838
ASTContext &Context;
3939
SmallVector<CapturedValue, 4> Captures;
4040
llvm::SmallDenseMap<ValueDecl*, unsigned, 4> captureEntryNumber;
41+
42+
/// A stack of pack element environments we're currently walking into.
43+
/// A reference to an element archetype defined by one of these is not
44+
/// a capture.
45+
llvm::SetVector<GenericEnvironment *> VisitingEnvironments;
46+
47+
/// A set of pack element environments we've encountered that were not
48+
/// in the above stack; those are the captures.
49+
llvm::SetVector<GenericEnvironment *> CapturedEnvironments;
50+
4151
SourceLoc GenericParamCaptureLoc;
4252
SourceLoc DynamicSelfCaptureLoc;
4353
DynamicSelfType *DynamicSelf = nullptr;
@@ -65,7 +75,8 @@ class FindCapturedVars : public ASTWalker {
6575
dynamicSelfToRecord = DynamicSelf;
6676
}
6777

68-
return CaptureInfo(Context, Captures, dynamicSelfToRecord, OpaqueValue,
78+
return CaptureInfo(Context, Captures, CapturedEnvironments.getArrayRef(),
79+
dynamicSelfToRecord, OpaqueValue,
6980
HasGenericParamCaptures);
7081
}
7182

@@ -148,9 +159,17 @@ class FindCapturedVars : public ASTWalker {
148159
// perform it accurately.
149160
if (type->hasArchetype() || type->hasTypeParameter()) {
150161
type.walk(TypeCaptureWalker(ObjC, [&](Type t) {
151-
if ((t->is<ArchetypeType>() ||
162+
// Record references to element archetypes that were bound
163+
// outside the body of the current closure.
164+
if (auto *element = t->getAs<ElementArchetypeType>()) {
165+
auto *env = element->getGenericEnvironment();
166+
if (VisitingEnvironments.count(env) == 0)
167+
CapturedEnvironments.insert(env);
168+
}
169+
170+
if ((t->is<PrimaryArchetypeType>() ||
171+
t->is<PackArchetypeType>() ||
152172
t->is<GenericTypeParamType>()) &&
153-
!t->isOpenedExistential() &&
154173
!HasGenericParamCaptures) {
155174
GenericParamCaptureLoc = loc;
156175
HasGenericParamCaptures = true;
@@ -619,6 +638,25 @@ class FindCapturedVars : public ASTWalker {
619638
}
620639
}
621640

641+
if (auto expansion = dyn_cast<PackExpansionExpr>(E)) {
642+
if (auto *env = expansion->getGenericEnvironment()) {
643+
assert(VisitingEnvironments.count(env) == 0);
644+
VisitingEnvironments.insert(env);
645+
}
646+
}
647+
648+
return Action::Continue(E);
649+
}
650+
651+
PostWalkResult<Expr *> walkToExprPost(Expr *E) override {
652+
if (auto expansion = dyn_cast<PackExpansionExpr>(E)) {
653+
if (auto *env = expansion->getGenericEnvironment()) {
654+
assert(env == VisitingEnvironments.back());
655+
VisitingEnvironments.pop_back();
656+
(void) env;
657+
}
658+
}
659+
622660
return Action::Continue(E);
623661
}
624662
};

0 commit comments

Comments
 (0)