fast_steal/
task.rs

1use core::{
2    hash::Hash,
3    ops::Range,
4    sync::atomic::{AtomicU64, Ordering},
5};
6
7#[derive(Debug)]
8pub struct Task {
9    start: AtomicU64,
10    end: AtomicU64,
11}
12
13impl Task {
14    pub fn remain(&self) -> u64 {
15        let start = self.start();
16        let end = self.end();
17        end.saturating_sub(start)
18    }
19
20    pub fn start(&self) -> u64 {
21        self.start.load(Ordering::Acquire)
22    }
23    pub fn set_start(&self, start: u64) {
24        self.start.store(start, Ordering::Release);
25    }
26    pub fn fetch_add_start(&self, value: u64) -> u64 {
27        self.start.fetch_add(value, Ordering::AcqRel)
28    }
29    pub fn fetch_sub_start(&self, value: u64) -> u64 {
30        self.start.fetch_sub(value, Ordering::AcqRel)
31    }
32    pub fn end(&self) -> u64 {
33        self.end.load(Ordering::Acquire)
34    }
35    pub fn set_end(&self, end: u64) {
36        self.end.store(end, Ordering::Release);
37    }
38    pub fn fetch_add_end(&self, value: u64) -> u64 {
39        self.end.fetch_add(value, Ordering::AcqRel)
40    }
41    pub fn fetch_sub_end(&self, value: u64) -> u64 {
42        self.end.fetch_sub(value, Ordering::AcqRel)
43    }
44
45    pub fn new(start: u64, end: u64) -> Self {
46        Self {
47            start: AtomicU64::new(start),
48            end: AtomicU64::new(end),
49        }
50    }
51
52    pub fn split_two(&self) -> (u64, u64) {
53        let start = self.start();
54        let end = self.end();
55        let mid = (start + end) / 2;
56        self.set_end(mid);
57        (mid, end)
58    }
59}
60
61impl PartialEq for Task {
62    fn eq(&self, other: &Self) -> bool {
63        self.start() == other.start() && self.end() == other.end()
64    }
65}
66impl Eq for Task {}
67impl PartialOrd for Task {
68    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
69        Some(self.cmp(other))
70    }
71}
72impl Ord for Task {
73    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
74        self.remain().cmp(&other.remain())
75    }
76}
77impl Hash for Task {
78    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
79        self.start().hash(state);
80        self.end().hash(state);
81    }
82}
83
84impl From<&(u64, u64)> for Task {
85    fn from(value: &(u64, u64)) -> Self {
86        Self::new(value.0, value.1)
87    }
88}
89
90impl From<&Range<u64>> for Task {
91    fn from(value: &Range<u64>) -> Self {
92        Self::new(value.start, value.end)
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn test_new_task() {
102        let task = Task::new(10, 20);
103        assert_eq!(task.start(), 10);
104        assert_eq!(task.end(), 20);
105        assert_eq!(task.remain(), 10);
106    }
107
108    #[test]
109    fn test_from_tuple() {
110        let task: Task = (&(5, 15)).into();
111        assert_eq!(task.start(), 5);
112        assert_eq!(task.end(), 15);
113    }
114
115    #[test]
116    fn test_from_range() {
117        let range = 3..8;
118        let task: Task = (&range).into();
119        assert_eq!(task.start(), range.start);
120        assert_eq!(task.end(), range.end);
121    }
122
123    #[test]
124    fn test_setters() {
125        let task = Task::new(0, 0);
126        task.set_start(7);
127        task.set_end(14);
128        assert_eq!(task.start(), 7);
129        assert_eq!(task.end(), 14);
130    }
131
132    #[test]
133    fn test_remain() {
134        let task = Task::new(10, 25);
135        assert_eq!(task.remain(), 15);
136
137        task.set_start(20);
138        assert_eq!(task.remain(), 5);
139    }
140
141    #[test]
142    fn test_partial_eq() {
143        let task1 = Task::new(1, 10);
144        let task2 = Task::new(1, 10);
145        let task3 = Task::new(2, 10);
146        let task4 = Task::new(1, 11);
147
148        assert_eq!(task1, task2);
149        assert_ne!(task1, task3);
150        assert_ne!(task1, task4);
151    }
152
153    #[test]
154    fn test_split_two() {
155        let task = Task::new(1, 6); // 1, 2, 3, 4, 5
156        let (mid, end) = task.split_two();
157        assert_eq!(task.start(), 1);
158        assert_eq!(task.end(), 3);
159        assert_eq!(mid, 3);
160        assert_eq!(end, 6);
161    }
162}