orchid/orchid-base/src/binary.rs

use std::mem;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

use crate::api;

type WideBox = Box<dyn Future<Output = ()>>;

static OWNED_VTABLE: RawWakerVTable = RawWakerVTable::new(
	|data| {
		let data = unsafe { Rc::<api::binary::OwnedWakerBin>::from_raw(data as *const _) };
		let val = RawWaker::new(Rc::into_raw(data.clone()) as *const (), &OWNED_VTABLE);
		// Clone must create a duplicate of the Rc, so it has to be un-leaked, cloned,
		// then leaked again.
		let _ = Rc::into_raw(data);
		val
	},
	|data| {
		// Wake must awaken the task and then clean up the state, so the waker must be
		// un-leaked
		let data = unsafe { Rc::<api::binary::OwnedWakerBin>::from_raw(data as *const _) };
		(data.wake)(data.data);
		mem::drop(data);
	},
	|data| {
		// Wake-by-ref must awaken the task while preserving the future, so the Rc is
		// untouched
		let data = unsafe { (data as *const api::binary::OwnedWakerBin).as_ref() }.unwrap();
		(data.wake_ref)(data.data);
	},
	|data| {
		// Drop must clean up the state, so the waker must be un-leaked
		let data = unsafe { Rc::<api::binary::OwnedWakerBin>::from_raw(data as *const _) };
		(data.drop)(data.data);
		mem::drop(data);
	},
);

struct BorrowedWakerData<'a> {
	go_around: &'a mut bool,
	cx: api::binary::FutureContextBin,
}
static BORROWED_VTABLE: RawWakerVTable = RawWakerVTable::new(
	|data| {
		let data = unsafe { (data as *mut BorrowedWakerData).as_mut() }.unwrap();
		let owned_data = Rc::<api::binary::OwnedWakerBin>::new((data.cx.waker)(data.cx.data));
		RawWaker::new(Rc::into_raw(owned_data) as *const (), &OWNED_VTABLE)
	},
	|data| *unsafe { (data as *mut BorrowedWakerData).as_mut() }.unwrap().go_around = true,
	|data| *unsafe { (data as *mut BorrowedWakerData).as_mut() }.unwrap().go_around = true,
	|_data| {},
);

/// Convert a future to a binary-compatible format that can be sent across
/// dynamic library boundaries
#[must_use]
pub fn future_to_vt<Fut: Future<Output = ()> + 'static>(fut: Fut) -> api::binary::FutureBin {
	let wide_box = Box::new(fut) as WideBox;
	let data = Box::into_raw(Box::new(wide_box));
	extern "C" fn drop(raw: *const ()) {
		mem::drop(unsafe { Box::<WideBox>::from_raw(raw as *mut _) })
	}
	extern "C" fn poll(raw: *const (), cx: api::binary::FutureContextBin) -> api::binary::UnitPoll {
		let mut this = unsafe { Pin::new_unchecked(&mut **(raw as *mut WideBox).as_mut().unwrap()) };
		loop {
			let mut go_around = false;
			let borrowed_waker = unsafe {
				Waker::from_raw(RawWaker::new(
					&mut BorrowedWakerData { go_around: &mut go_around, cx } as *mut _ as *const (),
					&BORROWED_VTABLE,
				))
			};
			let mut ctx = Context::from_waker(&borrowed_waker);
			let result = this.as_mut().poll(&mut ctx);
			if matches!(result, Poll::Ready(())) {
				break api::binary::UnitPoll::Ready;
			}
			if !go_around {
				break api::binary::UnitPoll::Pending;
			}
		}
	}
	api::binary::FutureBin { data: data as *const _, drop, poll }
}

struct VirtualFuture {
	vt: api::binary::FutureBin,
}
impl Unpin for VirtualFuture {}
impl Future for VirtualFuture {
	type Output = ();
	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
		extern "C" fn waker(raw: *const ()) -> api::binary::OwnedWakerBin {
			let waker = unsafe { (raw as *mut Context).as_mut() }.unwrap().waker().clone();
			let data = Box::into_raw(Box::<Waker>::new(waker)) as *const ();
			return api::binary::OwnedWakerBin { data, drop, wake, wake_ref };
			extern "C" fn drop(raw: *const ()) {
				mem::drop(unsafe { Box::<Waker>::from_raw(raw as *mut Waker) })
			}
			extern "C" fn wake(raw: *const ()) {
				unsafe { Box::<Waker>::from_raw(raw as *mut Waker) }.wake();
			}
			extern "C" fn wake_ref(raw: *const ()) {
				unsafe { (raw as *mut Waker).as_mut() }.unwrap().wake_by_ref();
			}
		}
		let cx = api::binary::FutureContextBin { data: cx as *mut Context as *const (), waker };
		let result = (self.vt.poll)(self.vt.data, cx);
		match result {
			api::binary::UnitPoll::Pending => Poll::Pending,
			api::binary::UnitPoll::Ready => Poll::Ready(()),
		}
	}
}
impl Drop for VirtualFuture {
	fn drop(&mut self) { (self.vt.drop)(self.vt.data) }
}

/// Receive a future sent across dynamic library boundaries and convert it into
/// an owned object
pub fn vt_to_future(vt: api::binary::FutureBin) -> impl Future<Output = ()> { VirtualFuture { vt } }