swiftui_native/

asyncfn.rs

//! Swift `async` closures, supplied from Rust.
//!
//! `refreshable(action: @Sendable () async -> ())` cannot take an ordinary
//! function pointer. An async closure's first word is not the function — it is
//! an **async function pointer**, a two-field record the compiler emits:
//!
//! | offset | contents |
//! |--------|----------|
//! | `+0`   | 32-bit relative pointer to the async function |
//! | `+4`   | the context size the function expects |
//!
//! The function itself runs under the async convention: its context arrives in
//! **x22**, and it returns by tail-calling the continuation stored at `+8` in
//! that context. An empty `{ }` compiles to exactly two instructions —
//! `ldr x0, [x22, #8]` then `br x0` — which is the shape the thunks below keep,
//! with the Rust callback run in between.
//!
//! The work happens synchronously and the continuation is called straight away:
//! this is an async closure that never suspends, which is a real one, not a
//! stub. Suspension would mean allocating a child context and driving the task
//! system, and nothing here needs that.
//!
//! The thunks below are known to *run*: building `.task` on top of this and
//! rendering it does reach the Rust callback. What is parity-*verified* is the
//! value `refreshable` builds, not the closure running — a refresh action only
//! fires on a pull, which the suite does not perform.
//!
//! Like [`crate::keypath`], the record is a `static mut`: its one field is a
//! 32-bit offset to code in this image, which the heap is too far from.

use core::ffi::c_void;

/// The record an async closure's first word points at.
#[repr(C, align(4))]
pub struct FunctionPointer {
    function: i32,
    context_size: u32,
}

impl FunctionPointer {
    /// An empty record, to be filled by [`prepare`].
    pub const fn new() -> Self {
        FunctionPointer { function: 0, context_size: 0 }
    }
}

impl Default for FunctionPointer {
    fn default() -> Self {
        Self::new()
    }
}

/// The context size to declare.
///
/// Thirty-two is what the compiler's partial-apply forwarder asks for — an empty
/// closure's own record says sixteen, but the forwarder is what a call site
/// actually passes, and asking for more than is used is free.
const EMPTY_CONTEXT: u32 = 32;

/// Point `record` at `thunk` and arm the thunk to call `action`.
///
/// Returns the two words of the closure: the record, and a null context —
/// exactly what the compiler passes for a closure that captures nothing.
///
/// # Safety
/// `record` must be a `static mut` in this image, and `slot` one of the three
/// below. Arming a slot twice is allowed and the later action wins — two
/// `.task` modifiers on two different views share a slot, and both run the
/// second one's action, which is what the `fn()` version did too.
pub unsafe fn prepare(
    record: &'static mut FunctionPointer,
    slot: usize,
    action: impl Fn() + Send + Sync + 'static,
) -> [u64; 2] {
    unsafe {
        let (code, target, shim) = thunks(slot);
        // The thunk `blr`s straight to the target, so the target has to be
        // executable code taking no arguments — which a boxed closure is not.
        // Each slot therefore owns a shim that calls through the box beside it.
        // The action has to be `Send + Sync` because this one is genuinely
        // async: Swift may resume the task on a cooperative thread, not the
        // main one, so unlike every other callback here it is not main-thread
        // by construction.
        // Last arming wins, as it did when this was a bare `fn` pointer; the
        // previous action is leaked rather than dropped, because SwiftUI may
        // already be running it.
        ACTIONS[slot].store(
            Box::into_raw(Box::new(Box::new(action) as Box<dyn Fn() + Send + Sync>)),
            Ordering::Release,
        );
        *target = shim;
        let at = &record.function as *const i32 as *const u8;
        let delta = code as isize - at as isize;
        record.function = i32::try_from(delta).expect("async function pointer out of range");
        record.context_size = EMPTY_CONTEXT;
        [record as *const FunctionPointer as u64, 0]
    }
}

