Skip to content

Commit e27ebf5

Browse files
Merge pull request #34471 from nate-chandler/concurrency/irgen/remaining-context-fields
[Async CC] Add remaining fields to context.
2 parents 75fc060 + 80d1f83 commit e27ebf5

File tree

2 files changed

+120
-24
lines changed

2 files changed

+120
-24
lines changed

lib/IRGen/GenCall.cpp

Lines changed: 65 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -97,12 +97,27 @@ AsyncContextLayout irgen::getAsyncContextLayout(
9797
SmallVector<const TypeInfo *, 4> typeInfos;
9898
SmallVector<SILType, 4> valTypes;
9999
SmallVector<AsyncContextLayout::ArgumentInfo, 4> paramInfos;
100+
bool isCoroutine = originalType->isCoroutine();
101+
SmallVector<SILYieldInfo, 4> yieldInfos;
100102
SmallVector<SILResultInfo, 4> indirectReturnInfos;
101103
SmallVector<SILResultInfo, 4> directReturnInfos;
102104

103105
auto parameters = substitutedType->getParameters();
104106
SILFunctionConventions fnConv(substitutedType, IGF.getSILModule());
105107

108+
auto addTaskContinuationFunction = [&]() {
109+
auto ty = SILType();
110+
auto &ti = IGF.IGM.getTaskContinuationFunctionPtrTypeInfo();
111+
valTypes.push_back(ty);
112+
typeInfos.push_back(&ti);
113+
};
114+
auto addExecutor = [&]() {
115+
auto ty = SILType();
116+
auto &ti = IGF.IGM.getSwiftExecutorPtrTypeInfo();
117+
valTypes.push_back(ty);
118+
typeInfos.push_back(&ti);
119+
};
120+
106121
// AsyncContext * __ptrauth_swift_async_context_parent Parent;
107122
{
108123
auto ty = SILType();
@@ -113,21 +128,26 @@ AsyncContextLayout irgen::getAsyncContextLayout(
113128

114129
// TaskContinuationFunction * __ptrauth_swift_async_context_resume
115130
// ResumeParent;
116-
{
117-
auto ty = SILType();
118-
auto &ti = IGF.IGM.getTaskContinuationFunctionPtrTypeInfo();
119-
valTypes.push_back(ty);
120-
typeInfos.push_back(&ti);
121-
}
131+
addTaskContinuationFunction();
122132

123133
// ExecutorRef ResumeParentExecutor;
134+
addExecutor();
135+
136+
// AsyncContextFlags Flags;
124137
{
125-
auto ty = SILType();
126-
auto &ti = IGF.IGM.getSwiftExecutorPtrTypeInfo();
138+
auto ty = SILType::getPrimitiveObjectType(
139+
BuiltinIntegerType::get(32, IGF.IGM.IRGen.SIL.getASTContext())
140+
->getCanonicalType());
141+
const auto &ti = IGF.IGM.getTypeInfo(ty);
127142
valTypes.push_back(ty);
128143
typeInfos.push_back(&ti);
129144
}
130145

146+
if (isCoroutine) {
147+
// SwiftPartialFunction * __ptrauth(...) yieldToCaller?;
148+
addTaskContinuationFunction();
149+
}
150+
131151
// SwiftError *errorResult;
132152
auto errorCanType = IGF.IGM.Context.getExceptionType();
133153
auto errorType = SILType::getPrimitiveObjectType(errorCanType);
@@ -147,15 +167,33 @@ AsyncContextLayout irgen::getAsyncContextLayout(
147167
indirectReturnInfos.push_back(indirectResult);
148168
}
149169

150-
// ResultTypes directResults...;
151-
auto directResults = fnConv.getDirectSILResults();
152-
for (auto result : directResults) {
153-
auto ty =
154-
fnConv.getSILType(result, IGF.IGM.getMaximalTypeExpansionContext());
155-
auto &ti = IGF.getTypeInfoForLowered(ty.getASTType());
156-
valTypes.push_back(ty);
157-
typeInfos.push_back(&ti);
158-
directReturnInfos.push_back(result);
170+
// union {
171+
if (isCoroutine) {
172+
// SwiftPartialFunction * __ptrauth(...) resumeFromYield?
173+
addTaskContinuationFunction();
174+
// SwiftPartialFunction * __ptrauth(...) abortFromYield?
175+
addTaskContinuationFunction();
176+
// SwiftActor * __ptrauth(...) calleeActorDuringYield?
177+
addExecutor();
178+
// YieldTypes yieldValues...
179+
for (auto yield : fnConv.getYields()) {
180+
auto ty =
181+
fnConv.getSILType(yield, IGF.IGM.getMaximalTypeExpansionContext());
182+
auto &ti = IGF.getTypeInfoForLowered(ty.getASTType());
183+
valTypes.push_back(ty);
184+
typeInfos.push_back(&ti);
185+
yieldInfos.push_back(yield);
186+
}
187+
} else {
188+
// ResultTypes directResults...;
189+
for (auto result : fnConv.getDirectSILResults()) {
190+
auto ty =
191+
fnConv.getSILType(result, IGF.IGM.getMaximalTypeExpansionContext());
192+
auto &ti = IGF.getTypeInfoForLowered(ty.getASTType());
193+
valTypes.push_back(ty);
194+
typeInfos.push_back(&ti);
195+
directReturnInfos.push_back(result);
196+
}
159197
}
160198

161199
// SelfType self?;
@@ -244,7 +282,8 @@ AsyncContextLayout irgen::getAsyncContextLayout(
244282
IGF.IGM, LayoutStrategy::Optimal, valTypes, typeInfos, IGF, originalType,
245283
substitutedType, substitutionMap, std::move(bindings),
246284
trailingWitnessInfo, errorType, canHaveValidError, paramInfos,
247-
indirectReturnInfos, directReturnInfos, localContextInfo);
285+
isCoroutine, yieldInfos, indirectReturnInfos, directReturnInfos,
286+
localContextInfo);
248287
}
249288

250289
AsyncContextLayout::AsyncContextLayout(
@@ -254,14 +293,16 @@ AsyncContextLayout::AsyncContextLayout(
254293
SubstitutionMap substitutionMap, NecessaryBindings &&bindings,
255294
Optional<TrailingWitnessInfo> trailingWitnessInfo, SILType errorType,
256295
bool canHaveValidError, ArrayRef<ArgumentInfo> argumentInfos,
296+
bool isCoroutine, ArrayRef<SILYieldInfo> yieldInfos,
257297
ArrayRef<SILResultInfo> indirectReturnInfos,
258298
ArrayRef<SILResultInfo> directReturnInfos,
259299
Optional<AsyncContextLayout::ArgumentInfo> localContextInfo)
260300
: StructLayout(IGM, /*decl=*/nullptr, LayoutKind::NonHeapObject, strategy,
261301
fieldTypeInfos, /*typeToFill*/ nullptr),
262302
IGF(IGF), originalType(originalType), substitutedType(substitutedType),
263303
substitutionMap(substitutionMap), errorType(errorType),
264-
canHaveValidError(canHaveValidError),
304+
canHaveValidError(canHaveValidError), isCoroutine(isCoroutine),
305+
yieldInfos(yieldInfos.begin(), yieldInfos.end()),
265306
directReturnInfos(directReturnInfos.begin(), directReturnInfos.end()),
266307
indirectReturnInfos(indirectReturnInfos.begin(),
267308
indirectReturnInfos.end()),
@@ -271,6 +312,11 @@ AsyncContextLayout::AsyncContextLayout(
271312
#ifndef NDEBUG
272313
assert(fieldTypeInfos.size() == fieldTypes.size() &&
273314
"type infos don't match types");
315+
if (isCoroutine) {
316+
assert(directReturnInfos.empty());
317+
} else {
318+
assert(yieldInfos.empty());
319+
}
274320
if (!bindings.empty()) {
275321
assert(fieldTypeInfos.size() >= 2 && "no field for bindings");
276322
auto fixedBindingsField =

lib/IRGen/GenCall.h

Lines changed: 55 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,8 @@ namespace irgen {
9696
Parent = 0,
9797
ResumeParent = 1,
9898
ResumeParentExecutor = 2,
99-
Error = 3,
99+
Flags = 3,
100+
YieldToParent = 4,
100101
};
101102
enum class FixedCount : unsigned {
102103
Parent = 1,
@@ -110,6 +111,8 @@ namespace irgen {
110111
SubstitutionMap substitutionMap;
111112
SILType errorType;
112113
bool canHaveValidError;
114+
bool isCoroutine;
115+
SmallVector<SILYieldInfo, 4> yieldInfos;
113116
SmallVector<SILResultInfo, 4> directReturnInfos;
114117
SmallVector<SILResultInfo, 4> indirectReturnInfos;
115118
Optional<ArgumentInfo> localContextInfo;
@@ -124,20 +127,57 @@ namespace irgen {
124127
unsigned getResumeParentExecutorIndex() {
125128
return (unsigned)FixedIndex::ResumeParentExecutor;
126129
}
127-
unsigned getErrorIndex() { return (unsigned)FixedIndex::Error; }
130+
unsigned getFlagsIndex() { return (unsigned)FixedIndex::Flags; }
131+
unsigned getYieldToParentIndex() {
132+
assert(isCoroutine);
133+
return (unsigned)FixedIndex::YieldToParent;
134+
}
135+
unsigned getErrorIndex() {
136+
return (isCoroutine ? getYieldToParentIndex() : getFlagsIndex()) + 1;
137+
}
128138
unsigned getFirstIndirectReturnIndex() {
129139
return getErrorIndex() + getErrorCount();
130140
}
131141
unsigned getIndexAfterIndirectReturns() {
132142
return getFirstIndirectReturnIndex() + getIndirectReturnCount();
133143
}
134144
unsigned getFirstDirectReturnIndex() {
145+
assert(!isCoroutine);
135146
return getIndexAfterIndirectReturns();
136147
}
137148
unsigned getIndexAfterDirectReturns() {
149+
assert(!isCoroutine);
138150
return getFirstDirectReturnIndex() + getDirectReturnCount();
139151
}
140-
unsigned getFirstArgumentIndex() { return getIndexAfterDirectReturns(); }
152+
unsigned getResumeFromYieldIndex() {
153+
assert(isCoroutine);
154+
return getIndexAfterIndirectReturns();
155+
}
156+
unsigned getAbortFromYieldIndex() {
157+
assert(isCoroutine);
158+
return getResumeFromYieldIndex() + 1;
159+
;
160+
}
161+
unsigned getCalleeExecutorDuringYieldIndex() {
162+
assert(isCoroutine);
163+
return getAbortFromYieldIndex() + 1;
164+
}
165+
unsigned getFirstYieldIndex() {
166+
assert(isCoroutine);
167+
return getCalleeExecutorDuringYieldIndex() + 1;
168+
}
169+
unsigned getIndexAfterYields() {
170+
assert(isCoroutine);
171+
return getFirstYieldIndex() + getYieldCount();
172+
}
173+
unsigned getIndexAfterUnion() {
174+
if (isCoroutine) {
175+
return getIndexAfterYields();
176+
} else {
177+
return getIndexAfterDirectReturns();
178+
}
179+
}
180+
unsigned getFirstArgumentIndex() { return getIndexAfterUnion(); }
141181
unsigned getIndexAfterArguments() {
142182
return getFirstArgumentIndex() + getArgumentCount();
143183
}
@@ -176,6 +216,7 @@ namespace irgen {
176216
ElementLayout getResumeParentExecutorLayout() {
177217
return getElement(getResumeParentExecutorIndex());
178218
}
219+
ElementLayout getFlagsLayout() { return getElement(getFlagsIndex()); }
179220
bool canHaveError() { return canHaveValidError; }
180221
ElementLayout getErrorLayout() { return getElement(getErrorIndex()); }
181222
unsigned getErrorCount() { return (unsigned)FixedCount::Error; }
@@ -227,10 +268,18 @@ namespace irgen {
227268
return getElement(getSelfWitnessTableIndex());
228269
}
229270

230-
unsigned getDirectReturnCount() { return directReturnInfos.size(); }
271+
unsigned getDirectReturnCount() {
272+
assert(!isCoroutine);
273+
return directReturnInfos.size();
274+
}
231275
ElementLayout getDirectReturnLayout(unsigned index) {
276+
assert(!isCoroutine);
232277
return getElement(getFirstDirectReturnIndex() + index);
233278
}
279+
unsigned getYieldCount() {
280+
assert(isCoroutine);
281+
return yieldInfos.size();
282+
}
234283

235284
AsyncContextLayout(
236285
IRGenModule &IGM, LayoutStrategy strategy, ArrayRef<SILType> fieldTypes,
@@ -239,8 +288,9 @@ namespace irgen {
239288
SubstitutionMap substitutionMap, NecessaryBindings &&bindings,
240289
Optional<TrailingWitnessInfo> trailingWitnessInfo, SILType errorType,
241290
bool canHaveValidError, ArrayRef<ArgumentInfo> argumentInfos,
242-
ArrayRef<SILResultInfo> directReturnInfos,
291+
bool isCoroutine, ArrayRef<SILYieldInfo> yieldInfos,
243292
ArrayRef<SILResultInfo> indirectReturnInfos,
293+
ArrayRef<SILResultInfo> directReturnInfos,
244294
Optional<ArgumentInfo> localContextInfo);
245295
};
246296

0 commit comments

Comments
 (0)