lace-ws 0.1.0

work-stealing library based on the lace deque
Documentation 
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

use crate::arena::picked::{Arena, ArenaBox};
use crate::common::Function;
use crate::Worker;
use std::mem::{replace, size_of, transmute};

pub struct Task<I, O> {
    f: Function<I, O>,
    io: TaskIO<I, O>,
    execute_stolen: fn(&mut TypeErased, &mut Worker),
}
enum TaskIO<I, O> {
    // the boxed input
    Input(ArenaBox<I>),
    // the output is generated by a thief, so they can't use the arena,
    // but the input must still be deallocated
    Output(ArenaBox<I>, Box<O>),
    Empty,
}
impl<I, O> Task<I, O> {
    #[inline(always)]
    pub fn new(arena: &mut Arena, f: Function<I, O>, input: I) -> Self {
        Self {
            f,
            io: TaskIO::Input(arena.alloc(input)),
            execute_stolen: Self::execute_stolen,
        }
    }

    #[inline(always)]
    pub fn take_input(&mut self, arena: &mut Arena) -> I {
        let TaskIO::Input(ipt) = replace(&mut self.io, TaskIO::Empty) else {
            panic!("Task Input Taken Twice");
        };
        arena.take(ipt)
    }
    #[inline(always)]
    pub fn take_output(&mut self, arena: &mut Arena) -> O {
        let TaskIO::Output(abox, obox) = replace(&mut self.io, TaskIO::Empty) else {
            panic!("Task Output Taken Twice or Not Written");
        };
        // the stolen value may have been copied out of the ArenaBox,
        // depending on what kind of arena we are using, so to be sure
        // not to leak memory when stealing we free it here
        arena.dealloc(abox);
        *obox
    }
    #[inline(always)]
    pub fn execute(mut self, worker: &mut Worker) -> O {
        let ipt = self.take_input(&mut worker.arena);
        (self.f)(worker, ipt)
    }

    fn execute_stolen(task_erased: &mut TypeErased, worker: &mut Worker) {
        let task: &mut Self = unsafe { transmute(task_erased) };
        task.type_check::<I, O>();
        let TaskIO::Input(mut abox) = replace(&mut task.io, TaskIO::Empty) else {
            panic!("Task Input Taken Already While Stealing");
        };
        let ipt: I = abox.steal();
        let out = (task.f)(worker, ipt);
        task.io = TaskIO::Output(abox, Box::from(out));
    }

    pub fn erase(self) -> TypeErased {
        unsafe { transmute(self) }
    }
    #[inline(always)]
    fn type_check<X, Y>(&self) {
        // check that the type is correct using the function pointer
        // as a proxy for the type ID
        if !std::ptr::fn_addr_eq(self.execute_stolen, Task::<X, Y>::execute_stolen as for<'a, 'b> fn(&'a mut TypeErased, &'b mut Worker)) {
            panic!("Task unerase()'d to incorrect base type")
        }
    }
}
pub struct TypeErased {
    _fill: [u8; size_of::<Task<u8, u8>>()],
}
impl TypeErased {
    #[inline(always)]
    pub fn unerase<I, O>(self) -> Task<I, O> {
        let task: Task<I, O> = unsafe { transmute(self) };
        task.type_check::<I, O>();
        task
    }
    #[inline(always)]
    pub fn execute_stolen(&mut self, worker: &mut Worker) {
        // the generic types don't change the layout of the task struct,
        // so we can safely access the execute_stolen pointer here
        let f = unsafe { transmute::<&mut Self, &mut Task<(), ()>>(self) }.execute_stolen;
        f(self, worker);
    }
}