async-ecs/async-ecs/src/misc/tokio_executor.rs

use std::mem::transmute;

use futures::{
    executor::block_on,
    future::{BoxFuture, Future, FutureExt},
    join,
};
use tokio::task::{block_in_place, spawn};

use parallel_iterator::{
    Consumer, Executor, Folder, IndexedConsumer, IndexedProducer, Producer, Reducer,
};

pub struct TokioExecutor;

impl Executor for TokioExecutor {
    fn exec<P, C>(self, producer: P, consumer: C) -> C::Result
    where
        P: Producer,
        C: Consumer<P::Item>,
    {
        block_in_place(|| block_on(exec(2 * num_cpus::get(), producer, consumer)))
    }

    fn exec_indexed<P, C>(self, producer: P, consumer: C) -> C::Result
    where
        P: IndexedProducer,
        C: IndexedConsumer<P::Item>,
    {
        let _producer = producer;
        let _consumer = consumer;

        unimplemented!()
    }
}

fn exec<'a, P, C>(splits: usize, producer: P, consumer: C) -> BoxFuture<'a, C::Result>
where
    P: Producer + 'a,
    C: Consumer<P::Item> + 'a,
{
    async move {
        if consumer.is_full() {
            consumer.into_folder().complete()
        } else if splits > 0 {
            match producer.split() {
                (left_producer, Some(right_producer)) => {
                    let ((left_consumer, reducer), right_consumer) =
                        (consumer.split_off_left(), consumer);

                    let left = run_as_task(exec(splits / 2, left_producer, left_consumer));
                    let right = run_as_task(exec(splits / 2, right_producer, right_consumer));

                    let (left_result, right_result) = join!(left, right);

                    reducer.reduce(left_result, right_result)
                }
                (producer, None) => producer.fold_with(consumer.into_folder()).complete(),
            }
        } else {
            producer.fold_with(consumer.into_folder()).complete()
        }
    }
    .boxed()
}

struct Pointer<T>(*mut T);

unsafe impl<T> Send for Pointer<T> {}

async fn run_as_task<'a, T, F>(f: F) -> T
where
    T: Send + 'a,
    F: Future<Output = T> + Send + 'a,
{
    let mut result = None;
    let r = Pointer(&mut result as *mut _);

    let task: BoxFuture<'a, ()> = async move {
        unsafe {
            *r.0 = Some(f.await);
        }
    }
    .boxed();
    let task: BoxFuture<'static, ()> = unsafe { transmute(task) };

    spawn(task).await.expect("Error in iterator executor");

    result.unwrap()
}