luby/crates/luaffi/src/future.rs

use crate::{
    __internal::{display, type_id},
    Cdef, CdefBuilder, FfiReturnConvention, IntoFfi, Metatype, MetatypeBuilder, Type, TypeBuilder,
    TypeType, UnsafeExternCFn,
};
use luaify::luaify;
use std::{
    fmt::Display,
    marker::PhantomPinned,
    mem,
    pin::Pin,
    ptr,
    task::{Context, Poll},
};

// guardrail bytes prepended to all lua_future values in the header part to hopefully try to prevent
// misinterpreting other kinds of cdata as a lua_future if the lua user code yields a cdata that is
// not a lua_future for some odd reason.
type Signature = u64;
const SIGNATURE: Signature = Signature::from_ne_bytes(*b"\x00lb_poll");

#[repr(C)]
#[allow(non_camel_case_types)]
pub struct lua_future<F: Future<Output: IntoFfi>> {
    //
    // SAFETY: LuaJIT guarantees that cdata payloads, which are GC-managed, are never relocated
    // (i.e. pinned). We can safely assume that we are pinned and poll the future inside this
    // wrapper. We use this to our advantage by storing the future value directly inside the cdata
    // payload instead of boxing the future and introducing indirection.
    //
    // https://github.com/LuaJIT/LuaJIT/issues/1167#issuecomment-1968047229
    //
    // .sig and .poll fields MUST come first. .poll is only to be called by the Rust async runtime
    // to advance the future without knowing its type (see `lua_pollable` below).
    //
    // This assumes that the crate containing the async runtime and the crate containing the future
    // type are ABI-compatible (compiled by the same compiler with the same target into the same
    // binary). This is always the case for luby because all modules are statically linked into one
    // binary.
    //
    // only .take and .drop are visible to Lua itself; other fields are opaque.
    //
    sig: Signature,
    poll: fn(Pin<&mut Self>, cx: &mut Context) -> Poll<()>,
    state: State<F>,
    take: unsafe extern "C" fn(&mut Self) -> <F::Output as IntoFfi>::Into,
    drop: unsafe extern "C" fn(&mut Self),
}

#[repr(C)]
#[allow(non_camel_case_types)]
pub struct lua_pollable {
    //
    // SAFETY: The only way to obtain a reference to a `lua_pollable` is by returning a
    // `lua_future<T>` from Rust to Lua, which LuaJIT boxes into cdata and `coroutine.yield`'s back
    // to Rust, then casting the yielded pointer value to `*mut lua_pollable`.
    //
    // This is the type-erased "header" part of a `lua_future<T>` which allows the async runtime to
    // poll the future without knowing its concrete type (essentially dynamic dispatch). It has the
    // same layout as `lua_future<T>` without the state part.
    //
    sig: Signature,
    poll: fn(Pin<&mut Self>, cx: &mut Context) -> Poll<()>,
    _phantom: PhantomPinned,
}

enum State<F: Future> {
    Pending(F),
    Fulfilled(F::Output),
    Complete,
}

impl<F: Future<Output: IntoFfi>> lua_future<F> {
    pub fn new(fut: F) -> Self {
        Self {
            sig: SIGNATURE,
            poll: Self::poll,
            state: State::Pending(fut),
            take: Self::take,
            drop: Self::drop,
        }
    }

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
        // SAFETY: we do not ever move the future here, so this is safe
        let this = unsafe { Pin::into_inner_unchecked(self) };
        match this.state {
            State::Pending(ref mut fut) => match unsafe { Pin::new_unchecked(fut) }.poll(cx) {
                Poll::Pending => Poll::Pending,
                Poll::Ready(value) => Poll::Ready(this.state = State::Fulfilled(value)), // drop the future in-place
            },
            State::Fulfilled(_) => Poll::Ready(()),
            State::Complete => unreachable!("lua_future::poll() called on a completed future"),
        }
    }

    unsafe extern "C" fn take(&mut self) -> <F::Output as IntoFfi>::Into {
        // `fut:__take()` returns the fulfilled value by-value (not by out-param) because if we
        // preallocate a cdata for the out-param and the thread for some reason gets dropped and
        // never resumed, the GC could call the destructor on an uninitialised cdata.
        match self.state {
            State::Fulfilled(_) => match mem::replace(&mut self.state, State::Complete) {
                State::Fulfilled(value) => value.convert(),
                _ => unreachable!(),
            },
            State::Pending(_) => panic!("lua_future::take() called on a pending future"),
            State::Complete => panic!("lua_future::take() called twice"),
        }
    }

    unsafe extern "C" fn drop(&mut self) {
        unsafe { ptr::drop_in_place(self) }
    }
}

impl lua_pollable {
    pub fn is_valid(&self) -> bool {
        // TODO: signature check can currently read out-of-bounds if lua code for some reason yields
        // a cdata of size less than 8 bytes that is not a lua_future. there is no easy way to fix
        // afaik this because there is no way to find the size of a cdata payload using the C API.
        self.sig == SIGNATURE
    }
}

impl Future for lua_pollable {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        assert!(self.is_valid(), "invalid lua_pollable value");
        (self.poll)(self, cx)
    }
}

unsafe impl<F: Future<Output: IntoFfi> + 'static> Type for lua_future<F> {
    fn name() -> impl Display {
        display!("future__{:x}", type_id::<F>())
    }

    fn ty() -> TypeType {
        TypeType::Aggregate
    }

    fn cdecl(name: impl Display) -> impl Display {
        display!("struct {} {name}", Self::name())
    }

    fn build(s: &mut TypeBuilder) {
        s.cdef::<Self>().metatype::<Self>();
    }
}

unsafe impl<F: Future<Output: IntoFfi> + 'static> Cdef for lua_future<F> {
    fn build(s: &mut CdefBuilder) {
        s.field_opaque(mem::offset_of!(Self, take)) // opaque .sig, .poll and .state
            .field::<UnsafeExternCFn<(&mut Self,), <F::Output as IntoFfi>::Into>>("__take")
            .field::<UnsafeExternCFn<(&mut Self,), ()>>("__drop");
    }
}

unsafe impl<F: Future<Output: IntoFfi> + 'static> Metatype for lua_future<F> {
    type Target = Self;

    fn build(s: &mut MetatypeBuilder) {
        s.metatable_raw("gc", luaify!(|self| self.__drop()));
    }
}

unsafe impl<F: Future<Output: IntoFfi> + 'static> IntoFfi for lua_future<F> {
    type Into = lua_future<F>;

    fn convention() -> FfiReturnConvention {
        // futures are always returned by-value due to rust type inference limitations
        FfiReturnConvention::ByValue
    }

    fn convert(self) -> Self::Into {
        self
    }

    fn postlude(ret: &str) -> impl Display {
        // When returning a future from Rust to Lua, yield it immediately to the runtime which will
        // poll it to completion in the background, then take the fulfilled value once the thread
        // gets resumed. Lua user code should never to worry about awaiting futures.
        //
        // Once the current thread gets resumed and we take the future's fulfilled value, we clear
        // the finaliser on the future and forget it (there is nothing to drop once the value is
        // taken).
        //
        // `coroutine.yield` is cached as `__yield` and `ffi.gc` as `__gc` in locals (see lib.rs)
        display!(
            "__yield({ret}); {ret} = __gc({ret}, nil):__take(); {}",
            <F::Output as IntoFfi>::postlude(ret)
        )
    }
}

impl<F: IntoFuture<Output: IntoFfi>> From<F> for lua_future<F::IntoFuture> {
    fn from(value: F) -> Self {
        Self::new(value.into_future())
    }
}