lib/mprompt/src/api.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 //! A typed Zig layer over the package's stack-switching runtime runs a function on a stack of its
2 //! own and returns the function's result as a Zig value. A body run this way can pause with a value
3 //! of a type it names and continue with a value of another type it names.
4 //!
5 //! The runtime underneath passes one untyped pointer each way when it switches stacks, and Zig
6 //! callers want their own types, error unions included, checked by the compiler. A body pauses in
7 //! one of two styles: a handler answers the pause at once on the caller's side, as a generator
8 //! hands out each value, or the pause goes back to the caller, which continues it later, as a
9 //! scheduler does with waiting workers.
10 //!
11 //! A value that crosses from one stack to the other has to stay in memory that both sides can reach
12 //! until the receiving side has read it. When a pause goes back to the caller, the caller gets one
13 //! pointer back and has to tell from it whether the body finished or paused. A handler that returns
14 //! without continuing the body leaves the call that started the body with no result to return.
15 //!
16 //! Each value travels by address: the sender stores it in a small record on its own stack and
17 //! passes the record's address, and the record stays valid because the sender waits inside its
18 //! resume or pause call until the receiver has read it. `SuspendedRun` keeps the body's result and
19 //! a one-byte marker (*done marker*) in the caller's own memory, and the body's wrapper returns the
20 //! marker's address when the body finishes, so that address means the run finished and any other
21 //! pointer is a handle to the paused body. `yieldWith` hands its handler a typed handle to the
22 //! paused rest of the body, a `Continuation`. The program stops with a panic when the handler
23 //! returns without continuing the body. Every context passes by pointer, and the compiler rejects
24 //! any other context type. The stack switch underneath exists for x86_64 and aarch64 targets other
25 //! than Windows, on 64-bit targets only.
26 //!
27 //! - *one-shot handle*: a handle resumable at most once, continuing in place with no copy
28 //! - *multi-shot handle*: a reference-counted handle resumable more than once
29 //! - *one-shot continuation*: a continuation resumable at most once, continuing in place with no
30 //! copy
31 //! - *multi-shot continuation*: a reference-counted continuation resumable more than once
32 //! - *tail resume*: a resume made as a function's last act, returning straight to the call that
33 //! entered the prompt
34 const std = @import("std");
35 const raw = @import("root.zig");
36
37 pub const Prompt = raw.Prompt;
38
39 const VoidSlot = struct {};
40
41 fn Slot(comptime T: type) type {
42 return if (T == void) VoidSlot else SlotValueType(T);
43 }
44
45 fn SlotValueType(comptime T: type) type {
46 return struct {
47 value: T = undefined,
48 };
49 }
50
51 /// Returns a tagged union type with two cases: the body returned, or the body paused. A caller
52 /// switches on it after starting or continuing a paused run, to learn whether the body finished or
53 /// paused again. `SuspendedRun.start` and `SuspendedPrompt.continueWith` return it. `ResumeValue`
54 /// is the type a caller passes to continue the body. `Result` is the type the body returns. An
55 /// error union passes through `Result` unchanged, so a failing body comes back as an error inside
56 /// the returned case.
57 pub fn PromptOutcome(comptime ResumeValue: type, comptime Result: type) type {
58 return union(enum) {
59 /// Holds the value the body returned. By the time this case comes back, the runtime has
60 /// released the body's stacklet, so no handle remains to drop.
61 returned: Result,
62 /// Holds the handle to the paused body. The caller either continues the handle with a value
63 /// or drops it, and a one-shot handle is used once.
64 suspended: SuspendedPrompt(ResumeValue, Result),
65 };
66 }
67
68 /// Returns the typed handle type for a body that paused inside a `SuspendedRun`. A caller keeps
69 /// this handle from the moment a run pauses until it continues or abandons the run, for example a
70 /// scheduler that holds one handle per waiting worker. The handle points into the `SuspendedRun`
71 /// that started the body, so that `SuspendedRun` must stay at the same address while the handle is
72 /// in use. A one-shot handle is continued once or dropped once, because continuing it again, or
73 /// after a drop, uses a stacklet the runtime has already freed or reused. `asMulti` turns it into a
74 /// multi-shot handle that the caller can continue once per reference.
75 pub fn SuspendedPrompt(comptime ResumeValue: type, comptime Result: type) type {
76 return struct {
77 /// Holds the runtime's resumption for the paused body. The methods pass it to the runtime's
78 /// `resumePrompt`, `resumeDrop`, `resumeMulti`, `resumeDup` and `resumeResumeCount`.
79 raw_resume: *raw.Resume,
80 /// Points to the done marker inside the `SuspendedRun` that started the body.
81 /// `continueWith` compares the runtime's result with this address to tell a finished body
82 /// from a new pause.
83 done_marker: *u8,
84 /// Points to the storage inside the `SuspendedRun` where the body's wrapper writes the
85 /// body's result. `continueWith` reads the result from it when the body finishes.
86 result: *Slot(Result),
87
88 const Self = @This();
89 /// Names the `PromptOutcome` type for this handle's `ResumeValue` and `Result`, which
90 /// `continueWith` returns.
91 pub const Outcome: type = PromptOutcome(ResumeValue, Result);
92
93 /// Resumes the paused body so that its `suspendPrompt` call returns `value`, for a caller
94 /// continuing the body with the value the pause is waiting for. The call runs the body on
95 /// its stacklet until it finishes or pauses again. The call returns `.returned` with the
96 /// body's result when the body finished, or `.suspended` with a new handle when it paused
97 /// again. The call uses up a one-shot handle, and a later pause comes back as a new handle
98 /// in `.suspended`. The call gives up one reference of a multi-shot handle. When other
99 /// references of a multi-shot handle remain, the call first copies the paused stack to the
100 /// heap, and a later resume copies it back.
101 pub fn continueWith(self: Self, value: ResumeValue) Outcome {
102 var slot: Slot(ResumeValue) = .{};
103 writeSlot(ResumeValue, &slot, value);
104 const raw_result = raw.resumePrompt(self.raw_resume, slotPtr(ResumeValue, &slot));
105 return outcomeFromRaw(ResumeValue, Result, raw_result, self.done_marker, self.result);
106 }
107
108 /// Continues the paused body with no value, for a `ResumeValue` of `void`. Any other
109 /// `ResumeValue` is a compile error.
110 pub fn continueWithoutValue(self: Self) Outcome {
111 if (ResumeValue != void) {
112 @compileError("continueWithoutValue requires ResumeValue to be void");
113 }
114 return self.continueWith({});
115 }
116
117 /// Returns a multi-shot handle for the same pause, for a caller that converts the handle
118 /// before continuing the same pause more than once. The caller uses the returned handle in
119 /// place of the old one. Converting a one-shot handle allocates a record from the process
120 /// allocator, and the record starts with one reference. The program stops with a panic if
121 /// that allocation fails. A handle that is already multi-shot comes back unchanged.
122 pub fn asMulti(self: Self) Self {
123 return .{
124 .raw_resume = raw.resumeMulti(self.raw_resume),
125 .done_marker = self.done_marker,
126 .result = self.result,
127 };
128 }
129
130 /// A caller takes one extra reference for each extra time it will continue the pause. The
131 /// call returns a second handle for the same pause and adds one reference, or returns null
132 /// for a one-shot handle. The call copies no stack memory. The runtime copies the paused
133 /// stack only when a resume finds other references still held.
134 pub fn dup(self: Self) ?Self {
135 const duplicated = raw.resumeDup(self.raw_resume) orelse return null;
136 return .{
137 .raw_resume = duplicated,
138 .done_marker = self.done_marker,
139 .result = self.result,
140 };
141 }
142
143 /// Gives up this handle without continuing the body, for a caller abandoning a paused body
144 /// it will never continue, as the Chic host does when it tears down a turn. For a one-shot
145 /// handle, and for the last reference of a multi-shot handle, the runtime frees the paused
146 /// stacklets, and the rest of the body never runs, so its `defer` statements never run
147 /// either.
148 pub fn drop(self: Self) void {
149 raw.resumeDrop(self.raw_resume);
150 }
151
152 /// Returns the number of resumes of a multi-shot handle, or 0 for a one-shot handle, for a
153 /// caller reading how many times the pause has been continued so far.
154 pub fn resumeCount(self: Self) c_long {
155 return raw.resumeResumeCount(self.raw_resume);
156 }
157 };
158 }
159
160 fn SuspendedRunStartCallbackType(
161 comptime RunState: type,
162 comptime Result: type,
163 comptime body: anytype,
164 ) type {
165 return struct {
166 fn run(prompt: *Prompt, arg: ?*anyopaque) callconv(.c) ?*anyopaque {
167 const run_state: *RunState = @ptrCast(@alignCast(arg.?));
168 writeSlot(Result, &run_state.result, body(prompt, run_state.context));
169 return @ptrCast(&run_state.done_marker);
170 }
171 };
172 }
173
174 /// Returns a struct type that holds a context pointer, storage for the body's result, and a done
175 /// marker, for a caller that starts a body that may pause and hand control back, then continues it
176 /// later from its own loop, as a scheduler does with workers or the Chic host does with an
177 /// interpreter turn. `Context` must be a pointer type, and any other type is a compile error.
178 /// `init(context)` builds the value, and `start(body)` runs `body(prompt, context)` on a new
179 /// stacklet and returns a `PromptOutcome`. The body pauses with `suspendPrompt`, and the pause
180 /// comes back to the caller as `.suspended`. The value must stay at one address from `start` until
181 /// the body finishes or its last handle is dropped, because every handle points to its result
182 /// storage and done marker. A caller may call `start` again after the previous run finished, as the
183 /// scheduler test does with one value per worker.
184 pub fn SuspendedRun(
185 comptime ResumeValue: type,
186 comptime Result: type,
187 comptime Context: type,
188 ) type {
189 requirePointer(Context, "SuspendedRun context");
190 return struct {
191 /// The pointer passed to `init`, which `start` hands to the body.
192 context: Context,
193 /// Storage where the body's wrapper writes the body's result when the body finishes. A
194 /// handle from this run reads the result from it. The default is empty storage, whose value
195 /// stays undefined until the body returns.
196 result: Slot(Result) = .{},
197 /// The done marker: a byte whose address the body's wrapper returns when the body finishes.
198 /// `start` and each handle compare the runtime's result with this address to tell a
199 /// finished body from a pause. Only its address is used, and its value stays at the
200 /// default 0.
201 done_marker: u8 = 0,
202
203 const Self = @This();
204 /// Names the `PromptOutcome` type for this run's `ResumeValue` and `Result`, which `start`
205 /// returns.
206 pub const Outcome: type = PromptOutcome(ResumeValue, Result);
207
208 /// Returns a run holding `context`, with empty result storage and the done marker at 0, for
209 /// a caller that builds the run once before starting the body. The call allocates nothing
210 /// and starts nothing.
211 pub fn init(context: Context) Self {
212 return .{ .context = context };
213 }
214
215 /// Creates a prompt with a new stacklet and runs `body(prompt, context)` on it, for a
216 /// caller that begins the body and learns whether it finished at once or paused. The call
217 /// returns `.returned` with the body's result when the body finishes without pausing, or
218 /// `.suspended` with a handle when it pauses through `suspendPrompt`. `self` must stay at
219 /// the same address while any handle from this run is in use.
220 pub fn start(
221 self: *Self,
222 comptime body: *const fn (*Prompt, Context) Result,
223 ) Outcome {
224 const Callback: type = SuspendedRunStartCallbackType(Self, Result, body);
225 const raw_result = raw.prompt(Callback.run, self);
226 return outcomeFromRaw(
227 ResumeValue,
228 Result,
229 raw_result,
230 &self.done_marker,
231 &self.result,
232 );
233 }
234 };
235 }
236
237 fn suspend_prompt_callback(raw_resume: *raw.Resume, _: ?*anyopaque) callconv(.c) ?*anyopaque {
238 return @ptrCast(raw_resume);
239 }
240
241 /// Suspends the body at `prompt`, so the pending `start` or `continueWith` call returns
242 /// `.suspended` with a handle, for a body started by `SuspendedRun` that pauses and hands control
243 /// back to the code that started or last continued it. The call returns the value the caller later
244 /// passes to `continueWith`. `prompt` must be the prompt the body received from `start`, and the
245 /// body must still be running on it, which assertions check in safe builds. Only a body started by
246 /// `SuspendedRun` can call it, because under `run` the handle would be read as the body's result.
247 pub fn suspendPrompt(comptime ResumeValue: type, prompt: *Prompt) ResumeValue {
248 return readSlot(ResumeValue, raw.yieldPrompt(prompt, suspend_prompt_callback, null));
249 }
250
251 /// Returns the typed handle type that a `yieldWith` handler gets for the paused rest of the body,
252 /// for a handler that uses it to continue the paused body with its answer. `ResumeValue` is the
253 /// type the paused `yieldWith` call returns. `PromptResult` is the type the body's `run` returns.
254 /// `PromptResult` must be the `Result` of the enclosing `run`, and the compiler does not check that
255 /// the two match. The handler must call one of the continue methods before it returns, or the
256 /// program stops with a panic.
257 pub fn Continuation(comptime ResumeValue: type, comptime PromptResult: type) type {
258 return struct {
259 /// Holds the runtime's resumption for the paused body.
260 raw_resume: *raw.Resume,
261 /// Points to a flag in `yieldWith` that a continue method sets once the body hands control
262 /// back. `yieldWith` stops the program with a panic when the handler returns and the flag
263 /// is still false.
264 continued: *bool,
265 /// Points to where a continue method stores the runtime's result pointer. `yieldWith`
266 /// returns that pointer to the call that entered the prompt.
267 raw_result: *?*anyopaque,
268
269 const Self = @This();
270
271 /// Resumes the paused body so that its `yieldWith` call returns `value`, for a handler
272 /// continuing the paused body with its answer and getting the body's result back. The call
273 /// returns the body's result once the body runs to its end, through any later pauses that
274 /// their own handlers continue. The call marks the continuation as used. A one-shot
275 /// continuation is continued once. When other references of a multi-shot continuation
276 /// remain, the call first copies the paused stack to the heap, and a later resume copies it
277 /// back.
278 pub fn continueWith(self: Self, value: ResumeValue) PromptResult {
279 return self.continueInternal(false, value);
280 }
281
282 /// Resumes the paused body with `value` as a tail resume, for a handler whose last act is
283 /// to continue the body, so repeated pauses do not pile up handler frames. For a one-shot
284 /// continuation, control never comes back to the handler, and the body's result goes
285 /// straight to the call that entered the prompt. The call must be the handler's last
286 /// action. For a multi-shot continuation, only the first tail resume works this way, and
287 /// later ones behave as `continueWith` does.
288 pub fn continueTailWith(self: Self, value: ResumeValue) PromptResult {
289 return self.continueInternal(true, value);
290 }
291
292 /// Continues the paused body with no value, for a `ResumeValue` of `void`. Any other
293 /// `ResumeValue` is a compile error.
294 pub fn continueWithoutValue(self: Self) PromptResult {
295 if (ResumeValue != void) {
296 @compileError("continueWithoutValue requires ResumeValue to be void");
297 }
298 return self.continueWith({});
299 }
300
301 /// Continues the paused body with no value as a tail resume, for a `ResumeValue` of `void`.
302 /// Any other `ResumeValue` is a compile error.
303 pub fn continueTailWithoutValue(self: Self) PromptResult {
304 if (ResumeValue != void) {
305 @compileError("continueTailWithoutValue requires ResumeValue to be void");
306 }
307 return self.continueTailWith({});
308 }
309
310 /// Returns a multi-shot continuation for the same pause, for a handler that converts the
311 /// continuation before continuing the same pause more than once, as a search that tries
312 /// each branch does. The handler uses the returned continuation in place of the old one.
313 /// Converting a one-shot continuation allocates a record from the process allocator, and
314 /// the program stops with a panic if that allocation fails. A continuation that is already
315 /// multi-shot comes back unchanged. Each continue call gives up one reference, so
316 /// continuing twice takes a `dup` first.
317 pub fn asMulti(self: Self) Self {
318 return .{
319 .raw_resume = raw.resumeMulti(self.raw_resume),
320 .continued = self.continued,
321 .raw_result = self.raw_result,
322 };
323 }
324
325 /// A handler takes one extra reference for each extra time it will continue the pause. The
326 /// call returns a second continuation for the same pause and adds one reference, or returns
327 /// null for a one-shot continuation. The call copies no stack memory.
328 pub fn dup(self: Self) ?Self {
329 const duplicated = raw.resumeDup(self.raw_resume) orelse return null;
330 return .{
331 .raw_resume = duplicated,
332 .continued = self.continued,
333 .raw_result = self.raw_result,
334 };
335 }
336
337 /// Gives up one reference without continuing the body, for a handler giving back an extra
338 /// reference it took and will not use. The handler must still continue the body through
339 /// another reference before it returns, because a handler that drops its only reference and
340 /// returns stops the program with a panic.
341 pub fn drop(self: Self) void {
342 raw.resumeDrop(self.raw_resume);
343 }
344
345 /// Returns the number of resumes of a multi-shot continuation, or 0 for a one-shot
346 /// continuation, for a handler reading how many times the pause has been continued so far.
347 pub fn resumeCount(self: Self) c_long {
348 return raw.resumeResumeCount(self.raw_resume);
349 }
350
351 fn continueInternal(self: Self, comptime tail: bool, value: ResumeValue) PromptResult {
352 var slot: Slot(ResumeValue) = .{};
353 writeSlot(ResumeValue, &slot, value);
354 const result = if (tail)
355 raw.resumeTailPrompt(self.raw_resume, slotPtr(ResumeValue, &slot))
356 else
357 raw.resumePrompt(self.raw_resume, slotPtr(ResumeValue, &slot));
358 self.continued.* = true;
359 self.raw_result.* = result;
360 return readSlot(PromptResult, result);
361 }
362 };
363 }
364
365 fn RunEnvironmentType(
366 comptime Result: type,
367 comptime Context: type,
368 comptime body: anytype,
369 ) type {
370 return struct {
371 context: Context,
372 result: Slot(Result) = .{},
373
374 fn start(prompt: *Prompt, arg: ?*anyopaque) callconv(.c) ?*anyopaque {
375 const environment: *@This() = @ptrCast(@alignCast(arg.?));
376 writeSlot(Result, &environment.result, body(prompt, environment.context));
377 return slotPtr(Result, &environment.result);
378 }
379 };
380 }
381
382 /// A caller runs the body to its end in one call, and a handler answers each pause before the body
383 /// goes on. The call creates a prompt with a new stacklet, runs `body(prompt, context)` on it, and
384 /// returns what the body returns. The prompt marks where the body entered its stacklet, and the
385 /// body passes it to `yieldWith` to pause back to that point. `context` must be a pointer, and any
386 /// other type is a compile error. An error union passes through: a body that returns an error makes
387 /// `run` return that error. The stacklet comes from the calling thread's cache or from a new
388 /// reservation. The first prompt in a process initializes the runtime with the default
389 /// configuration when `init` has not run. The runtime releases the stacklet when the body returns.
390 /// The body pauses with `yieldWith`. `suspendPrompt` is only for bodies started by `SuspendedRun`.
391 pub fn run(
392 comptime Result: type,
393 context: anytype,
394 comptime body: *const fn (*Prompt, @TypeOf(context)) Result,
395 ) Result {
396 const Context = @TypeOf(context);
397 requirePointer(Context, "run context");
398 const Environment: type = RunEnvironmentType(Result, Context, body);
399 var environment: Environment = .{ .context = context };
400 return readSlot(Result, raw.prompt(Environment.start, &environment));
401 }
402
403 fn RunWithoutContextType(comptime Result: type, comptime body: anytype) type {
404 return struct {
405 fn start(prompt: *Prompt, _: *@This()) Result {
406 return body(prompt);
407 }
408 };
409 }
410
411 /// Runs `body(prompt)` the way `run` does, with no context argument, so a caller whose body needs
412 /// no context skips the pointer.
413 pub fn runWithoutContext(
414 comptime Result: type,
415 comptime body: *const fn (*Prompt) Result,
416 ) Result {
417 const Context: type = RunWithoutContextType(Result, body);
418 var context: Context = .{};
419 return run(Result, &context, Context.start);
420 }
421
422 fn YieldWithEnvironmentType(
423 comptime ResumeValue: type,
424 comptime PromptResult: type,
425 comptime YieldValue: type,
426 comptime HandlerContext: type,
427 comptime handler: anytype,
428 ) type {
429 return struct {
430 value: YieldValue,
431 handler_context: HandlerContext,
432
433 fn onYield(raw_resume: *raw.Resume, arg: ?*anyopaque) callconv(.c) ?*anyopaque {
434 const environment: *@This() = @ptrCast(@alignCast(arg.?));
435 var continued = false;
436 var raw_result: ?*anyopaque = null;
437 const continuation: Continuation(ResumeValue, PromptResult) = .{
438 .raw_resume = raw_resume,
439 .continued = &continued,
440 .raw_result = &raw_result,
441 };
442 handler(continuation, environment.handler_context, environment.value);
443 if (!continued) {
444 std.debug.panic("mprompt.yieldWith handlers must continue the prompt before returning", .{});
445 }
446 return raw_result;
447 }
448 };
449 }
450
451 /// Suspends the body at `prompt` and runs `handler(continuation, handler_context, value)` on the
452 /// stack of the code that called `run`, so that a handler outside the body computes its next value,
453 /// as a generator hands out each value or a worker asks for input. The call returns the value the
454 /// handler continues the body with. The handler receives `value` by copy. `handler_context` must be
455 /// a pointer, and any other type is a compile error. The handler must continue the body before it
456 /// returns, and a handler that returns without continuing stops the program with a panic. `prompt`
457 /// must be the prompt the body received from `run`, and the body must still be running on it, which
458 /// assertions check in safe builds. `PromptResult` must be the `Result` type of the enclosing
459 /// `run`.
460 pub fn yieldWith(
461 comptime ResumeValue: type,
462 comptime PromptResult: type,
463 prompt: *Prompt,
464 value: anytype,
465 handler_context: anytype,
466 comptime handler: *const fn (Continuation(ResumeValue, PromptResult), @TypeOf(handler_context), @TypeOf(value)) void,
467 ) ResumeValue {
468 const YieldValue = @TypeOf(value);
469 const HandlerContext = @TypeOf(handler_context);
470 requirePointer(HandlerContext, "yieldWith handler context");
471 const Environment: type = YieldWithEnvironmentType(
472 ResumeValue,
473 PromptResult,
474 YieldValue,
475 HandlerContext,
476 handler,
477 );
478 var environment: Environment = .{
479 .value = value,
480 .handler_context = handler_context,
481 };
482 return readSlot(ResumeValue, raw.yieldPrompt(prompt, Environment.onYield, &environment));
483 }
484
485 fn writeSlot(comptime T: type, slot: *Slot(T), value: T) void {
486 if (comptime T != void) {
487 slot.value = value;
488 }
489 }
490
491 fn slotPtr(comptime T: type, slot: *Slot(T)) ?*anyopaque {
492 if (comptime T != void) {
493 return @ptrCast(slot);
494 }
495 return null;
496 }
497
498 fn readSlot(comptime T: type, ptr: ?*anyopaque) T {
499 if (comptime T != void) {
500 const slot: *Slot(T) = @ptrCast(@alignCast(ptr.?));
501 return slot.value;
502 }
503 return {};
504 }
505
506 fn outcomeFromRaw(
507 comptime ResumeValue: type,
508 comptime Result: type,
509 raw_result: ?*anyopaque,
510 done_marker: *u8,
511 result: *Slot(Result),
512 ) PromptOutcome(ResumeValue, Result) {
513 const done_ptr: ?*anyopaque = @ptrCast(done_marker);
514 if (raw_result == done_ptr) {
515 return .{ .returned = readSlot(Result, slotPtr(Result, result)) };
516 }
517
518 return .{
519 .suspended = .{
520 .raw_resume = @ptrCast(@alignCast(raw_result.?)),
521 .done_marker = done_marker,
522 .result = result,
523 },
524 };
525 }
526
527 fn requirePointer(comptime T: type, comptime name: []const u8) void {
528 switch (@typeInfo(T)) {
529 .pointer => {},
530 else => @compileError(name ++ " must be a pointer"),
531 }
532 }
533
534 const Counter = struct {
535 seen: usize = 0,
536 };
537
538 fn returnCount(_: *raw.Prompt, counter: *Counter) usize {
539 counter.seen += 1;
540 return counter.seen + 40;
541 }
542
543 test "typed prompt run returns a Zig value" {
544 var counter: Counter = .{};
545 try std.testing.expectEqual(@as(usize, 41), run(usize, &counter, returnCount));
546 try std.testing.expectEqual(@as(usize, 1), counter.seen);
547 }
548
549 fn fallibleBody(_: *raw.Prompt, counter: *Counter) error{Done}!usize {
550 counter.seen += 1;
551 return error.Done;
552 }
553
554 test "typed prompt run preserves error unions" {
555 var counter: Counter = .{};
556 try std.testing.expectError(error.Done, run(error{Done}!usize, &counter, fallibleBody));
557 try std.testing.expectEqual(@as(usize, 1), counter.seen);
558 }
559
560 fn noContextBody(_: *raw.Prompt) usize {
561 return 42;
562 }
563
564 test "typed prompt run supports no-context bodies" {
565 try std.testing.expectEqual(@as(usize, 42), runWithoutContext(usize, noContextBody));
566 }
567
568 const SuspendedContext = struct {
569 seen: usize = 0,
570 };
571
572 fn suspendOnce(prompt: *raw.Prompt, context: *SuspendedContext) usize {
573 const resumed = suspendPrompt(usize, prompt);
574 context.seen = resumed;
575 return resumed + 1;
576 }
577
578 test "typed suspended run resumes an escaped prompt" {
579 var context: SuspendedContext = .{};
580 var prompt_run = SuspendedRun(usize, usize, *SuspendedContext).init(&context);
581
582 const first = prompt_run.start(suspendOnce);
583 const suspended = switch (first) {
584 .suspended => |continuation| continuation,
585 .returned => return error.ExpectedSuspension,
586 };
587
588 const second = suspended.continueWith(41);
589 const result = switch (second) {
590 .returned => |value| value,
591 .suspended => return error.UnexpectedSuspension,
592 };
593
594 try std.testing.expectEqual(@as(usize, 42), result);
595 try std.testing.expectEqual(@as(usize, 41), context.seen);
596 }
597
598 fn suspendTwice(prompt: *raw.Prompt, context: *SuspendedContext) usize {
599 context.seen += suspendPrompt(usize, prompt);
600 context.seen += suspendPrompt(usize, prompt);
601 return context.seen;
602 }
603
604 test "typed suspended run can suspend again after resume" {
605 var context: SuspendedContext = .{};
606 var prompt_run = SuspendedRun(usize, usize, *SuspendedContext).init(&context);
607
608 const first = prompt_run.start(suspendTwice);
609 const first_suspended = switch (first) {
610 .suspended => |continuation| continuation,
611 .returned => return error.ExpectedSuspension,
612 };
613
614 const second = first_suspended.continueWith(13);
615 const second_suspended = switch (second) {
616 .suspended => |continuation| continuation,
617 .returned => return error.ExpectedSuspension,
618 };
619
620 const third = second_suspended.continueWith(29);
621 const result = switch (third) {
622 .returned => |value| value,
623 .suspended => return error.UnexpectedSuspension,
624 };
625
626 try std.testing.expectEqual(@as(usize, 42), result);
627 try std.testing.expectEqual(@as(usize, 42), context.seen);
628 }
629
630 fn useStackPages(kb: usize) void {
631 var top: u8 = 0;
632 const sp = @intFromPtr(&top);
633 const page_size = 4096;
634 const page_count = (kb * 1024 + page_size - 1) / page_size;
635 var checksum: u8 = 0;
636 var page: usize = 0;
637 while (page < page_count) : (page += 1) {
638 const address: *volatile u8 = @ptrFromInt(sp - page * page_size);
639 checksum +%= address.*;
640 }
641 std.mem.doNotOptimizeAway(checksum);
642 }
643
644 const StackWorkerEnv = struct {
645 completed: usize = 0,
646 };
647
648 fn stackUsingAsyncWorker(prompt: *raw.Prompt, env: *StackWorkerEnv) usize {
649 const stack_kb = suspendPrompt(usize, prompt);
650 useStackPages(stack_kb);
651 env.completed += 1;
652 return 1;
653 }
654
655 test "scheduler-style async prompts resume active workers" {
656 const worker_count = 16;
657 const request_count = 256;
658 const stack_kb = 8;
659 const WorkerRun = SuspendedRun(usize, usize, *StackWorkerEnv);
660 const Worker = SuspendedPrompt(usize, usize);
661
662 var envs: [worker_count]StackWorkerEnv = @as([worker_count]StackWorkerEnv, @splat(.{}));
663 var runs: [worker_count]WorkerRun = undefined;
664 for (&runs, &envs) |*worker_run, *env| {
665 worker_run.* = WorkerRun.init(env);
666 }
667
668 var workers: [worker_count]?Worker = @splat(null);
669 var completed: usize = 0;
670 var i: usize = 0;
671 while (i < request_count + worker_count) : (i += 1) {
672 const slot = i % worker_count;
673 if (workers[slot]) |continuation| {
674 const outcome = continuation.continueWith(stack_kb);
675 completed += switch (outcome) {
676 .returned => |value| value,
677 .suspended => return error.UnexpectedSuspension,
678 };
679 workers[slot] = null;
680 }
681
682 if (i < request_count) {
683 const outcome = runs[slot].start(stackUsingAsyncWorker);
684 workers[slot] = switch (outcome) {
685 .suspended => |continuation| continuation,
686 .returned => return error.ExpectedSuspension,
687 };
688 }
689 }
690
691 try std.testing.expectEqual(@as(usize, request_count), completed);
692 var observed: usize = 0;
693 for (envs) |env| {
694 observed += env.completed;
695 }
696 try std.testing.expectEqual(@as(usize, request_count), observed);
697 for (workers) |worker| {
698 try std.testing.expect(worker == null);
699 }
700 }
701
702 const YieldContext = struct {
703 yielded: usize = 0,
704 };
705
706 fn continueSuspended(
707 continuation: Continuation(usize, usize),
708 context: *YieldContext,
709 value: usize,
710 ) void {
711 context.yielded = value;
712 _ = continuation.continueWith(value + 1);
713 }
714
715 fn addAfterSuspend(prompt: *raw.Prompt, context: *YieldContext) usize {
716 const resumed = yieldWith(
717 usize,
718 usize,
719 prompt,
720 @as(usize, 40),
721 context,
722 continueSuspended,
723 );
724 return resumed + 1;
725 }
726
727 test "typed yield passes values through a continuation" {
728 var context: YieldContext = .{};
729 try std.testing.expectEqual(@as(usize, 42), run(usize, &context, addAfterSuspend));
730 try std.testing.expectEqual(@as(usize, 40), context.yielded);
731 }
732
733 const Collector = struct {
734 values: [8]usize = @splat(0),
735 count: usize = 0,
736 };
737
738 fn collectAndContinueTail(
739 continuation: Continuation(void, void),
740 collector: *Collector,
741 value: usize,
742 ) void {
743 collector.values[collector.count] = value;
744 collector.count += 1;
745 continuation.continueTailWithoutValue();
746 }
747
748 fn produceValues(prompt: *raw.Prompt, collector: *Collector) void {
749 var i: usize = 0;
750 while (i < 4) : (i += 1) {
751 yieldWith(
752 void,
753 void,
754 prompt,
755 i,
756 collector,
757 collectAndContinueTail,
758 );
759 }
760 }
761
762 test "typed yield supports tail continuation without payloads" {
763 var collector: Collector = .{};
764 run(void, &collector, produceValues);
765
766 try std.testing.expectEqual(@as(usize, 4), collector.count);
767 for (collector.values[0..collector.count], 0..) |value, index| {
768 try std.testing.expectEqual(index, value);
769 }
770 }