dicej · February 19, 2025 16:34
diff --git a/snippets.rs b/snippets.rs
 thread_local! {
    static FINISH_HOST: std::cell::Cell<*mut u8> = std::cell::Cell::new(std::ptr::null_mut());
    static FINISH_SPAWNED: std::cell::RefCell<Vec<std::pin::Pin<Box<dyn std::future::Future<Output = wasmtime::component::HostTaskOutput> + Send + Sync + 'static>>>> = std::cell::RefCell::new(Vec::new());
 }

 fn poll<T, G: for<'a> GetHost<&'a mut T>, F: std::future::Future + ?Sized>(
    getter: G,
    store: wasmtime::VMStoreRawPtr,
    cx: &mut std::task::Context,
    future: std::pin::Pin<&mut F>,
 ) -> std::task::Poll<F::Output> {
    let mut store_cx =
        unsafe { wasmtime::StoreContextMut::new(&mut *store.0.as_ptr().cast()) };
    let result = {
        let host = &mut getter(store_cx.data_mut());
        let old = FINISH_HOST.with(|v| v.replace((host as *mut G::Host).cast()));
        let result = future.poll(cx);
        // TODO: use RAII to reset this:
        FINISH_HOST.with(|v| v.set(old));
        result
    };
    for mut future in FINISH_SPAWNED.with(|v| {
        std::mem::take(std::ops::DerefMut::deref_mut(&mut v.borrow_mut()))
    }) {
        store_cx.spawn(wasmtime::component::__internal::Box::pin(
            futures::future::poll_fn(move |cx| {
                poll(getter, store, cx, future.as_mut())
            }),
        ))
    }
    result
 }

 inst.func_wrap_concurrent(
    "[static]body.finish",
    move |caller: wasmtime::StoreContextMut<'_, T>,
          (arg0,): (wasmtime::component::Resource<Body>,)| {
        let mut accessor = unsafe {
            wasmtime::component::Accessor::<G::Host>::new(
                || FINISH_HOST.with(|v| v.get()).cast(),
                |future| FINISH_SPAWNED.with(|v| v.borrow_mut().push(future)),
            )
        };
        let mut future = wasmtime::component::__internal::Box::pin(async move {
            let r = <G::Host as HostBody>::finish(&mut accessor, arg0).await;
            Ok((r?,))
        });
        let store = wasmtime::VMStoreRawPtr(caller.traitobj());
        wasmtime::component::__internal::Box::pin(futures::future::poll_fn(
            move |cx| poll(host_getter, store, cx, future.as_mut()),
        ))
    },
 )?;
	thread_local! {
	static FINISH_HOST: std::cell::Cell<*mut u8> = std::cell::Cell::new(std::ptr::null_mut());
	static FINISH_SPAWNED: std::cell::RefCell<Vec<std::pin::Pin<Box<dyn std::future::Future<Output = wasmtime::component::HostTaskOutput> + Send + Sync + 'static>>>> = std::cell::RefCell::new(Vec::new());
	}

	fn poll<T, G: for<'a> GetHost<&'a mut T>, F: std::future::Future + ?Sized>(
	getter: G,
	store: wasmtime::VMStoreRawPtr,
	cx: &mut std::task::Context,
	future: std::pin::Pin<&mut F>,
	) -> std::task::Poll<F::Output> {
	let mut store_cx =
	unsafe { wasmtime::StoreContextMut::new(&mut *store.0.as_ptr().cast()) };
	let result = {
	let host = &mut getter(store_cx.data_mut());
	let old = FINISH_HOST.with(\|v\| v.replace((host as *mut G::Host).cast()));
	let result = future.poll(cx);
	// TODO: use RAII to reset this:
	FINISH_HOST.with(\|v\| v.set(old));
	result
	};
	for mut future in FINISH_SPAWNED.with(\|v\| {
	std::mem::take(std::ops::DerefMut::deref_mut(&mut v.borrow_mut()))
	}) {
	store_cx.spawn(wasmtime::component::__internal::Box::pin(
	futures::future::poll_fn(move \|cx\| {
	poll(getter, store, cx, future.as_mut())
	}),
	))
	}
	result
	}

	inst.func_wrap_concurrent(
	"[static]body.finish",
	move \|caller: wasmtime::StoreContextMut<'_, T>,
	(arg0,): (wasmtime::component::Resource<Body>,)\| {
	let mut accessor = unsafe {
	wasmtime::component::Accessor::<G::Host>::new(
	\|\| FINISH_HOST.with(\|v\| v.get()).cast(),
	\|future\| FINISH_SPAWNED.with(\|v\| v.borrow_mut().push(future)),
	)
	};
	let mut future = wasmtime::component::__internal::Box::pin(async move {
	let r = <G::Host as HostBody>::finish(&mut accessor, arg0).await;
	Ok((r?,))
	});
	let store = wasmtime::VMStoreRawPtr(caller.traitobj());
	wasmtime::component::__internal::Box::pin(futures::future::poll_fn(
	move \|cx\| poll(host_getter, store, cx, future.as_mut()),
	))
	},
	)?;