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
132
133
134
135
136
137
138
139
140

use std::{
    ptr::{self, addr_of},
    sync::atomic::{AtomicU64, Ordering},
    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
    usize,
};

use const_array_init::const_arr;
use once_cell::race::OnceBox;

use crate::{
    handle::TaskHandle,
    runtime::{Scheduler, RUNTIME},
};

const IDX: usize = (1 << 6) - 1;
const IDX_MASK: usize = !IDX;

#[derive(Clone, Copy)]
pub(crate) struct Index<'a> {
    pub(crate) arena: &'a TaskArena,
    idx: usize,
}

impl<'a> Index<'a> {
    pub(crate) fn handle(&self) -> &TaskHandle {
        &self.arena.tasks[self.idx]
    }

    fn raw_waker(&self) -> RawWaker {
        unsafe fn clone(ptr: *const ()) -> RawWaker {
            RawWaker::new(ptr, &VTABLE)
        }

        unsafe fn wake(ptr: *const ()) {
            let slot = Index::from_raw(ptr as *mut ());
            let handle = slot.handle();

            let next = handle.next.load(Ordering::Acquire) as *const ();

            if next.is_null() {
                let tail = RUNTIME.tail.swap(ptr as *mut (), Ordering::AcqRel);

                if tail.is_null() {
                    Scheduler::schedule_polling(ptr as *mut ());
                } else if tail.ne(&(ptr as *mut ())) {
                    handle.next.store(tail, Ordering::Release);
                }
            }
        }

        unsafe fn drop(_ptr: *const ()) {}

        static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake, drop);

        RawWaker::new(self.into_raw(), &VTABLE)
    }

    pub(crate) unsafe fn poll(&self) -> *mut () {
        let handle = self.handle();

        let poll = match handle.task_mut() {
            Some(task) => {
                let waker = Waker::from_raw(self.raw_waker());
                let mut cx = Context::from_waker(&waker);

                Some(task.as_mut().poll(&mut cx))
            }
            None => None,
        };

        if poll.eq(&Some(Poll::Ready(()))) {
            *handle.task_mut() = None;
            self.arena
                .occupancy
                .fetch_and(!(1 << self.idx), Ordering::Release);
        }

        handle.next.swap(ptr::null_mut(), Ordering::AcqRel)
    }

    pub(crate) unsafe fn from_raw(ptr: *mut ()) -> Self {
        Self {
            arena: &*((ptr as usize & IDX_MASK) as *const TaskArena),
            idx: ptr as usize & IDX,
        }
    }

    pub(crate) fn into_raw(self) -> *mut () {
        ((addr_of!(*self.arena) as usize) | self.idx) as *mut ()
    }
}

#[repr(align(64))]
pub(crate) struct TaskArena {
    occupancy: AtomicU64,
    tasks: [TaskHandle; 64],
    next: OnceBox<TaskArena>,
}

impl TaskArena {
    pub(crate) const fn new() -> Self {
        TaskArena {
            occupancy: AtomicU64::new(0),
            tasks: const_arr!([TaskHandle; 64], |_| TaskHandle::new()),
            next: OnceBox::new(),
        }
    }
}

impl TaskArena {
    pub(crate) fn next_index(&self) -> Index {
        let mut occupancy = self.occupancy.load(Ordering::Acquire);

        let idx = loop {
            // Isolate lowest clear bit. See https://docs.rs/bitintr/latest/bitintr/trait.Blcic.html
            let least_significant_bit = !occupancy & (occupancy.wrapping_add(1));

            if least_significant_bit.ne(&0) {
                occupancy = self
                    .occupancy
                    .fetch_or(least_significant_bit, Ordering::AcqRel);

                if (occupancy & least_significant_bit).eq(&0) {
                    break least_significant_bit.trailing_zeros();
                }
            } else {
                return self
                    .next
                    .get_or_init(|| Box::new(TaskArena::new()))
                    .next_index();
            }
        };

        Index {
            arena: self,
            idx: idx as usize,
        }
    }
}