mina_tree/port_ocaml/
hash_table.rs

1use std::collections::BTreeMap;
2
3use super::hash::OCamlHash;
4
5/// <https://github.com/janestreet/base/blob/v0.14/src/hashtbl.ml>
6pub struct HashTable<K, V> {
7    /// Note: OCaml uses AVL trees, but we just need an ordered map
8    table: Vec<BTreeMap<K, V>>,
9    length: usize,
10}
11
12impl<K: std::fmt::Debug, V: std::fmt::Debug> std::fmt::Debug for HashTable<K, V> {
13    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
14        let table: Vec<_> = self.table.iter().map(BTreeMap::len).collect();
15
16        f.debug_struct("HashTable")
17            .field("table", &table)
18            .field("length", &self.length)
19            .finish()
20    }
21}
22
23impl<K: Ord + OCamlHash + Clone, V: Clone> HashTable<K, V> {
24    pub fn create() -> Self {
25        let size = 1u64.next_power_of_two();
26
27        let table = (0..size).map(|_| BTreeMap::new()).collect();
28
29        Self { table, length: 0 }
30    }
31
32    fn slot(&self, key: &K) -> usize {
33        let hash = key.ocaml_hash();
34        let hash = hash as usize;
35        hash & (self.table.len() - 1)
36    }
37
38    fn find(&self, key: &K) -> Option<&V> {
39        let slot = self.slot(key);
40        let tree = &self.table[slot];
41
42        tree.get(key)
43    }
44
45    fn set_impl(&mut self, key: K, data: V) {
46        let slot = self.slot(&key);
47        let tree = &mut self.table[slot];
48
49        let old = tree.insert(key, data);
50        if old.is_none() {
51            self.length += 1;
52        }
53    }
54
55    fn set(&mut self, key: K, data: V) {
56        self.set_impl(key, data);
57        self.maybe_resize_table();
58    }
59
60    fn maybe_resize_table(&mut self) {
61        let len = self.table.len();
62        let should_grow = self.length > len;
63
64        if !should_grow {
65            return;
66        }
67
68        let new_array_length = len * 2;
69        if new_array_length <= len {
70            return;
71        }
72
73        let new_table = (0..new_array_length).map(|_| BTreeMap::new()).collect();
74        let old_table = std::mem::replace(&mut self.table, new_table);
75        self.length = 0;
76
77        for (key, value) in old_table.into_iter().flat_map(BTreeMap::into_iter) {
78            self.set_impl(key, value);
79        }
80    }
81
82    /// <https://github.com/janestreet/base/blob/v0.14/src/hashtbl.ml#L450>
83    pub fn update<F>(&mut self, key: K, fun: F)
84    where
85        F: FnOnce(Option<&V>) -> V,
86    {
87        let value = fun(self.find(&key));
88        self.set(key, value)
89    }
90}
91
92impl<K: Clone, V: Clone> HashTable<K, V> {
93    /// <https://github.com/janestreet/base/blob/v0.14/src/hashtbl.ml#L259-L281>
94    pub fn to_alist(&self) -> Vec<(K, V)> {
95        self.table
96            .iter()
97            .flat_map(BTreeMap::iter)
98            .map(|(key, value)| (key.clone(), value.clone()))
99            .rev() // rev because OCaml builds the list in reverse: `(key, data) :: list`
100            .collect()
101    }
102}