      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131

    use crate::runtime::task::{Cell, Harness, Header, Schedule, State};

    use std::future::Future;
    use std::ptr::NonNull;
    use std::task::{Poll, Waker};

    /// Raw task handle
    pub(super) struct RawTask {
        ptr: NonNull<Header>,
    }

    pub(super) struct Vtable {
        /// Poll the future
        pub(super) poll: unsafe fn(NonNull<Header>),

        /// Deallocate the memory
        pub(super) dealloc: unsafe fn(NonNull<Header>),

        /// Read the task output, if complete
        pub(super) try_read_output: unsafe fn(NonNull<Header>, *mut (), &Waker),

        /// The join handle has been dropped
        pub(super) drop_join_handle_slow: unsafe fn(NonNull<Header>),

        /// Scheduler is being shutdown
        pub(super) shutdown: unsafe fn(NonNull<Header>),
    }

    /// Get the vtable for the requested `T` and `S` generics.
    pub(super) fn vtable<T: Future, S: Schedule>() -> &'static Vtable {
        &Vtable {
            poll: poll::<T, S>,
            dealloc: dealloc::<T, S>,
            try_read_output: try_read_output::<T, S>,
            drop_join_handle_slow: drop_join_handle_slow::<T, S>,
            shutdown: shutdown::<T, S>,
        }
    }

    impl RawTask {
        pub(super) fn new<T, S>(task: T) -> RawTask
        where
            T: Future,
            S: Schedule,
        {
            let ptr = Box::into_raw(Cell::<_, S>::new(task, State::new()));
            let ptr = unsafe { NonNull::new_unchecked(ptr as *mut Header) };

            RawTask { ptr }
        }

        pub(super) unsafe fn from_raw(ptr: NonNull<Header>) -> RawTask {
            RawTask { ptr }
        }

        /// Returns a reference to the task's meta structure.
        ///
        /// Safe as `Header` is `Sync`.
        pub(super) fn header(&self) -> &Header {
            unsafe { self.ptr.as_ref() }
        }

        /// Safety: mutual exclusion is required to call this function.
        pub(super) fn poll(self) {
            let vtable = self.header().vtable;
            unsafe { (vtable.poll)(self.ptr) }
        }

        pub(super) fn dealloc(self) {
            let vtable = self.header().vtable;
            unsafe {
                (vtable.dealloc)(self.ptr);
            }
        }

        /// Safety: `dst` must be a `*mut Poll<super::Result<T::Output>>` where `T`
        /// is the future stored by the task.
        pub(super) unsafe fn try_read_output(self, dst: *mut (), waker: &Waker) {
            let vtable = self.header().vtable;
            (vtable.try_read_output)(self.ptr, dst, waker);
        }

        pub(super) fn drop_join_handle_slow(self) {
            let vtable = self.header().vtable;
            unsafe { (vtable.drop_join_handle_slow)(self.ptr) }
        }

        pub(super) fn shutdown(self) {
            let vtable = self.header().vtable;
            unsafe { (vtable.shutdown)(self.ptr) }
        }
    }

    impl Clone for RawTask {
        fn clone(&self) -> Self {
            RawTask { ptr: self.ptr }
        }
    }

    impl Copy for RawTask {}

    unsafe fn poll<T: Future, S: Schedule>(ptr: NonNull<Header>) {
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.poll();
    }

    unsafe fn dealloc<T: Future, S: Schedule>(ptr: NonNull<Header>) {
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.dealloc();
    }

    unsafe fn try_read_output<T: Future, S: Schedule>(
        ptr: NonNull<Header>,
        dst: *mut (),
        waker: &Waker,
    ) {
        let out = &mut *(dst as *mut Poll<super::Result<T::Output>>);

        let harness = Harness::<T, S>::from_raw(ptr);
        harness.try_read_output(out, waker);
    }

    unsafe fn drop_join_handle_slow<T: Future, S: Schedule>(ptr: NonNull<Header>) {
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.drop_join_handle_slow()
    }

    unsafe fn shutdown<T: Future, S: Schedule>(ptr: NonNull<Header>) {
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.shutdown()
    }