Skip to content

Commit 5c5f012

Browse files
slavapestovslavapestov
committed
AST: Add withUnsafe[Throwing]Continuation built-ins
These lower as follows, where %fn is the passed-in function value: %c = get_async_continuation_addr ... apply %fn(%c) await_async_continuation %c
1 parent 361eef2 commit 5c5f012

File tree

4 files changed

+202
-0
lines changed

4 files changed

+202
-0
lines changed

‎include/swift/AST/Builtins.def

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -504,6 +504,20 @@ BUILTIN_SIL_OPERATION(DifferentiableFunction, "differentiableFunction", Special)
504504
/// linearFunction
505505
BUILTIN_SIL_OPERATION(LinearFunction, "linearFunction", Special)
506506

507+
/// withUnsafeContinuation<T> : (Builtin.RawUnsafeContinuation -> ()) async -> T
508+
///
509+
/// Unsafely capture the current continuation and pass it to the given
510+
/// function value. Returns a value of type T when the continuation is
511+
/// resumed.
512+
BUILTIN_SIL_OPERATION(WithUnsafeContinuation, "withUnsafeContinuation", Special)
513+
514+
/// withUnsafeThrowingContinuation<T> : (Builtin.RawUnsafeContinuation -> ()) async throws -> T
515+
///
516+
/// Unsafely capture the current continuation and pass it to the given
517+
/// function value. Returns a value of type T or throws an error when
518+
/// the continuation is resumed.
519+
BUILTIN_SIL_OPERATION(WithUnsafeThrowingContinuation, "withUnsafeThrowingContinuation", Special)
520+
507521
#undef BUILTIN_SIL_OPERATION
508522

509523
// BUILTIN_RUNTIME_CALL - A call into a runtime function.

‎lib/AST/Builtins.cpp

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1655,6 +1655,29 @@ static ValueDecl *getPolymorphicBinaryOperation(ASTContext &ctx,
16551655
return builder.build(id);
16561656
}
16571657

