Skip to main content

compio_driver/
cancel.rs

1use std::collections::HashSet;
2
3use crate::{Key, OpCode, key::ErasedKey};
4
5pub(crate) struct CancelRegistry {
6    cancellers: HashSet<ErasedKey>,
7}
8
9/// A type-erased cancel token.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
11#[repr(transparent)]
12pub struct Cancel(usize);
13
14impl Cancel {
15    /// Check if this cancel token cancels the given key.
16    pub fn cancels<T: OpCode>(&self, key: &Key<T>) -> bool {
17        self.0 == key.as_raw()
18    }
19}
20
21impl CancelRegistry {
22    pub fn new() -> Self {
23        Self {
24            cancellers: HashSet::new(),
25        }
26    }
27
28    pub fn register<T>(&mut self, key: &Key<T>) -> Cancel {
29        let raw = key.as_raw();
30        if self.cancellers.contains(&raw) {
31            return Cancel(raw);
32        }
33        self.cancellers.insert(key.clone().erase());
34        Cancel(raw)
35    }
36
37    pub fn take(&mut self, token: Cancel) -> Option<ErasedKey> {
38        self.cancellers.take(&token.0)
39    }
40
41    pub fn remove(&mut self, key: &ErasedKey) -> bool {
42        self.cancellers.remove(key)
43    }
44}