// x22 holds the async context and `[x22, #8]` is the continuation to tail-call
// on return — which is the whole of what an empty `{ }` compiles to.
//
// The prologue is the one the compiler emits for an async function that calls
// out and then returns. Bit 60 of x29 marks the frame as async so the unwinder
// walks the context chain rather than the stack; x22 is saved because the
// continuation is read from it after the call, and x30 because `blr` overwrites
// it. Leaving any of that out lets the closure run to completion and then kills
// the process on the way back — which `refreshable` never notices, because a
// refresh action only fires on a pull, and `task` notices immediately.
macro_rules! async_thunks {
    ($($name:ident, $sym:literal, $target:ident;)*) => {$(
        crate::thunk::runtime_thunk! {
            /// An `async` closure body that runs a Rust callback and resumes.
            fn $name = $sym via $target {
                prologue: [
                    "orr x29, x29, #0x1000000000000000",
                    "sub sp, sp, #0x20",
                    "stp x29, x30, [sp, #0x10]",
                    "str x22, [sp, #0x8]",
                    "add x29, sp, #0x10",
                ],
                branch: "blr",
                epilogue: [
                    "ldr x22, [sp, #0x8]",
                    "ldp x29, x30, [sp, #0x10]",
                    "and x29, x29, #0xefffffffffffffff",
                    "add sp, sp, #0x20",
                    "ldr x0, [x22, #8]",
                    "br x0",
                ],
            }
        }
    )*};
}

async_thunks! {
    async0, "rswift_async0", RSWIFT_ASYNC0;
    async1, "rswift_async1", RSWIFT_ASYNC1;
    async2, "rswift_async2", RSWIFT_ASYNC2;
}

/// Where each slot's action lives.
///
/// A slot, not a single static, because a thunk carries no context of its own:
/// an async closure is a *record* plus a context word, and the record only
/// holds a 32-bit offset to code. The one way back to the right closure is to
/// have branched from the right thunk, so there are as many simultaneous
/// actions as there are thunks — three, one per call site: `refreshable`,
/// `task`, and `task_id`. Slot 2 used to fall through to slot 1's thunk, so a
/// view carrying both a `.task` and a `.task(id:)` ran whichever was applied
/// second, twice.
///
/// **Within a slot, the last arming wins.** Two `.task` modifiers on two
/// different views share slot 1, and both will run the second one's action.
/// This is what the `fn()` version did as well — it stored the bare pointer in
/// the same one slot — so nothing regressed, but it is a real ceiling, and the
/// place to lift it is the context word: a per-closure Swift heap object,
/// passed as the second word instead of the null below, would let one thunk
/// serve every call. That needs the context to survive SwiftUI's retain and
/// release of it, which a Rust box does not.
static ACTIONS: [AtomicPtr<Box<dyn Fn() + Send + Sync>>; SLOTS] =
    [const { AtomicPtr::new(core::ptr::null_mut()) }; SLOTS];

use core::sync::atomic::{AtomicPtr, Ordering};

const SLOTS: usize = 3;

/// One shim per slot, since each has to name its own index.
macro_rules! shims {
    ($($name:ident = $i:expr;)*) => {$(
        unsafe extern "C" fn $name() {
            let f = ACTIONS[$i].load(Ordering::Acquire);
            if !f.is_null() {
                unsafe { (**f)() };
            }
        }
    )*};
}

shims! { call0 = 0; call1 = 1; call2 = 2; }

/// The async thunks, by slot — the code to branch from, the slot to write, and
/// the shim to write into it.
fn thunks(slot: usize) -> (*const c_void, *mut *const c_void, *const c_void) {
    match slot {
        0 => (async0 as *const c_void, &raw mut RSWIFT_ASYNC0, call0 as *const c_void),
        1 => (async1 as *const c_void, &raw mut RSWIFT_ASYNC1, call1 as *const c_void),
        _ => (async2 as *const c_void, &raw mut RSWIFT_ASYNC2, call2 as *const c_void),
    }
}

