Add ffi crate

crates/luaffi/src/future.rs

@@ -0,0 +1,172 @@
+use crate::{
+    __internal::{display, type_id},
+    CDef, CDefBuilder, Metatype, MetatypeBuilder, ToFfi, Type, TypeBuilder,
+};
+use luaify::luaify;
+use std::{
+    fmt::Display,
+    mem,
+    pin::Pin,
+    ptr,
+    task::{Context, Poll},
+};
+
+#[repr(C)]
+#[allow(non_camel_case_types)]
+pub struct lua_future<F: Future<Output: ToFfi>> {
+    //
+    // SAFETY: .poll MUST be the first field. It 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.
+    //
+    // .poll and .state are opaque to Lua itself.
+    //
+    poll: fn(&mut Self, cx: &mut Context) -> Poll<()>,
+    state: State<F>,
+    take: unsafe extern "C" fn(&mut Self) -> <F::Output as ToFfi>::To,
+    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.
+    //
+    poll: fn(&mut Self, cx: &mut Context) -> Poll<()>,
+}
+
+enum State<F: Future> {
+    Pending(F),
+    Fulfilled(F::Output),
+    Complete,
+}
+
+impl<F: Future<Output: ToFfi>> lua_future<F> {
+    pub fn new(fut: F) -> Self {
+        Self {
+            poll: Self::poll,
+            state: State::Pending(fut),
+            take: Self::take,
+            drop: Self::drop,
+        }
+    }
+
+    fn poll(&mut self, cx: &mut Context) -> Poll<()> {
+        //
+        // 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.
+        //
+        // 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
+        //
+        match self.state {
+            State::Pending(ref mut fut) => match unsafe { Pin::new_unchecked(fut) }.poll(cx) {
+                Poll::Pending => Poll::Pending,
+                Poll::Ready(value) => Poll::Ready(self.state = State::Fulfilled(value)),
+            },
+            State::Fulfilled(_) => Poll::Ready(()),
+            State::Complete => unreachable!("lua_future::poll() called on completed future"),
+        }
+    }
+
+    unsafe extern "C" fn take(&mut self) -> <F::Output as ToFfi>::To {
+        // `fut:__take()` returns the fulfilled value by-value because it is the lowest common
+        // denominator for supported return conventions (all `ToFfi` impls support return by-value;
+        // primitives e.g. don't support return by out-param because they get boxed in cdata).
+        //
+        // Plus, if we preallocate a cdata for out-param and the thread for some reason gets dropped
+        // and never resumed, 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 pending future"),
+            State::Complete => panic!("lua_future::take() called twice"),
+        }
+    }
+
+    unsafe extern "C" fn drop(&mut self) {
+        unsafe { ptr::drop_in_place(self) }
+    }
+}
+
+impl Future for lua_pollable {
+    type Output = ();
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
+        // SAFETY: see comment above in `lua_future::poll()`
+        (self.poll)(Pin::into_inner(self), cx)
+    }
+}
+
+unsafe impl<F: Future<Output: ToFfi> + 'static> Type for lua_future<F> {
+    fn name() -> impl Display {
+        display!("future__{:x}", type_id::<F>())
+    }
+
+    fn cdecl(name: impl Display) -> impl Display {
+        display!("struct future__{:x} {name}", type_id::<F>())
+    }
+
+    fn build(s: &mut TypeBuilder) {
+        s.cdef::<Self>().metatype::<Self>();
+    }
+}
+
+unsafe impl<F: Future<Output: ToFfi> + 'static> CDef for lua_future<F> {
+    fn build(s: &mut CDefBuilder) {
+        s.field_opaque(mem::offset_of!(Self, take))
+            .field::<unsafe extern "C" fn(*mut Self) -> <F::Output as ToFfi>::To>("__take")
+            .field::<unsafe extern "C" fn(*mut Self)>("__drop");
+    }
+}
+
+unsafe impl<F: Future<Output: ToFfi> + '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: ToFfi> + 'static> ToFfi for lua_future<F> {
+    type To = lua_future<F>;
+
+    fn convert(self) -> Self::To {
+        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 call drop on).
+        //
+        // `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 ToFfi>::postlude(ret)
+        )
+    }
+}
+
+impl<F: IntoFuture<Output: ToFfi>> From<F> for lua_future<F::IntoFuture> {
+    fn from(value: F) -> Self {
+        Self::new(value.into_future())
+    }
+}