Skip to content

Commit 704b370

Browse files
etcwildeetcwilde
authored
Merge pull request #38604 from etcwilde/ewilde/rewrite-main
Re-working async-main
2 parents 2ccf3a2 + bb80f9f commit 704b370

File tree

16 files changed

+396
-93
lines changed

16 files changed

+396
-93
lines changed

‎include/swift/SIL/SILDeclRef.h

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,8 +156,11 @@ struct SILDeclRef {
156156
/// The main entry-point function. This may reference a SourceFile for a
157157
/// top-level main, or a decl for e.g an @main decl.
158158
EntryPoint,
159+
160+
/// The asynchronous main entry-point function.
161+
AsyncEntryPoint,
159162
};
160-
163+
161164
/// The AST node represented by this SILDeclRef.
162165
Loc loc;
163166
/// The Kind of this SILDeclRef.
@@ -231,6 +234,9 @@ struct SILDeclRef {
231234
/// Produces a SILDeclRef for a synthetic main entry-point such as @main.
232235
static SILDeclRef getMainDeclEntryPoint(ValueDecl *decl);
233236

237+
/// Produces a SILDeclRef for the synthesized async main entry-point
238+
static SILDeclRef getAsyncMainDeclEntryPoint(ValueDecl *decl);
239+
234240
/// Produces a SILDeclRef for the entry-point of a main FileUnit.
235241
static SILDeclRef getMainFileEntryPoint(FileUnit *file);
236242

‎lib/IRGen/GenObjC.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,7 @@ namespace {
474474
case SILDeclRef::Kind::PropertyWrapperBackingInitializer:
475475
case SILDeclRef::Kind::PropertyWrapperInitFromProjectedValue:
476476
case SILDeclRef::Kind::EntryPoint:
477+
case SILDeclRef::Kind::AsyncEntryPoint:
477478
llvm_unreachable("Method does not have a selector");
478479

479480
case SILDeclRef::Kind::Destroyer:

‎lib/SIL/IR/SILDeclRef.cpp

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,8 @@ SILLinkage SILDeclRef::getLinkage(ForDefinition_t forDefinition) const {
273273
// The main entry-point is public.
274274
if (kind == Kind::EntryPoint)
275275
return SILLinkage::Public;
276+
if (kind == Kind::AsyncEntryPoint)
277+
return SILLinkage::Hidden;
276278

277279
// Add External to the linkage (e.g. Public -> PublicExternal) if this is a
278280
// declaration not a definition.
@@ -448,6 +450,15 @@ SILDeclRef SILDeclRef::getMainDeclEntryPoint(ValueDecl *decl) {
448450
return result;
449451
}
450452

453+
SILDeclRef SILDeclRef::getAsyncMainDeclEntryPoint(ValueDecl *decl) {
454+
auto *file = cast<FileUnit>(decl->getDeclContext()->getModuleScopeContext());
455+
assert(file->getMainDecl() == decl);
456+
SILDeclRef result;
457+
result.loc = decl;
458+
result.kind = Kind::AsyncEntryPoint;
459+
return result;
460+
}
461+
451462
SILDeclRef SILDeclRef::getMainFileEntryPoint(FileUnit *file) {
452463
assert(file->hasEntryPoint() && !file->getMainDecl());
453464
SILDeclRef result;
@@ -539,7 +550,7 @@ IsSerialized_t SILDeclRef::isSerialized() const {
539550
return IsNotSerialized;
540551
}
541552

542-
if (kind == Kind::EntryPoint)
553+
if (kind == Kind::EntryPoint || kind == Kind::AsyncEntryPoint)
543554
return IsNotSerialized;
544555

545556
if (isIVarInitializerOrDestroyer())
@@ -925,6 +936,9 @@ std::string SILDeclRef::mangle(ManglingKind MKind) const {
925936
return mangler.mangleInitFromProjectedValueEntity(cast<VarDecl>(getDecl()),
926937
SKind);
927938

939+
case SILDeclRef::Kind::AsyncEntryPoint: {
940+
return "async_Main";
941+
}
928942
case SILDeclRef::Kind::EntryPoint: {
929943
return getASTContext().getEntryPointFunctionName();
930944
}
@@ -1257,7 +1271,8 @@ unsigned SILDeclRef::getParameterListCount() const {
12571271
return 1;
12581272

12591273
// Always uncurried even if the underlying function is curried.
1260-
if (kind == Kind::DefaultArgGenerator || kind == Kind::EntryPoint)
1274+
if (kind == Kind::DefaultArgGenerator || kind == Kind::EntryPoint ||
1275+
kind == Kind::AsyncEntryPoint)
12611276
return 1;
12621277

12631278
auto *vd = getDecl();

‎lib/SIL/IR/SILFunctionType.cpp

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2526,6 +2526,9 @@ static CanSILFunctionType getNativeSILFunctionType(
25262526
case SILDeclRef::Kind::Deallocator:
25272527
return getSILFunctionTypeForConventions(DeallocatorConventions());
25282528

2529+
case SILDeclRef::Kind::AsyncEntryPoint:
2530+
return getSILFunctionTypeForConventions(
2531+
DefaultConventions(NormalParameterConvention::Guaranteed));
25292532
case SILDeclRef::Kind::EntryPoint:
25302533
llvm_unreachable("Handled by getSILFunctionTypeForAbstractCFunction");
25312534
}
@@ -3071,6 +3074,7 @@ static ObjCSelectorFamily getObjCSelectorFamily(SILDeclRef c) {
30713074
case SILDeclRef::Kind::PropertyWrapperBackingInitializer:
30723075
case SILDeclRef::Kind::PropertyWrapperInitFromProjectedValue:
30733076
case SILDeclRef::Kind::EntryPoint:
3077+
case SILDeclRef::Kind::AsyncEntryPoint:
30743078
llvm_unreachable("Unexpected Kind of foreign SILDeclRef");
30753079
}
30763080

@@ -3343,6 +3347,8 @@ TypeConverter::getDeclRefRepresentation(SILDeclRef c) {
33433347
case SILDeclRef::Kind::IVarDestroyer:
33443348
return SILFunctionTypeRepresentation::Method;
33453349

3350+
case SILDeclRef::Kind::AsyncEntryPoint:
3351+
return SILFunctionTypeRepresentation::Thin;
33463352
case SILDeclRef::Kind::EntryPoint:
33473353
return SILFunctionTypeRepresentation::CFunctionPointer;
33483354
}
@@ -4170,6 +4176,7 @@ static AbstractFunctionDecl *getBridgedFunction(SILDeclRef declRef) {
41704176
case SILDeclRef::Kind::IVarInitializer:
41714177
case SILDeclRef::Kind::IVarDestroyer:
41724178
case SILDeclRef::Kind::EntryPoint:
4179+
case SILDeclRef::Kind::AsyncEntryPoint:
41734180
return nullptr;
41744181
}
41754182
llvm_unreachable("bad SILDeclRef kind");

‎lib/SIL/IR/SILLocation.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ SourceLoc SILLocation::getSourceLoc() const {
4242

4343
// Don't crash if the location is a FilenameAndLocation.
4444
// TODO: this is a workaround until rdar://problem/25225083 is implemented.
45-
if (isFilenameAndLocation())
45+
if (getStorageKind() == FilenameAndLocationKind)
4646
return SourceLoc();
4747

4848
return getSourceLoc(getPrimaryASTNode());

‎lib/SIL/IR/SILPrinter.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,7 @@ void SILDeclRef::print(raw_ostream &OS) const {
316316
switch (kind) {
317317
case SILDeclRef::Kind::Func:
318318
case SILDeclRef::Kind::EntryPoint:
319+
case SILDeclRef::Kind::AsyncEntryPoint:
319320
break;
320321
case SILDeclRef::Kind::Allocator:
321322
OS << "!allocator";

‎lib/SIL/IR/TypeLowering.cpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2556,6 +2556,32 @@ getFunctionInterfaceTypeWithCaptures(TypeConverter &TC,
25562556
innerExtInfo);
25572557
}
25582558

2559+
static CanAnyFunctionType getAsyncEntryPoint(ASTContext &C) {
2560+
2561+
// @main struct Main {
2562+
// static func main() async throws {}
2563+
// static func $main() async throws { try await main() }
2564+
// }
2565+
//
2566+
// func @async_main() async -> Void {
2567+
// do {
2568+
// try await Main.$main()
2569+
// exit(0)
2570+
// } catch {
2571+
// _emitErrorInMain(error)
2572+
// }
2573+
// }
2574+
//
2575+
// This generates the type signature for @async_main
2576+
// TODO: 'Never' return type would be more accurate.
2577+
2578+
CanType returnType = C.getVoidType()->getCanonicalType();
2579+
FunctionType::ExtInfo extInfo =
2580+
FunctionType::ExtInfoBuilder().withAsync(true).withThrows(false).build();
2581+
return CanAnyFunctionType::get(/*genericSig*/ nullptr, {}, returnType,
2582+
extInfo);
2583+
}
2584+
25592585
static CanAnyFunctionType getEntryPointInterfaceType(ASTContext &C) {
25602586
// Use standard library types if we have them; otherwise, fall back to
25612587
// builtins.
@@ -2675,6 +2701,8 @@ CanAnyFunctionType TypeConverter::makeConstantInterfaceType(SILDeclRef c) {
26752701
case SILDeclRef::Kind::IVarDestroyer:
26762702
return getIVarInitDestroyerInterfaceType(cast<ClassDecl>(vd),
26772703
c.isForeign, true);
2704+
case SILDeclRef::Kind::AsyncEntryPoint:
2705+
return getAsyncEntryPoint(Context);
26782706
case SILDeclRef::Kind::EntryPoint:
26792707
return getEntryPointInterfaceType(Context);
26802708
}
@@ -2729,6 +2757,7 @@ TypeConverter::getConstantGenericSignature(SILDeclRef c) {
27292757
case SILDeclRef::Kind::StoredPropertyInitializer:
27302758
return vd->getDeclContext()->getGenericSignatureOfContext();
27312759
case SILDeclRef::Kind::EntryPoint:
2760+
case SILDeclRef::Kind::AsyncEntryPoint:
27322761
llvm_unreachable("Doesn't have generic signature");
27332762
}
27342763

‎lib/SILGen/SILGen.cpp

Lines changed: 78 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -338,24 +338,15 @@ SILGenModule::getConformanceToBridgedStoredNSError(SILLocation loc, Type type) {
338338
return SwiftModule->lookupConformance(type, proto);
339339
}
340340

341-
static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
342-
Optional<FuncDecl*> &cache,
343-
StringRef name) {
341+
static FuncDecl *lookupIntrinsic(ModuleDecl &module,
342+
Optional<FuncDecl *> &cache, Identifier name) {
344343
if (cache)
345344
return *cache;
346-
347-
auto *module = C.getLoadedModule(C.Id_Concurrency);
348-
if (!module) {
349-
cache = nullptr;
350-
return nullptr;
351-
}
352-
353-
SmallVector<ValueDecl *, 1> decls;
354-
module->lookupQualified(module,
355-
DeclNameRef(C.getIdentifier(name)),
356-
NL_QualifiedDefault | NL_IncludeUsableFromInline,
357-
decls);
358345

346+
SmallVector<ValueDecl *, 1> decls;
347+
module.lookupQualified(&module, DeclNameRef(name),
348+
NL_QualifiedDefault | NL_IncludeUsableFromInline,
349+
decls);
359350
if (decls.size() != 1) {
360351
cache = nullptr;
361352
return nullptr;
@@ -365,6 +356,18 @@ static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
365356
return func;
366357
}
367358

359+
static FuncDecl *lookupConcurrencyIntrinsic(ASTContext &C,
360+
Optional<FuncDecl *> &cache,
361+
StringRef name) {
362+
auto *module = C.getLoadedModule(C.Id_Concurrency);
363+
if (!module) {
364+
cache = nullptr;
365+
return nullptr;
366+
}
367+
368+
return lookupIntrinsic(*module, cache, C.getIdentifier(name));
369+
}
370+
368371
FuncDecl *
369372
SILGenModule::getAsyncLetStart() {
370373
return lookupConcurrencyIntrinsic(getASTContext(),
@@ -437,6 +440,37 @@ SILGenModule::getCheckExpectedExecutor() {
437440
"_checkExpectedExecutor");
438441
}
439442

443+
FuncDecl *SILGenModule::getAsyncMainDrainQueue() {
444+
return lookupConcurrencyIntrinsic(getASTContext(), AsyncMainDrainQueue,
445+
"_asyncMainDrainQueue");
446+
}
447+
448+
FuncDecl *SILGenModule::getGetMainExecutor() {
449+
return lookupConcurrencyIntrinsic(getASTContext(), GetMainExecutor,
450+
"_getMainExecutor");
451+
}
452+
453+
FuncDecl *SILGenModule::getSwiftJobRun() {
454+
return lookupConcurrencyIntrinsic(getASTContext(), SwiftJobRun,
455+
"_swiftJobRun");
456+
}
457+
458+
FuncDecl *SILGenModule::getExit() {
459+
if (ExitFunc)
460+
return *ExitFunc;
461+
462+
ASTContext &C = getASTContext();
463+
ModuleDecl *concurrencyShims =
464+
C.getModuleByIdentifier(C.getIdentifier("_SwiftConcurrencyShims"));
465+
466+
if (!concurrencyShims) {
467+
ExitFunc = nullptr;
468+
return nullptr;
469+
}
470+
471+
return lookupIntrinsic(*concurrencyShims, ExitFunc, C.getIdentifier("exit"));
472+
}
473+
440474
ProtocolConformance *SILGenModule::getNSErrorConformanceToError() {
441475
if (NSErrorConformanceToError)
442476
return *NSErrorConformanceToError;
@@ -1036,6 +1070,7 @@ void SILGenModule::emitFunctionDefinition(SILDeclRef constant, SILFunction *f) {
10361070
postEmitFunction(constant, f);
10371071
return;
10381072
}
1073+
case SILDeclRef::Kind::AsyncEntryPoint:
10391074
case SILDeclRef::Kind::EntryPoint: {
10401075
f->setBare(IsBare);
10411076

@@ -1046,7 +1081,27 @@ void SILGenModule::emitFunctionDefinition(SILDeclRef constant, SILFunction *f) {
10461081
auto *decl = constant.getDecl();
10471082
auto *dc = decl->getDeclContext();
10481083
PrettyStackTraceSILFunction X("silgen emitArtificialTopLevel", f);
1049-
SILGenFunction(*this, *f, dc).emitArtificialTopLevel(decl);
1084+
// In all cases, a constant.kind == EntryPoint indicates the main entrypoint
1085+
// to the program, @main.
1086+
// In the synchronous case, the decl is not async, so emitArtificialTopLevel
1087+
// emits the error unwrapping and call to MainType.$main() into @main.
1088+
//
1089+
// In the async case, emitAsyncMainThreadStart is responsible for generating
1090+
// the contents of @main. This wraps @async_main in a task, passes that task
1091+
// to swift_job_run to execute the first thunk, and starts the runloop to
1092+
// run any additional continuations. The kind is EntryPoint, and the decl is
1093+
// async.
1094+
// When the kind is 'AsyncMain', we are generating @async_main. In this
1095+
// case, emitArtificialTopLevel emits the code for calling MaintType.$main,
1096+
// unwrapping errors, and calling exit(0) into @async_main to run the
1097+
// user-specified main function.
1098+
if (constant.kind == SILDeclRef::Kind::EntryPoint && isa<FuncDecl>(decl) &&
1099+
static_cast<FuncDecl *>(decl)->hasAsync()) {
1100+
SILDeclRef mainEntryPoint = SILDeclRef::getAsyncMainDeclEntryPoint(decl);
1101+
SILGenFunction(*this, *f, dc).emitAsyncMainThreadStart(mainEntryPoint);
1102+
} else {
1103+
SILGenFunction(*this, *f, dc).emitArtificialTopLevel(decl);
1104+
}
10501105
postEmitFunction(constant, f);
10511106
return;
10521107
}
@@ -2029,10 +2084,15 @@ class SILGenModuleRAII {
20292084

20302085
// If the source file contains an artificial main, emit the implicit
20312086
// top-level code.
2032-
if (auto *mainDecl = sf->getMainDecl())
2087+
if (auto *mainDecl = sf->getMainDecl()) {
2088+
if (isa<FuncDecl>(mainDecl) &&
2089+
static_cast<FuncDecl *>(mainDecl)->hasAsync())
2090+
emitSILFunctionDefinition(
2091+
SILDeclRef::getAsyncMainDeclEntryPoint(mainDecl));
20332092
emitSILFunctionDefinition(SILDeclRef::getMainDeclEntryPoint(mainDecl));
2093+
}
20342094
}
2035-
2095+
20362096
void emitSILFunctionDefinition(SILDeclRef ref) {
20372097
SGM.emitFunctionDefinition(ref, SGM.getFunction(ref, ForDefinition));
20382098
}

‎lib/SILGen/SILGen.h

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,11 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
131131
Optional<FuncDecl*> ResumeUnsafeThrowingContinuationWithError;
132132
Optional<FuncDecl*> CheckExpectedExecutor;
133133

134+
Optional<FuncDecl *> AsyncMainDrainQueue;
135+
Optional<FuncDecl *> GetMainExecutor;
136+
Optional<FuncDecl *> SwiftJobRun;
137+
Optional<FuncDecl *> ExitFunc;
138+
134139
public:
135140
SILGenModule(SILModule &M, ModuleDecl *SM);
136141

@@ -517,6 +522,15 @@ class LLVM_LIBRARY_VISIBILITY SILGenModule : public ASTVisitor<SILGenModule> {
517522
/// Retrieve the _Concurrency._checkExpectedExecutor intrinsic.
518523
FuncDecl *getCheckExpectedExecutor();
519524

525+
/// Retrieve the _Concurrency._asyncMainDrainQueue intrinsic.
526+
FuncDecl *getAsyncMainDrainQueue();
527+
/// Retrieve the _Concurrency._getMainExecutor intrinsic.
528+
FuncDecl *getGetMainExecutor();
529+
/// Retrieve the _Concurrency._swiftJobRun intrinsic.
530+
FuncDecl *getSwiftJobRun();
531+
// Retrieve the _SwiftConcurrencyShims.exit intrinsic.
532+
FuncDecl *getExit();
533+
520534
SILFunction *getKeyPathProjectionCoroutine(bool isReadAccess,
521535
KeyPathTypeKind typeKind);
522536

‎lib/SILGen/SILGenBuiltin.cpp

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1402,9 +1402,10 @@ emitFunctionArgumentForAsyncTaskEntryPoint(SILGenFunction &SGF,
14021402
}
14031403

14041404
// Emit SIL for the named builtin: createAsyncTask.
1405-
static ManagedValue emitBuiltinCreateAsyncTask(
1406-
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
1407-
ArrayRef<ManagedValue> args, SGFContext C) {
1405+
ManagedValue emitBuiltinCreateAsyncTask(SILGenFunction &SGF, SILLocation loc,
1406+
SubstitutionMap subs,
1407+
ArrayRef<ManagedValue> args,
1408+
SGFContext C) {
14081409
ASTContext &ctx = SGF.getASTContext();
14091410
auto flags = args[0].forward(SGF);
14101411

0 commit comments

Comments
 (0)