// ── A task's entry, which *does* get a context ──
//
// The three slots above exist because a SwiftUI `@ViewBuilder` async closure is
// handed `(record, context)` and SwiftUI passes back the context it was given —
// which for a Rust closure is nothing it can own, so the closure has to live in
// a slot the thunk names by its own address. That is a real ceiling: three
// simultaneous actions, last arming wins within a slot.
//
// `swift_task_create` is not like that. Its fifth argument *is* the closure
// context, and the entry receives it in **x20** — read off the disassembly of
// `Task.detached { }`, whose entry starts `ldr x9, [x20, #0x10]` before x20
// could have been set by anything else. So a task carries its own closure and
// there is no slot and no ceiling: as many tasks as you like, at once.
//
// The frame is the same as above; the one addition is moving x20 into the first
// argument register before the call, which nothing in the prologue disturbs.
crate::thunk::runtime_thunk! {
    /// A `Task`'s entry: run the closure the context names, then resume.
    fn task_entry = "rswift_task_entry" via RSWIFT_TASK_ENTRY {
        prologue: [
            "orr x29, x29, #0x1000000000000000",
            "sub sp, sp, #0x20",
            "stp x29, x30, [sp, #0x10]",
            "str x22, [sp, #0x8]",
            "add x29, sp, #0x10",
            "mov x0, x20",
        ],
        branch: "blr",
        epilogue: [
            "ldr x22, [sp, #0x8]",
            "ldp x29, x30, [sp, #0x10]",
            "and x29, x29, #0xefffffffffffffff",
            "add sp, sp, #0x20",
            "ldr x0, [x22, #8]",
            "br x0",
        ],
    }
}

/// How a task's completion is reported back to whoever holds its handle.
///
/// A Swift `Task`'s result is delivered by awaiting it, and Rust has nothing to
/// await from — so the entry signals this instead, and
/// [`crate::runtime::Task::wait`] blocks on it. Shared by `Arc` between the
/// handle and the context, so a handle dropped early frees it and a task that
/// outlives its handle still has somewhere to write.
pub(crate) struct TaskState {
    done: std::sync::Mutex<bool>,
    woken: std::sync::Condvar,
}

impl TaskState {
    pub(crate) fn new() -> Self {
        Self { done: std::sync::Mutex::new(false), woken: std::sync::Condvar::new() }
    }

    /// Has the task's closure returned?
    pub(crate) fn is_finished(&self) -> bool {
        *self.done.lock().expect("task state")
    }

    /// Block until it has, or until `limit` elapses. `true` if it finished.
    pub(crate) fn wait(&self, limit: Option<std::time::Duration>) -> bool {
        let mut done = self.done.lock().expect("task state");
        match limit {
            None => {
                while !*done {
                    done = self.woken.wait(done).expect("task state");
                }
                true
            }
            Some(limit) => {
                let deadline = std::time::Instant::now() + limit;
                while !*done {
                    let left = deadline.saturating_duration_since(std::time::Instant::now());
                    if left.is_zero() {
                        return false;
                    }
                    let (next, timeout) = self
                        .woken
                        .wait_timeout(done, left)
                        .expect("task state");
                    done = next;
                    if timeout.timed_out() && !*done {
                        return false;
                    }
                }
                true
            }
        }
    }

    fn finish(&self) {
        *self.done.lock().expect("task state") = true;
        self.woken.notify_all();
    }
}

/// What a task's context carries.
///
/// Laid out like [`crate::abi::ClosureCtx`] and for the same reason: the runtime
/// treats it as a `HeapObject`, so it needs an `isa` and a refcount, and the
/// refcount is the immortal pattern so the retains and releases around it are
/// no-ops. A Rust `Box` would be freed by the first release.
#[repr(C)]
struct TaskCtx {
    isa: *const c_void,
    refcount: u64,
    /// The closure, owned. Null once the entry has taken it.
    body: *mut Box<dyn FnOnce() + Send>,
    /// The completion signal, as a raw `Arc`. Null once the entry has used it.
    state: *const TaskState,
}

