tf_idf_vectorizer/utils/datastruct/arena/
mod.rs

1#[derive(Debug, Clone)]
2pub struct Arena<T> {
3    pool: Vec<Entry<T>>,
4    free_list: Vec<usize>,
5}
6
7#[derive(Debug, Clone)]
8pub enum Entry<T> {
9    None {
10        before_gen_id: u32,
11    },
12    Some {
13        value: T,
14        gen_id: u32,
15    },
16}
17
18impl<T> Entry<T> {
19    fn new_entry(&mut self, value: T) {
20        match self {
21            Entry::None { before_gen_id } => {
22                *self = Entry::Some {
23                    value,
24                    gen_id: *before_gen_id + 1,
25                };
26            }
27            Entry::Some { .. } => panic!("Inconsistent state: trying to allocate on occupied entry"),
28        }
29    }
30
31    fn new(value: T) -> Self {
32        Entry::Some { value, gen_id: 0 }
33    }
34
35    fn get_gen_id(&self) -> u32 {
36        match self {
37            Entry::None { before_gen_id } => *before_gen_id,
38            Entry::Some { gen_id, .. } => *gen_id,
39        }
40    }
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44pub struct ArenaSlot {
45    idx: u32,
46    gen_id: u32,
47}
48
49impl ArenaSlot {
50    fn new(index: usize, gen_id: u32) -> Self {
51        Self {
52            idx: index as u32,
53            gen_id,
54        }
55    }
56}
57
58impl<T> Arena<T> {
59    pub fn new() -> Self {
60        Self {
61            pool: Vec::new(),
62            free_list: Vec::new(),
63        }
64    }
65
66    pub fn alloc(&mut self, value: T) -> ArenaSlot {
67        if let Some(free_idx) = self.free_list.pop() {
68            let entry = &mut self.pool[free_idx];
69            entry.new_entry(value);
70            ArenaSlot::new(free_idx, entry.get_gen_id())
71        } else {
72            let entry = Entry::new(value);
73            self.pool.push(entry);
74            let idx = self.pool.len() - 1;
75            ArenaSlot::new(idx, self.pool[idx].get_gen_id())
76        }
77    }
78
79    pub fn dealloc(&mut self, slot: &ArenaSlot) {
80        let entry = &mut self.pool[slot.idx as usize];
81        match entry {
82            Entry::Some { gen_id, .. } if *gen_id == slot.gen_id => {
83                let gen_id = entry.get_gen_id();
84                *entry = Entry::None {
85                    before_gen_id: gen_id,
86                };
87                self.free_list.push(slot.idx as usize);
88            }
89            _ => {
90                // slot is invalid or already freed
91            }
92        }
93    }
94
95    pub fn get(&self, slot: &ArenaSlot) -> Option<&T> {
96        let entry = &self.pool[slot.idx as usize];
97        match entry {
98            Entry::Some { value, gen_id } if *gen_id == slot.gen_id => Some(value),
99            _ => None,
100        }
101    }
102
103    pub fn get_mut(&mut self, slot: &ArenaSlot) -> Option<&mut T> {
104        let entry = &mut self.pool[slot.idx as usize];
105        match entry {
106            Entry::Some { value, gen_id } if *gen_id == slot.gen_id => Some(value),
107            _ => None,
108        }
109    }
110}
tf_idf_vectorizer/utils/datastruct/arena/mod.rs

tf_idf_vectorizer/utils/datastruct/arena/
mod.rs
