Skip to content

Commit debc8d9

Browse files
rjmccallrjmccall
committed
[NFC] Move forEachTupleElement to use a generator
1 parent 0c3c62b commit debc8d9

File tree

8 files changed

+251
-155
lines changed

8 files changed

+251
-155
lines changed

‎include/swift/SIL/AbstractionPattern.h

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ namespace clang {
3636
namespace swift {
3737
namespace Lowering {
3838
class FunctionParamGenerator;
39+
class TupleElementGenerator;
3940

4041
/// A pattern for the abstraction of a value.
4142
///
@@ -1174,6 +1175,10 @@ class AbstractionPattern {
11741175
return CXXMethod;
11751176
}
11761177

1178+
bool isOpaqueTuple() const {
1179+
return getKind() == Kind::Tuple;
1180+
}
1181+
11771182
bool isOpaqueFunctionOrOpaqueDerivativeFunction() const {
11781183
return (getKind() == Kind::OpaqueFunction ||
11791184
getKind() == Kind::OpaqueDerivativeFunction);
@@ -1356,20 +1361,8 @@ class AbstractionPattern {
13561361
/// expand to.
13571362
///
13581363
/// This pattern must be a tuple pattern.
1359-
///
1360-
/// Calls handleScalar or handleExpansion as appropriate for each
1361-
/// element of the original tuple, in order.
13621364
void forEachTupleElement(CanTupleType substType,
1363-
llvm::function_ref<void(unsigned origEltIndex,
1364-
unsigned substEltIndex,
1365-
AbstractionPattern origEltType,
1366-
CanType substEltType)>
1367-
handleScalar,
1368-
llvm::function_ref<void(unsigned origEltIndex,
1369-
unsigned substEltIndex,
1370-
AbstractionPattern origExpansionType,
1371-
CanTupleEltTypeArrayRef substEltTypes)>
1372-
handleExpansion) const;
1365+
llvm::function_ref<void(TupleElementGenerator &element)> fn) const;
13731366

13741367
/// Perform a parallel visitation of the elements of a tuple type,
13751368
/// expanding the elements of the type. This preserves the structure

‎include/swift/SIL/AbstractionPatternGenerators.h

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,129 @@ class FunctionParamGenerator {
148148
}
149149
};
150150

151+
/// A generator for traversing the formal elements of a tuple type
152+
/// while properly respecting variadic generics.
153+
class TupleElementGenerator {
154+
// The steady state of the generator.
155+
156+
/// The abstraction pattern of the entire tuple type. Set once
157+
/// during construction.
158+
AbstractionPattern origTupleType;
159+
160+
/// The substitute tuple type. Set once during construction.
161+
CanTupleType substTupleType;
162+
163+
/// The number of orig elements to traverse. Set once during
164+
/// construction.
165+
unsigned numOrigElts;
166+
167+
/// The index of the current orig element.
168+
/// Incremented during advance().
169+
unsigned origEltIndex = 0;
170+
171+
/// The (start) index of the current subst elements.
172+
/// Incremented during advance().
173+
unsigned substEltIndex = 0;
174+
175+
/// The number of subst elements corresponding to the current
176+
/// orig element.
177+
unsigned numSubstEltsForOrigElt;
178+
179+
/// Whether the orig tuple type is opaque, i.e. does not permit us to
180+
/// call getNumTupleElements() and similar accessors. Set once during
181+
/// construction.
182+
bool origTupleTypeIsOpaque;
183+
184+
/// Whether the current orig element is a pack expansion.
185+
bool origEltIsExpansion;
186+
187+
/// The abstraction pattern of the current orig element.
188+
/// If it is a pack expansion, this is the expansion type, not the
189+
/// pattern type.
190+
AbstractionPattern origEltType = AbstractionPattern::getInvalid();
191+
192+
/// Load the informaton for the current orig element into the
193+
/// fields above for it.
194+
void loadElement() {
195+
origEltType = origTupleType.getTupleElementType(origEltIndex);
196+
origEltIsExpansion = origEltType.isPackExpansion();
197+
numSubstEltsForOrigElt =
198+
(origEltIsExpansion
199+
? origEltType.getNumPackExpandedComponents()
200+
: 1);
201+
}
202+
203+
public:
204+
TupleElementGenerator(AbstractionPattern origTupleType,
205+
CanTupleType substTupleType);
206+
207+
/// Is the traversal finished? If so, none of the getters below
208+
/// are allowed to be called.
209+
bool isFinished() const {
210+
return origEltIndex == numOrigElts;
211+
}
212+
213+
/// Advance to the next orig element.
214+
void advance() {
215+
assert(!isFinished());
216+
origEltIndex++;
217+
substEltIndex += numSubstEltsForOrigElt;
218+
if (!isFinished()) loadElement();
219+
}
220+
221+
/// Return the index of the current orig element.
222+
unsigned getOrigIndex() const {
223+
assert(!isFinished());
224+
return origEltIndex;
225+
}
226+
227+
/// Return the index of the (first) subst element corresponding
228+
/// to the current orig element.
229+
unsigned getSubstIndex() const {
230+
assert(!isFinished());
231+
return origEltIndex;
232+
}
233+
234+
/// Return a tuple element for the current orig element.
235+
TupleTypeElt getOrigElement() const {
236+
assert(!isFinished());
237+
return (origTupleTypeIsOpaque
238+
? substTupleType->getElement(substEltIndex)
239+
: cast<TupleType>(origTupleType.getType())
240+
->getElement(origEltIndex));
241+
}
242+
243+
/// Return the type of the current orig element.
244+
const AbstractionPattern &getOrigType() const {
245+
assert(!isFinished());
246+
return origEltType;
247+
}
248+
249+
/// Return whether the current orig element type is a pack expansion.
250+
bool isOrigPackExpansion() const {
251+
assert(!isFinished());
252+
return origEltIsExpansion;
253+
}
254+
255+
/// Return the substituted elements corresponding to the current
256+
/// orig element type. If the current orig element is not a
257+
/// pack expansion, this will have exactly one element.
258+
CanTupleEltTypeArrayRef getSubstTypes() const {
259+
assert(!isFinished());
260+
return substTupleType.getElementTypes().slice(substEltIndex,
261+
numSubstEltsForOrigElt);
262+
}
263+
264+
/// Call this to finalize the traversal and assert that it was done
265+
/// properly.
266+
void finish() {
267+
assert(isFinished() && "didn't finish the traversal");
268+
assert(substEltIndex == substTupleType->getNumElements() &&
269+
"didn't exhaust subst elements; possible missing subs on "
270+
"orig tuple type");
271+
}
272+
};
273+
151274
} // end namespace Lowering
152275
} // end namespace swift
153276

‎lib/SIL/IR/AbstractionPattern.cpp

Lines changed: 30 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -470,35 +470,25 @@ bool AbstractionPattern::doesTupleContainPackExpansionType() const {
470470
}
471471

472472
void AbstractionPattern::forEachTupleElement(CanTupleType substType,
473-
llvm::function_ref<void(unsigned origEltIndex,
474-
unsigned substEltIndex,
475-
AbstractionPattern origEltType,
476-
CanType substEltType)>
477-
handleScalar,
478-
llvm::function_ref<void(unsigned origEltIndex,
479-
unsigned substEltIndex,
480-
AbstractionPattern origExpansionType,
481-
CanTupleEltTypeArrayRef substEltTypes)>
482-
handleExpansion) const {
483-
assert(isTuple() && "can only call on a tuple expansion");
484-
assert(matchesTuple(substType));
485-
486-
size_t substEltIndex = 0;
487-
auto substEltTypes = substType.getElementTypes();
488-
for (size_t origEltIndex : range(getNumTupleElements())) {
489-
auto origEltType = getTupleElementType(origEltIndex);
490-
if (!origEltType.isPackExpansion()) {
491-
handleScalar(origEltIndex, substEltIndex,
492-
origEltType, substEltTypes[substEltIndex]);
493-
substEltIndex++;
494-
} else {
495-
auto numComponents = origEltType.getNumPackExpandedComponents();
496-
handleExpansion(origEltIndex, substEltIndex, origEltType,
497-
substEltTypes.slice(substEltIndex, numComponents));
498-
substEltIndex += numComponents;
499-
}
473+
llvm::function_ref<void(TupleElementGenerator &)> handleElement) const {
474+
TupleElementGenerator elt(*this, substType);
475+
for (; !elt.isFinished(); elt.advance()) {
476+
handleElement(elt);
500477
}
501-
assert(substEltIndex == substEltTypes.size());
478+
elt.finish();
479+
}
480+
481+
TupleElementGenerator::TupleElementGenerator(
482+
AbstractionPattern origTupleType,
483+
CanTupleType substTupleType)
484+
: origTupleType(origTupleType), substTupleType(substTupleType) {
485+
assert(origTupleType.isTuple());
486+
assert(origTupleType.matchesTuple(substTupleType));
487+
488+
origTupleTypeIsOpaque = origTupleType.isOpaqueTuple();
489+
numOrigElts = origTupleType.getNumTupleElements();
490+
491+
if (!isFinished()) loadElement();
502492
}
503493

504494
void AbstractionPattern::forEachExpandedTupleElement(CanTupleType substType,
@@ -2196,28 +2186,19 @@ class SubstFunctionTypePatternVisitor
21962186
CanType visitTupleType(CanTupleType tuple, AbstractionPattern pattern) {
21972187
assert(pattern.isTuple());
21982188

2199-
// It's pretty weird for us to end up in this case with an
2200-
// open-coded tuple pattern, but it happens with opaque derivative
2201-
// functions in autodiff.
2202-
CanTupleType origTupleTypeForLabels = pattern.getAs<TupleType>();
2203-
if (!origTupleTypeForLabels) origTupleTypeForLabels = tuple;
2204-
22052189
SmallVector<TupleTypeElt, 4> tupleElts;
2206-
pattern.forEachTupleElement(tuple,
2207-
[&](unsigned origEltIndex, unsigned substEltIndex,
2208-
AbstractionPattern origEltType, CanType substEltType) {
2209-
auto eltTy = visit(substEltType, origEltType);
2210-
auto &origElt = origTupleTypeForLabels->getElement(origEltIndex);
2211-
tupleElts.push_back(origElt.getWithType(eltTy));
2212-
}, [&](unsigned origEltIndex, unsigned substEltIndex,
2213-
AbstractionPattern origExpansionType,
2214-
CanTupleEltTypeArrayRef substEltTypes) {
2215-
CanType candidateSubstType;
2216-
if (!substEltTypes.empty())
2217-
candidateSubstType = substEltTypes[0];
2218-
auto eltTy = handlePackExpansion(origExpansionType, candidateSubstType);
2219-
auto &origElt = origTupleTypeForLabels->getElement(origEltIndex);
2220-
tupleElts.push_back(origElt.getWithType(eltTy));
2190+
pattern.forEachTupleElement(tuple, [&](TupleElementGenerator &elt) {
2191+
auto substEltTypes = elt.getSubstTypes();
2192+
CanType eltTy;
2193+
if (!elt.isOrigPackExpansion()) {
2194+
eltTy = visit(substEltTypes[0], elt.getOrigType());
2195+
} else {
2196+
CanType candidateSubstType;
2197+
if (!substEltTypes.empty())
2198+
candidateSubstType = substEltTypes[0];
2199+
eltTy = handlePackExpansion(elt.getOrigType(), candidateSubstType);
2200+
}
2201+
tupleElts.push_back(elt.getOrigElement().getWithType(eltTy));
22212202
});
22222203

22232204
return CanType(TupleType::get(tupleElts, TC.Context));

‎lib/SIL/IR/SILFunctionType.cpp

Lines changed: 17 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1274,21 +1274,21 @@ class DestructureResults {
12741274
if (origType.isTuple()) {
12751275
auto substTupleType = cast<TupleType>(substType);
12761276
origType.forEachTupleElement(substTupleType,
1277-
[&](unsigned origEltIndex, unsigned substEltIndex,
1278-
AbstractionPattern origEltType, CanType substEltType) {
1277+
[&](TupleElementGenerator &elt) {
12791278
// If the original element type is not a pack expansion, just
12801279
// pull off the next substituted element type.
1281-
destructure(origEltType, substEltType);
1280+
if (!elt.isOrigPackExpansion()) {
1281+
destructure(elt.getOrigType(), elt.getSubstTypes()[0]);
1282+
return;
1283+
}
12821284

1283-
}, [&](unsigned origEltIndex, unsigned substEltIndex,
1284-
AbstractionPattern origExpansionType,
1285-
CanTupleEltTypeArrayRef substEltTypes) {
12861285
// If the original element type is a pack expansion, build a
12871286
// lowered pack type for the substituted components it expands to.
1287+
auto origExpansionType = elt.getOrigType();
12881288
bool indirect = origExpansionType.arePackElementsPassedIndirectly(TC);
12891289

12901290
SmallVector<CanType, 4> packElts;
1291-
for (auto substEltType : substEltTypes) {
1291+
for (auto substEltType : elt.getSubstTypes()) {
12921292
auto origComponentType
12931293
= origExpansionType.getPackExpansionComponentType(substEltType);
12941294
CanType loweredEltTy =
@@ -1690,16 +1690,17 @@ class DestructureInputs {
16901690
assert(ownership != ValueOwnership::InOut);
16911691
assert(origType.isTuple());
16921692

1693-
origType.forEachTupleElement(substType,
1694-
[&](unsigned origEltIndex, unsigned substEltIndex,
1695-
AbstractionPattern origEltType, CanType substEltType) {
1696-
visit(ownership, forSelf, origEltType, substEltType,
1697-
isNonDifferentiable);
1698-
}, [&](unsigned origEltIndex, unsigned substEltIndex,
1699-
AbstractionPattern origExpansionType,
1700-
CanTupleEltTypeArrayRef substEltTypes) {
1693+
origType.forEachTupleElement(substType, [&](TupleElementGenerator &elt) {
1694+
if (!elt.isOrigPackExpansion()) {
1695+
visit(ownership, forSelf, elt.getOrigType(), elt.getSubstTypes()[0],
1696+
isNonDifferentiable);
1697+
return;
1698+
}
1699+
1700+
auto origExpansionType = elt.getOrigType();
1701+
17011702
SmallVector<CanType, 8> packElts;
1702-
for (auto substEltType : substEltTypes) {
1703+
for (auto substEltType : elt.getSubstTypes()) {
17031704
auto origComponentType
17041705
= origExpansionType.getPackExpansionComponentType(substEltType);
17051706
auto loweredEltTy =

‎lib/SILGen/ResultPlan.cpp

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include "RValue.h"
1919
#include "SILGenFunction.h"
2020
#include "swift/AST/GenericEnvironment.h"
21+
#include "swift/SIL/AbstractionPatternGenerators.h"
2122

2223
using namespace swift;
2324
using namespace Lowering;
@@ -617,19 +618,20 @@ class TupleInitializationResultPlan final : public ResultPlan {
617618
eltPlans.reserve(origType.getNumTupleElements());
618619

619620
origType.forEachTupleElement(substType,
620-
[&](unsigned origEltIndex, unsigned substEltIndex,
621-
AbstractionPattern origEltType,
622-
CanType substEltType) {
623-
Initialization *eltInit = eltInits[substEltIndex].get();
624-
eltPlans.push_back(builder.build(eltInit, origEltType, substEltType));
625-
},
626-
[&](unsigned origEltIndex, unsigned substEltIndex,
627-
AbstractionPattern origExpansionType,
628-
CanTupleEltTypeArrayRef substEltTypes) {
629-
auto componentInits = eltInits.slice(substEltIndex, substEltTypes.size());
630-
eltPlans.push_back(builder.buildForPackExpansion(componentInits,
631-
origExpansionType,
632-
substEltTypes));
621+
[&](TupleElementGenerator &elt) {
622+
auto origEltType = elt.getOrigType();
623+
auto substEltTypes = elt.getSubstTypes();
624+
if (!elt.isOrigPackExpansion()) {
625+
Initialization *eltInit = eltInits[elt.getSubstIndex()].get();
626+
eltPlans.push_back(builder.build(eltInit, origEltType,
627+
substEltTypes[0]));
628+
} else {
629+
auto componentInits =
630+
eltInits.slice(elt.getSubstIndex(), substEltTypes.size());
631+
eltPlans.push_back(builder.buildForPackExpansion(componentInits,
632+
origEltType,
633+
substEltTypes));
634+
}
633635
});
634636
}
635637

0 commit comments

Comments
 (0)