hpt_allocator/storage/
cpu.rs

1use std::sync::Mutex;
2
3use hashbrown::HashMap;
4use once_cell::sync::Lazy;
5
6use super::{SafePtr, Storage};
7
8/// This is a global variable that stores the allocated ptrs and their reference count
9pub static CPU_STORAGE: Lazy<Mutex<CpuStorage>> = Lazy::new(|| Mutex::new(CpuStorage::new()));
10
11#[derive(Debug)]
12pub struct CpuStorage {
13    storage: HashMap<SafePtr, usize>,
14}
15
16impl CpuStorage {
17    pub fn new() -> Self {
18        CpuStorage {
19            storage: HashMap::new(),
20        }
21    }
22}
23
24impl Storage for CpuStorage {
25    fn increment_ref(&mut self, ptr: SafePtr) {
26        if let Some(cnt) = self.storage.get_mut(&ptr) {
27            *cnt = match cnt.checked_add(1) {
28                Some(cnt) => cnt,
29                None => {
30                    panic!(
31                        "Reference count overflow for ptr {:p} in cpu storage",
32                        ptr.ptr
33                    );
34                }
35            };
36        } else {
37            self.storage.insert(ptr, 1);
38        }
39    }
40
41    fn decrement_ref(&mut self, ptr: SafePtr) -> bool {
42        if let Some(cnt) = self.storage.get_mut(&ptr) {
43            *cnt = cnt.checked_sub(1).expect("Reference count underflow");
44            if *cnt == 0 {
45                self.storage.remove(&ptr);
46                true
47            } else {
48                false
49            }
50        } else {
51            panic!("ptr {:p} not found in cpu storage", ptr.ptr);
52        }
53    }
54}
55
56/// # Clone Storage
57///
58/// increment the reference count of the ptr in the storage
59pub fn clone_storage(ptr: *mut u8) {
60    if let Ok(mut storage) = CPU_STORAGE.lock() {
61        storage.increment_ref(SafePtr { ptr });
62    }
63}