hpt-allocator 0.1.2

An internal library for memory allocator for hpt
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

pub mod cpu;
#[cfg(feature = "cuda")]
pub mod cuda;

use std::collections::HashMap;

use crate::ptr::SafePtr;

pub trait Storage {
    fn increment_ref(&mut self, ptr: SafePtr);
    fn decrement_ref(&mut self, ptr: SafePtr) -> bool;
}

#[derive(Debug)]
pub struct CommonStorage {
    pub(crate) storage: HashMap<SafePtr, usize>,
}

impl CommonStorage {
    pub fn new() -> Self {
        CommonStorage {
            storage: HashMap::new(),
        }
    }
}

impl Storage for CommonStorage {
    fn increment_ref(&mut self, ptr: SafePtr) {
        if let Some(cnt) = self.storage.get_mut(&ptr) {
            *cnt = match cnt.checked_add(1) {
                Some(cnt) => cnt,
                None => {
                    panic!(
                        "Reference count overflow for ptr {:p} in cpu storage",
                        ptr.ptr
                    );
                }
            };
        } else {
            self.storage.insert(ptr, 1);
        }
    }

    fn decrement_ref(&mut self, ptr: SafePtr) -> bool {
        if let Some(cnt) = self.storage.get_mut(&ptr) {
            *cnt = cnt.checked_sub(1).expect("Reference count underflow");
            if *cnt == 0 {
                self.storage.remove(&ptr);
                true
            } else {
                false
            }
        } else {
            false
        }
    }
}

/// # Clone Storage
///
/// increment the reference count of the ptr in the storage
pub fn clone_storage(ptr: *mut u8, device_id: usize, map: &mut HashMap<usize, CommonStorage>) {
    if let Some(storage) = map.get_mut(&device_id) {
        storage.increment_ref(SafePtr { ptr });
    } else {
        map.insert(device_id, CommonStorage::new());
        map.get_mut(&device_id)
            .unwrap()
            .increment_ref(SafePtr { ptr });
    }
}