Skip to content

Commit a721ef8

Browse files
slavapestovslavapestov
committed
AST: Record captured generic environments in CaptureInfo
1 parent a8bd4e7 commit a721ef8

File tree

4 files changed

+172
-25
lines changed

4 files changed

+172
-25
lines changed

‎include/swift/AST/CaptureInfo.h

Lines changed: 34 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.
@@ -140,19 +141,27 @@ class DynamicSelfType;
140141
/// Stores information about captured variables.
141142
class CaptureInfo {
142143
class CaptureInfoStorage final
143-
: public llvm::TrailingObjects<CaptureInfoStorage, CapturedValue> {
144+
: public llvm::TrailingObjects<CaptureInfoStorage,
145+
CapturedValue,
146+
GenericEnvironment *> {
144147

145148
DynamicSelfType *DynamicSelf;
146149
OpaqueValueExpr *OpaqueValue;
147-
unsigned Count;
150+
unsigned NumCapturedValues;
151+
unsigned NumGenericEnvironments;
152+
148153
public:
149-
explicit CaptureInfoStorage(unsigned count, DynamicSelfType *dynamicSelf,
150-
OpaqueValueExpr *opaqueValue)
151-
: DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue), Count(count) { }
154+
explicit CaptureInfoStorage(DynamicSelfType *dynamicSelf,
155+
OpaqueValueExpr *opaqueValue,
156+
unsigned numCapturedValues,
157+
unsigned numGenericEnvironments)
158+
: DynamicSelf(dynamicSelf), OpaqueValue(opaqueValue),
159+
NumCapturedValues(numCapturedValues),
160+
NumGenericEnvironments(numGenericEnvironments) { }
152161

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

157166
DynamicSelfType *getDynamicSelfType() const {
158167
return DynamicSelf;
@@ -161,6 +170,10 @@ class CaptureInfo {
161170
OpaqueValueExpr *getOpaqueValue() const {
162171
return OpaqueValue;
163172
}
173+
174+
unsigned numTrailingObjects(OverloadToken<CapturedValue>) const {
175+
return NumCapturedValues;
176+
}
164177
};
165178

166179
enum class Flags : unsigned {
@@ -173,9 +186,11 @@ class CaptureInfo {
173186
public:
174187
/// The default-constructed CaptureInfo is "not yet computed".
175188
CaptureInfo() = default;
176-
CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
189+
CaptureInfo(ASTContext &ctx,
190+
ArrayRef<CapturedValue> captures,
177191
DynamicSelfType *dynamicSelf, OpaqueValueExpr *opaqueValue,
178-
bool genericParamCaptures);
192+
bool genericParamCaptures,
193+
ArrayRef<GenericEnvironment *> genericEnv=ArrayRef<GenericEnvironment*>());
179194

180195
/// A CaptureInfo representing no captures at all.
181196
static CaptureInfo empty();
@@ -189,6 +204,7 @@ class CaptureInfo {
189204
!hasDynamicSelfCapture() && !hasOpaqueValueCapture();
190205
}
191206

207+
/// Returns all captured values and opaque expressions.
192208
ArrayRef<CapturedValue> getCaptures() const {
193209
// FIXME: Ideally, everywhere that synthesizes a function should include
194210
// its capture info.
@@ -206,7 +222,14 @@ class CaptureInfo {
206222
/// \returns true if getLocalCaptures() will return a non-empty list.
207223
bool hasLocalCaptures() const;
208224

209-
/// \returns true if the function captures any generic type parameters.
225+
/// Returns all captured pack element environments.
226+
ArrayRef<GenericEnvironment *> getGenericEnvironments() const {
227+
assert(hasBeenComputed());
228+
return StorageAndFlags.getPointer()->getGenericEnvironments();
229+
}
230+
231+
/// \returns true if the function captures the primary generic environment
232+
/// from its innermost declaration context.
210233
bool hasGenericParamCaptures() const {
211234
// FIXME: Ideally, everywhere that synthesizes a function should include
212235
// its capture info.

‎lib/AST/CaptureInfo.cpp

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,47 +13,72 @@
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

2435
CaptureInfo::CaptureInfo(ASTContext &ctx, ArrayRef<CapturedValue> captures,
2536
DynamicSelfType *dynamicSelf,
2637
OpaqueValueExpr *opaqueValue,
27-
bool genericParamCaptures) {
38+
bool genericParamCaptures,
39+
ArrayRef<GenericEnvironment *> genericEnv) {
2840
static_assert(IsTriviallyDestructible<CapturedValue>::value,
2941
"Capture info is alloc'd on the ASTContext and not destroyed");
3042
static_assert(IsTriviallyDestructible<CaptureInfo::CaptureInfoStorage>::value,
3143
"Capture info is alloc'd on the ASTContext and not destroyed");
3244

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

37-
if (captures.empty() && !dynamicSelf && !opaqueValue) {
56+
if (captures.empty() && genericEnv.empty() && !dynamicSelf && !opaqueValue) {
3857
*this = CaptureInfo::empty();
3958
StorageAndFlags.setInt(flags);
4059
return;
4160
}
4261

4362
size_t storageToAlloc =
44-
CaptureInfoStorage::totalSizeToAlloc<CapturedValue>(captures.size());
63+
CaptureInfoStorage::totalSizeToAlloc<CapturedValue,
64+
GenericEnvironment *>(captures.size(),
65+
genericEnv.size());
4566
void *storageBuf = ctx.Allocate(storageToAlloc, alignof(CaptureInfoStorage));
46-
auto *storage = new (storageBuf) CaptureInfoStorage(captures.size(),
47-
dynamicSelf,
48-
opaqueValue);
67+
auto *storage = new (storageBuf) CaptureInfoStorage(dynamicSelf,
68+
opaqueValue,
69+
captures.size(),
70+
genericEnv.size());
4971
StorageAndFlags.setPointerAndInt(storage, flags);
5072
std::uninitialized_copy(captures.begin(), captures.end(),
5173
storage->getTrailingObjects<CapturedValue>());
74+
std::uninitialized_copy(genericEnv.begin(), genericEnv.end(),
75+
storage->getTrailingObjects<GenericEnvironment *>());
5276
}
5377

5478
CaptureInfo CaptureInfo::empty() {
55-
static const CaptureInfoStorage empty{0, /*dynamicSelf*/nullptr,
56-
/*opaqueValue*/nullptr};
79+
static const CaptureInfoStorage empty{/*dynamicSelf*/nullptr,
80+
/*opaqueValue*/nullptr,
81+
0, 0};
5782
CaptureInfo result;
5883
result.StorageAndFlags.setPointer(&empty);
5984
return result;
@@ -136,6 +161,13 @@ void CaptureInfo::print(raw_ostream &OS) const {
136161
OS << "<noescape>";
137162
},
138163
[&] { OS << ", "; });
164+
165+
interleave(getGenericEnvironments(),
166+
[&](GenericEnvironment *genericEnv) {
167+
OS << " shape_class=";
168+
OS << genericEnv->getOpenedElementShapeClass();
169+
},
170+
[&] { OS << ","; });
139171
OS << ')';
140172
}
141173

‎lib/SIL/IR/TypeLowering.cpp

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

4165+
// Captured pack element environments.
4166+
llvm::SetVector<GenericEnvironment *> genericEnv;
4167+
41654168
bool capturesGenericParams = false;
41664169
DynamicSelfType *capturesDynamicSelf = nullptr;
41674170
OpaqueValueExpr *capturesOpaqueValue = nullptr;
@@ -4180,6 +4183,13 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
41804183
if (captureInfo.hasOpaqueValueCapture())
41814184
capturesOpaqueValue = captureInfo.getOpaqueValue();
41824185

4186+
// Any set of mutually-recursive local functions will capture the same
4187+
// element environments, because we can't "cross" a pack expansion
4188+
// expression.
4189+
for (auto *env : captureInfo.getGenericEnvironments()) {
4190+
genericEnv.insert(env);
4191+
}
4192+
41834193
SmallVector<CapturedValue, 4> localCaptures;
41844194
captureInfo.getLocalCaptures(localCaptures);
41854195
for (auto capture : localCaptures) {
@@ -4397,8 +4407,9 @@ TypeConverter::getLoweredLocalCaptures(SILDeclRef fn) {
43974407
}
43984408

43994409
// Cache the uniqued set of transitive captures.
4400-
CaptureInfo info{Context, resultingCaptures, capturesDynamicSelf,
4401-
capturesOpaqueValue, capturesGenericParams};
4410+
CaptureInfo info(Context, resultingCaptures,
4411+
capturesDynamicSelf, capturesOpaqueValue,
4412+
capturesGenericParams, genericEnv.getArrayRef());
44024413
auto inserted = LoweredCaptures.insert({fn, info});
44034414
assert(inserted.second && "already in map?!");
44044415
(void)inserted;

‎lib/Sema/TypeCheckCaptures.cpp

Lines changed: 85 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,20 @@ class FindCapturedVars : public ASTWalker {
3737
ASTContext &Context;
3838
SmallVector<CapturedValue, 4> Captures;
3939
llvm::SmallDenseMap<ValueDecl*, unsigned, 4> captureEntryNumber;
40+
41+
/// We track the pack expansion expressions in ForEachStmts, because
42+
/// their local generics remain in scope until the end of the statement.
43+
llvm::DenseSet<PackExpansionExpr *> ForEachPatternSequences;
44+
45+
/// A stack of pack element environments we're currently walking into.
46+
/// A reference to an element archetype defined by one of these is not
47+
/// a capture.
48+
llvm::SetVector<GenericEnvironment *> VisitingEnvironments;
49+
50+
/// A set of pack element environments we've encountered that were not
51+
/// in the above stack; those are the captures.
52+
llvm::SetVector<GenericEnvironment *> CapturedEnvironments;
53+
4054
SourceLoc GenericParamCaptureLoc;
4155
SourceLoc DynamicSelfCaptureLoc;
4256
DynamicSelfType *DynamicSelf = nullptr;
@@ -64,8 +78,9 @@ class FindCapturedVars : public ASTWalker {
6478
dynamicSelfToRecord = DynamicSelf;
6579
}
6680

67-
return CaptureInfo(Context, Captures, dynamicSelfToRecord, OpaqueValue,
68-
HasGenericParamCaptures);
81+
return CaptureInfo(Context, Captures, dynamicSelfToRecord,
82+
OpaqueValue, HasGenericParamCaptures,
83+
CapturedEnvironments.getArrayRef());
6984
}
7085

7186
bool hasGenericParamCaptures() const {
@@ -147,9 +162,17 @@ class FindCapturedVars : public ASTWalker {
147162
// perform it accurately.
148163
if (type->hasArchetype() || type->hasTypeParameter()) {
149164
type.walk(TypeCaptureWalker(ObjC, [&](Type t) {
150-
if ((t->is<ArchetypeType>() ||
165+
// Record references to element archetypes that were bound
166+
// outside the body of the current closure.
167+
if (auto *element = t->getAs<ElementArchetypeType>()) {
168+
auto *env = element->getGenericEnvironment();
169+
if (VisitingEnvironments.count(env) == 0)
170+
CapturedEnvironments.insert(env);
171+
}
172+
173+
if ((t->is<PrimaryArchetypeType>() ||
174+
t->is<PackArchetypeType>() ||
151175
t->is<GenericTypeParamType>()) &&
152-
!t->isOpenedExistential() &&
153176
!HasGenericParamCaptures) {
154177
GenericParamCaptureLoc = loc;
155178
HasGenericParamCaptures = true;
@@ -613,8 +636,66 @@ class FindCapturedVars : public ASTWalker {
613636
}
614637
}
615638

639+
if (auto expansion = dyn_cast<PackExpansionExpr>(E)) {
640+
if (auto *env = expansion->getGenericEnvironment()) {
641+
assert(VisitingEnvironments.count(env) == 0);
642+
VisitingEnvironments.insert(env);
643+
}
644+
}
645+
646+
return Action::Continue(E);
647+
}
648+
649+
PostWalkResult<Expr *> walkToExprPost(Expr *E) override {
650+
if (auto expansion = dyn_cast<PackExpansionExpr>(E)) {
651+
if (auto *env = expansion->getGenericEnvironment()) {
652+
assert(env == VisitingEnvironments.back());
653+
(void) env;
654+
655+
// If this is the pack expansion of a for .. in loop, the generic
656+
// environment remains in scope until the end of the loop.
657+
if (ForEachPatternSequences.count(expansion) == 0)
658+
VisitingEnvironments.pop_back();
659+
}
660+
}
661+
616662
return Action::Continue(E);
617663
}
664+
665+
PreWalkResult<Stmt *> walkToStmtPre(Stmt *S) override {
666+
if (auto *forEachStmt = dyn_cast<ForEachStmt>(S)) {
667+
if (auto *expansion =
668+
dyn_cast<PackExpansionExpr>(forEachStmt->getParsedSequence())) {
669+
if (auto *env = expansion->getGenericEnvironment()) {
670+
// Remember this generic environment, so that it remains on the
671+
// visited stack until the end of the for .. in loop.
672+
assert(ForEachPatternSequences.count(expansion) == 0);
673+
ForEachPatternSequences.insert(expansion);
674+
}
675+
}
676+
}
677+
678+
return Action::Continue(S);
679+
}
680+
681+
PostWalkResult<Stmt *> walkToStmtPost(Stmt *S) override {
682+
if (auto *forEachStmt = dyn_cast<ForEachStmt>(S)) {
683+
if (auto *expansion =
684+
dyn_cast<PackExpansionExpr>(forEachStmt->getParsedSequence())) {
685+
if (auto *env = expansion->getGenericEnvironment()) {
686+
assert(ForEachPatternSequences.count(expansion) != 0);
687+
ForEachPatternSequences.erase(expansion);
688+
689+
// Clean up the generic environment bound by the for loop.
690+
assert(env == VisitingEnvironments.back());
691+
VisitingEnvironments.pop_back();
692+
(void) env;
693+
}
694+
}
695+
}
696+
697+
return Action::Continue(S);
698+
}
618699
};
619700

620701
} // end anonymous namespace

0 commit comments

Comments
 (0)