1658+
static ValueDecl *getWithUnsafeContinuation(ASTContext &ctx,
1659+
Identifier id,
1660+
bool throws) {
1661+
BuiltinFunctionBuilder builder(ctx);
1662+
1663+
auto contTy = ctx.TheRawUnsafeContinuationType;
1664+
SmallVector<AnyFunctionType::Param, 1> params;
1665+
params.emplace_back(contTy);
1666+
1667+
auto voidTy = ctx.TheEmptyTupleType;
1668+
auto extInfo = FunctionType::ExtInfoBuilder().withNoEscape().build();
1669+
auto *fnTy = FunctionType::get(params, voidTy, extInfo);
1670+
1671+
builder.addParameter(makeConcrete(fnTy));
1672+
builder.setResult(makeGenericParam());
1673+
1674+
builder.setAsync();
1675+
if (throws)
1676+
builder.setThrows();
1677+
1678+
return builder.build(id);
1679+
}
1680+
16581681
/// An array of the overloaded builtin kinds.
16591682
static const OverloadedBuiltinKind OverloadedBuiltinKinds[] = {
16601683
OverloadedBuiltinKind::None,
@@ -2593,6 +2616,12 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
25932616
case BuiltinValueKind::TriggerFallbackDiagnostic:
25942617
return getTriggerFallbackDiagnosticOperation(Context, Id);
25952618

2619+
case BuiltinValueKind::WithUnsafeContinuation:
2620+
return getWithUnsafeContinuation(Context, Id, /*throws=*/false);
2621+
2622+
case BuiltinValueKind::WithUnsafeThrowingContinuation:
2623+
return getWithUnsafeContinuation(Context, Id, /*throws=*/true);
2624+
25962625
case BuiltinValueKind::AutoDiffCreateLinearMapContext:
25972626
return getAutoDiffCreateLinearMapContext(Context, Id);
25982627

‎lib/SILGen/SILGenBuiltin.cpp

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1462,6 +1462,85 @@ static ManagedValue emitBuiltinCreateAsyncTaskFuture(
14621462
return SGF.emitManagedRValueWithCleanup(apply);
14631463
}
14641464

1465+
// Shared implementation of withUnsafeContinuation and
1466+
// withUnsafe[Throwing]Continuation.
1467+
static ManagedValue emitBuiltinWithUnsafeContinuation(
1468+
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
1469+
ArrayRef<ManagedValue> args, SGFContext C, bool throws) {
1470+
// Allocate space to receive the resume value when the continuation is
1471+
// resumed.
1472+
auto substResultType = subs.getReplacementTypes()[0]->getCanonicalType();
1473+
auto opaqueResumeType = SGF.getLoweredType(AbstractionPattern::getOpaque(),
1474+
substResultType);
1475+
auto resumeBuf = SGF.emitTemporaryAllocation(loc, opaqueResumeType);
1476+
1477+
// Capture the current continuation.
1478+
auto continuation = SGF.B.createGetAsyncContinuationAddr(loc, resumeBuf,
1479+
substResultType,
1480+
throws);
1481+
1482+
// Get the callee value.
1483+
auto substFnType = args[0].getType().castTo<SILFunctionType>();
1484+
SILValue fnValue = (substFnType->isCalleeConsumed()
1485+
? args[0].forward(SGF)
1486+
: args[0].getValue());
1487+
1488+
// Call the provided function value.
1489+
SGF.B.createApply(loc, fnValue, {}, {continuation});
1490+
1491+
// Await the continuation.
1492+
SILBasicBlock *resumeBlock = SGF.createBasicBlock();
1493+
SILBasicBlock *errorBlock = nullptr;
1494+
1495+
if (throws)
1496+
errorBlock = SGF.createBasicBlock(FunctionSection::Postmatter);
1497+
1498+
SGF.B.createAwaitAsyncContinuation(loc, continuation, resumeBlock, errorBlock);
1499+
1500+
// Propagate an error if we have one.
1501+
if (throws) {
1502+
SGF.B.emitBlock(errorBlock);
1503+
1504+
Scope errorScope(SGF, loc);
1505+
1506+
auto errorTy = SGF.getASTContext().getErrorDecl()->getDeclaredType()
1507+
->getCanonicalType();
1508+
auto errorVal
1509+
= SGF.B.createOwnedPhiArgument(SILType::getPrimitiveObjectType(errorTy));
1510+
1511+
SGF.emitThrow(loc, errorVal, true);
1512+
}
1513+
1514+
SGF.B.emitBlock(resumeBlock);
1515+
1516+
// The incoming value is the maximally-abstracted result type of the
1517+
// continuation. Move it out of the resume buffer and reabstract it if
1518+
// necessary.
1519+
auto resumeResult = SGF.emitLoad(loc, resumeBuf,
1520+
AbstractionPattern::getOpaque(),
1521+
substResultType,
1522+
SGF.getTypeLowering(substResultType),
1523+
SGFContext(), IsTake);
1524+
1525+
return resumeResult;
1526+
}
1527+
1528+
// Emit SIL for the named builtin: withUnsafeContinuation
1529+
static ManagedValue emitBuiltinWithUnsafeContinuation(
1530+
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
1531+
ArrayRef<ManagedValue> args, SGFContext C) {
1532+
return emitBuiltinWithUnsafeContinuation(SGF, loc, subs, args, C,
1533+
/*throws=*/false);
1534+
}
1535+
1536+
// Emit SIL for the named builtin: withUnsafeThrowingContinuation
1537+
static ManagedValue emitBuiltinWithUnsafeThrowingContinuation(
1538+
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
1539+
ArrayRef<ManagedValue> args, SGFContext C) {
1540+
return emitBuiltinWithUnsafeContinuation(SGF, loc, subs, args, C,
1541+
/*throws=*/true);
1542+
}
1543+
14651544
static ManagedValue emitBuiltinAutoDiffCreateLinearMapContext(
14661545
SILGenFunction &SGF, SILLocation loc, SubstitutionMap subs,
14671546
ArrayRef<ManagedValue> args, SGFContext C) {

‎test/SILGen/async_builtins.swift

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,3 +38,83 @@ public struct X {
3838
_ = Builtin.createAsyncTaskFuture(0, nil, closure)
3939
}
4040
}
41+
42+
// CHECK-LABEL: sil [ossa] @$s4test26usesWithUnsafeContinuationyyYF : $@convention(thin) @async () -> () {
43+
public func usesWithUnsafeContinuation() async {
44+
// trivial resume type
45+
let _: Int = await Builtin.withUnsafeContinuation { c in }
46+
47+
// CHECK: [[FN:%.*]] = function_ref @$s4test26usesWithUnsafeContinuationyyYFyBcXEfU_ : $@convention(thin) (Builtin.RawUnsafeContinuation) -> ()
48+
// CHECK: [[TMP:%.*]] = convert_function [[FN]] : $@convention(thin) (Builtin.RawUnsafeContinuation) -> () to $@convention(thin) @noescape (Builtin.RawUnsafeContinuation) -> ()
49+
// CHECK: [[CLOSURE:%.*]] = thin_to_thick_function [[TMP]]
50+
// CHECK: [[BOX:%.*]] = alloc_stack $Int
51+
// CHECK: [[CC:%.*]] = get_async_continuation_addr Int, [[BOX]] : $*Int
52+
// CHECK: apply [[CLOSURE]]([[CC]]) : $@noescape @callee_guaranteed (Builtin.RawUnsafeContinuation) -> ()
53+
// CHECK: await_async_continuation [[CC]] : $Builtin.RawUnsafeContinuation, resume bb1
54+
55+
// CHECK: bb1:
56+
// CHECK: [[RESULT:%.*]] = load [trivial] [[BOX]] : $*Int
57+
// CHECK: dealloc_stack [[BOX]]
58+
59+
// loadable resume type
60+
let _: String = await Builtin.withUnsafeContinuation { c in }
61+
62+
// CHECK: [[FN:%.*]] = function_ref @$s4test26usesWithUnsafeContinuationyyYFyBcXEfU0_ : $@convention(thin) (Builtin.RawUnsafeContinuation) -> ()
63+
// CHECK: [[TMP:%.*]] = convert_function [[FN]] : $@convention(thin) (Builtin.RawUnsafeContinuation) -> () to $@convention(thin) @noescape (Builtin.RawUnsafeContinuation) -> ()
64+
// CHECK: [[CLOSURE:%.*]] = thin_to_thick_function [[TMP]]
65+
// CHECK: [[BOX:%.*]] = alloc_stack $String
66+
// CHECK: [[CC:%.*]] = get_async_continuation_addr String, [[BOX]] : $*String
67+
// CHECK: apply [[CLOSURE]]([[CC]]) : $@noescape @callee_guaranteed (Builtin.RawUnsafeContinuation) -> ()
68+
// CHECK: await_async_continuation [[CC]] : $Builtin.RawUnsafeContinuation, resume bb2
69+
70+
// CHECK: bb2:
71+
// CHECK: [[RESULT:%.*]] = load [take] [[BOX]] : $*String
72+
// CHECK: destroy_value [[RESULT]]
73+
// CHECK: dealloc_stack [[BOX]]
74+
75+
// address-only resume type
76+
let _: Any = await Builtin.withUnsafeContinuation { c in }
77+
78+
// CHECK: [[FN:%.*]] = function_ref @$s4test26usesWithUnsafeContinuationyyYFyBcXEfU1_ : $@convention(thin) (Builtin.RawUnsafeContinuation) -> ()
79+
// CHECK: [[TMP:%.*]] = convert_function [[FN]] : $@convention(thin) (Builtin.RawUnsafeContinuation) -> () to $@convention(thin) @noescape (Builtin.RawUnsafeContinuation) -> ()
80+
// CHECK: [[CLOSURE:%.*]] = thin_to_thick_function [[TMP]]
81+
// CHECK: [[BOX:%.*]] = alloc_stack $Any
82+
// CHECK: [[CC:%.*]] = get_async_continuation_addr Any, [[BOX]] : $*Any
83+
// CHECK: apply [[CLOSURE]]([[CC]]) : $@noescape @callee_guaranteed (Builtin.RawUnsafeContinuation) -> ()
84+
// CHECK: await_async_continuation [[CC]] : $Builtin.RawUnsafeContinuation, resume bb3
85+
86+
// CHECK: bb3:
87+
// CHECK: [[COPY:%.*]] = alloc_stack $Any
88+
// CHECK: copy_addr [take] [[BOX]] to [initialization] [[COPY]]
89+
// CHECK: destroy_addr [[COPY]]
90+
// CHECK: dealloc_stack [[COPY]]
91+
// CHECK: dealloc_stack [[BOX]]
92+
}
93+
94+
// CHECK-LABEL: sil [ossa] @$s4test34usesWithUnsafeThrowingContinuationyyYKF : $@convention(thin) @async () -> @error Error {
95+
public func usesWithUnsafeThrowingContinuation() async throws {
96+
let _: Int = try await Builtin.withUnsafeThrowingContinuation { c in }
97+
98+
// CHECK: [[FN:%.*]] = function_ref @$s4test34usesWithUnsafeThrowingContinuationyyYKFyBcXEfU_ : $@convention(thin) (Builtin.RawUnsafeContinuation) -> ()
99+
// CHECK: [[TMP:%.*]] = convert_function [[FN]] : $@convention(thin) (Builtin.RawUnsafeContinuation) -> () to $@convention(thin) @noescape (Builtin.RawUnsafeContinuation) -> ()
100+
// CHECK: [[CLOSURE:%.*]] = thin_to_thick_function [[TMP]]
101+
// CHECK: [[BOX:%.*]] = alloc_stack $Int
102+
// CHECK: [[CC:%.*]] = get_async_continuation_addr [throws] Int, [[BOX]] : $*Int
103+
// CHECK: apply [[CLOSURE]]([[CC]]) : $@noescape @callee_guaranteed (Builtin.RawUnsafeContinuation) -> ()
104+
// CHECK: await_async_continuation [[CC]] : $Builtin.RawUnsafeContinuation, resume bb1, error bb2
105+
106+
// CHECK: bb1:
107+
// CHECK: [[RESULT:%.*]] = load [trivial] [[BOX]] : $*Int
108+
// CHECK: dealloc_stack [[BOX]]
109+
110+
// CHECK: bb2([[ERROR:%.*]] : @owned $Error):
111+
// CHECK: builtin "willThrow"([[ERROR]] : $Error) : $()
112+
// CHECK: dealloc_stack [[BOX]]
113+
// CHECK: throw [[ERROR]]
114+
}
115+
116+
// Make sure we do the right thing when the closure value is non-trivial,
117+
// because it has captures and was formed by a partial_apply.
118+
public func usesWithUnsafeContinuationCaptures(fn: (Builtin.RawUnsafeContinuation) -> ()) async throws {
119+
let _: Int = await Builtin.withUnsafeContinuation { c in fn(c) }
120+
}

0 commit comments

Comments
 (0)