/// Contexts whose task has finished, ready to carry another.
///
/// Not freed, and that is the point: the runtime releases the closure context
/// after the task completes, and with the immortal refcount that release is a
/// no-op *that still reads the header*. So the allocation has to outlive the
/// task — but only the header does, and everything after it can be reused. The
/// alternative was a fresh thirty-two bytes leaked per task forever; this bounds
/// it at the high-water mark of tasks in flight.
fn retired() -> &'static std::sync::Mutex<Vec<usize>> {
    static R: std::sync::OnceLock<std::sync::Mutex<Vec<usize>>> = std::sync::OnceLock::new();
    R.get_or_init(|| std::sync::Mutex::new(Vec::new()))
}

/// Branch target of [`task_entry`]: run the closure the context names.
///
/// # Safety
/// `ctx` must be a `TaskCtx` from [`task_entry_point`].
#[unsafe(no_mangle)]
unsafe extern "C" fn rswift_task_body(ctx: *mut c_void) {
    let c = unsafe { &mut *(ctx as *mut TaskCtx) };

    let body = core::mem::replace(&mut c.body, core::ptr::null_mut());
    if !body.is_null() {
        let body = unsafe { Box::from_raw(body) };
        // A Rust panic must not unwind into the Swift async frame this was
        // called from: there is no landing pad there, and what happens instead
        // is unspecified. Caught and reported, and the task still completes —
        // which is what a Swift task does when its body throws.
        if std::panic::catch_unwind(std::panic::AssertUnwindSafe(body)).is_err() {
            eprintln!("rswift: a detached task's closure panicked; the task completed anyway");
        }
    }

    let state = core::mem::replace(&mut c.state, core::ptr::null());
    if !state.is_null() {
        // Reclaiming the `Arc` the context was holding, and signalling before
        // dropping it: a waiter is holding the other reference.
        let state = unsafe { std::sync::Arc::from_raw(state) };
        state.finish();
    }

    retired().lock().expect("retired contexts").push(ctx as usize);
}

/// The async function pointer and context for a task running `action`.
///
/// `swift_task_create` wants the two-field *record*, not the function, and the
/// context separately.
///
/// # Safety
/// The record is a `static mut` in this image, shared by every task; only its
/// two fields are written, and they are the same values every time.
pub(crate) unsafe fn task_entry_point(
    action: Box<dyn FnOnce() + Send>,
    state: &std::sync::Arc<TaskState>,
) -> (*const c_void, *mut c_void) {
    static mut RECORD: FunctionPointer = FunctionPointer::new();
    static ONCE: std::sync::Once = std::sync::Once::new();
    unsafe {
        ONCE.call_once(|| {
            RSWIFT_TASK_ENTRY = rswift_task_body as *const c_void;
            let record = crate::static_mut(&raw mut RECORD);
            let at = &record.function as *const i32 as *const u8;
            let delta = task_entry as *const u8 as isize - at as isize;
            record.function =
                i32::try_from(delta).expect("async function pointer out of range");
            record.context_size = EMPTY_CONTEXT;
        });

        let body = Box::into_raw(Box::new(action));
        let state = std::sync::Arc::into_raw(std::sync::Arc::clone(state));

        let reused = retired().lock().expect("retired contexts").pop();
        let ctx = match reused {
            Some(addr) => addr as *mut TaskCtx,
            None => Box::leak(Box::new(TaskCtx {
                isa: crate::abi::immortal_isa(),
                refcount: crate::abi::IMMORTAL_REFCOUNT,
                body: core::ptr::null_mut(),
                state: core::ptr::null(),
            })) as *mut TaskCtx,
        };
        // Only the payload is written; the header a reused context already has
        // is the one the runtime may still be releasing against.
        (*ctx).body = body;
        (*ctx).state = state;

        (&raw const RECORD as *const c_void, ctx as *mut c_void)
    }
}