Skip to main content

surrealdb_common/ids/
id_set.rs

1use std::fmt::{self, Debug};
2use std::hash::{BuildHasher, Hash, RandomState};
3use std::ops::{Index, IndexMut};
4
5use hashbrown::HashTable;
6
7use super::Id;
8
9pub trait SetEntry<T>: Hash {
10	fn into_owned(self) -> T;
11
12	fn equal(&self, other: &T) -> bool;
13}
14
15impl<T: Hash + Eq> SetEntry<T> for T {
16	fn into_owned(self) -> T {
17		self
18	}
19
20	fn equal(&self, other: &T) -> bool {
21		self == other
22	}
23}
24
25impl SetEntry<String> for &str {
26	fn into_owned(self) -> String {
27		self.to_owned()
28	}
29
30	fn equal(&self, other: &String) -> bool {
31		self == other
32	}
33}
34
35/// A collection which will ensure that the storage only contains unique values.
36/// If two values are pushed which are equal to each-other this collection will instead return the
37/// id of the previous value.
38#[derive(Default)]
39pub struct IdSet<I, V, S = RandomState> {
40	map: HashTable<I>,
41	storage: Vec<V>,
42	hasher: S,
43}
44
45impl<I, V> IdSet<I, V>
46where
47	I: Id,
48	V: Eq + Hash,
49{
50	pub fn new() -> Self {
51		IdSet {
52			map: HashTable::new(),
53			storage: Vec::new(),
54			hasher: RandomState::new(),
55		}
56	}
57}
58
59impl<I, V> Debug for IdSet<I, V>
60where
61	I: Id + Debug,
62	V: Eq + Hash + Debug,
63{
64	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65		let mut fmt = f.debug_list();
66		for idx in self.map.iter() {
67			fmt.entry(&self.storage[idx.idx()]);
68		}
69		fmt.finish()
70	}
71}
72
73impl<I, V, S> IdSet<I, V, S>
74where
75	I: Id,
76	V: Eq + Hash,
77	S: BuildHasher,
78{
79	pub fn push<T>(&mut self, v: T) -> Option<I>
80	where
81		T: SetEntry<V>,
82	{
83		let hash = self.hasher.hash_one(&v);
84		match self.map.find_entry(hash, |s| v.equal(&self.storage[s.idx()])) {
85			Ok(x) => Some(*x.get()),
86			Err(slot) => {
87				let idx = I::from_idx(self.storage.len())?;
88				self.storage.push(v.into_owned());
89				slot.into_table()
90					.insert_unique(hash, idx, |x| self.hasher.hash_one(&self.storage[x.idx()]));
91				Some(idx)
92			}
93		}
94	}
95
96	pub fn clear(&mut self) {
97		self.map.clear();
98		self.storage.clear();
99	}
100
101	pub fn len(&self) -> usize {
102		self.storage.len()
103	}
104
105	pub fn is_empty(&self) -> bool {
106		self.storage.is_empty()
107	}
108}
109
110impl<I: Id, V, S> IdSet<I, V, S> {
111	pub fn get(&self, index: I) -> Option<&V> {
112		self.storage.get(index.idx())
113	}
114
115	pub fn get_mut(&mut self, index: I) -> Option<&mut V> {
116		self.storage.get_mut(index.idx())
117	}
118}
119
120impl<I, V, S> Index<I> for IdSet<I, V, S>
121where
122	I: Id,
123	V: Eq + Hash,
124	S: BuildHasher,
125{
126	type Output = V;
127
128	fn index(&self, index: I) -> &Self::Output {
129		let Some(x) = self.get(index) else {
130			panic!("Tired to index into id set with out of range index {}", index.idx())
131		};
132		x
133	}
134}
135
136impl<I, V, S> IndexMut<I> for IdSet<I, V, S>
137where
138	I: Id,
139	V: Eq + Hash,
140	S: BuildHasher,
141{
142	#[track_caller]
143	fn index_mut(&mut self, index: I) -> &mut Self::Output {
144		let Some(x) = self.get_mut(index) else {
145			panic!("Tired to index into id set with out of range index {}", index.idx())
146		};
147		x
148	}
149}