1use std::collections::HashSet;
2
3use crate::fingerprint::{SimHash64, DEFAULT_NEAR_DUP_DISTANCE};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6pub struct IndexConfig {
7 pub max_distance: u32,
8 pub tables: Vec<TableSpec>,
9}
10
11impl IndexConfig {
12 pub fn google64_k3() -> Self {
13 Self::from_block_widths(&[11, 11, 11, 11, 10, 10], 3, DEFAULT_NEAR_DUP_DISTANCE)
14 }
15
16 pub fn from_block_widths(
17 block_widths: &[u8],
18 matching_blocks: usize,
19 max_distance: u32,
20 ) -> Self {
21 assert_eq!(
22 block_widths
23 .iter()
24 .map(|width| u16::from(*width))
25 .sum::<u16>(),
26 64
27 );
28 assert!(matching_blocks > 0);
29 assert!(matching_blocks <= block_widths.len());
30
31 let blocks = block_positions(block_widths);
32 let mut tables = Vec::new();
33 let mut chosen = Vec::new();
34 combinations(
35 block_widths.len(),
36 matching_blocks,
37 0,
38 &mut chosen,
39 &mut |combo| {
40 tables.push(TableSpec::from_selected_blocks(&blocks, combo));
41 },
42 );
43
44 Self {
45 max_distance,
46 tables,
47 }
48 }
49}
50
51#[derive(Debug, Clone, PartialEq, Eq)]
52pub struct TableSpec {
53 pub selected_blocks: Vec<usize>,
54 pub permutation: [u8; 64],
55 pub prefix_bits: u8,
56}
57
58impl TableSpec {
59 fn from_selected_blocks(blocks: &[Vec<u8>], selected_blocks: &[usize]) -> Self {
60 let selected: HashSet<_> = selected_blocks.iter().copied().collect();
61 let mut order = Vec::with_capacity(64);
62
63 for block in selected_blocks {
64 order.extend(blocks[*block].iter().copied());
65 }
66 for (idx, block) in blocks.iter().enumerate() {
67 if !selected.contains(&idx) {
68 order.extend(block.iter().copied());
69 }
70 }
71
72 let mut permutation = [0u8; 64];
73 permutation.copy_from_slice(&order);
74 let prefix_bits = selected_blocks
75 .iter()
76 .map(|block| blocks[*block].len() as u8)
77 .sum();
78
79 Self {
80 selected_blocks: selected_blocks.to_vec(),
81 permutation,
82 prefix_bits,
83 }
84 }
85
86 pub fn permute(&self, hash: SimHash64) -> u64 {
87 let mut value = 0u64;
88 for source_bit in self.permutation {
89 value <<= 1;
90 value |= (hash.0 >> source_bit) & 1;
91 }
92 value
93 }
94
95 pub fn prefix(&self, hash: SimHash64) -> u64 {
96 leading_prefix(self.permute(hash), self.prefix_bits)
97 }
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
101pub struct Match {
102 pub id: String,
103 pub fingerprint: SimHash64,
104 pub distance: u32,
105}
106
107#[derive(Debug, Clone)]
108struct IndexedItem {
109 id: String,
110 fingerprint: SimHash64,
111}
112
113#[derive(Debug, Clone)]
114struct TableEntry {
115 permuted: u64,
116 item_index: usize,
117}
118
119#[derive(Debug, Clone)]
120struct Table {
121 spec: TableSpec,
122 entries: Vec<TableEntry>,
123}
124
125#[derive(Debug, Clone)]
126pub struct SimHashIndex {
127 config: IndexConfig,
128 items: Vec<IndexedItem>,
129 tables: Vec<Table>,
130}
131
132impl SimHashIndex {
133 pub fn new(config: IndexConfig) -> Self {
134 let tables = config
135 .tables
136 .iter()
137 .cloned()
138 .map(|spec| Table {
139 spec,
140 entries: Vec::new(),
141 })
142 .collect();
143
144 Self {
145 config,
146 items: Vec::new(),
147 tables,
148 }
149 }
150
151 pub fn new_google64_k3() -> Self {
152 Self::new(IndexConfig::google64_k3())
153 }
154
155 pub fn from_items<I, S>(config: IndexConfig, items: I) -> Self
156 where
157 I: IntoIterator<Item = (S, SimHash64)>,
158 S: Into<String>,
159 {
160 let mut index = Self::new(config);
161 for (id, fingerprint) in items {
162 index.insert(id, fingerprint);
163 }
164 index
165 }
166
167 pub fn config(&self) -> &IndexConfig {
168 &self.config
169 }
170
171 pub fn len(&self) -> usize {
172 self.items.len()
173 }
174
175 pub fn is_empty(&self) -> bool {
176 self.items.is_empty()
177 }
178
179 pub fn insert<S: Into<String>>(&mut self, id: S, fingerprint: SimHash64) {
180 let item_index = self.items.len();
181 self.items.push(IndexedItem {
182 id: id.into(),
183 fingerprint,
184 });
185
186 for table in &mut self.tables {
187 table.entries.push(TableEntry {
188 permuted: table.spec.permute(fingerprint),
189 item_index,
190 });
191 table
192 .entries
193 .sort_by_key(|entry| (entry.permuted, entry.item_index));
194 }
195 }
196
197 pub fn query(&self, fingerprint: SimHash64) -> Vec<Match> {
198 let mut candidates = HashSet::new();
199
200 for table in &self.tables {
201 let permuted = table.spec.permute(fingerprint);
202 let (start, end) = prefix_range(permuted, table.spec.prefix_bits);
203 let start_idx = lower_bound(&table.entries, start);
204 let end_idx = upper_bound(&table.entries, end);
205
206 for entry in &table.entries[start_idx..end_idx] {
207 candidates.insert(entry.item_index);
208 }
209 }
210
211 let mut matches = Vec::new();
212 for item_index in candidates {
213 let item = &self.items[item_index];
214 let distance = fingerprint.hamming_distance(item.fingerprint);
215 if distance <= self.config.max_distance {
216 matches.push(Match {
217 id: item.id.clone(),
218 fingerprint: item.fingerprint,
219 distance,
220 });
221 }
222 }
223
224 matches.sort_by(|a, b| a.distance.cmp(&b.distance).then_with(|| a.id.cmp(&b.id)));
225 matches
226 }
227}
228
229fn block_positions(widths: &[u8]) -> Vec<Vec<u8>> {
230 let mut blocks = Vec::new();
231 let mut next_msb = 63i16;
232
233 for width in widths {
234 let mut block = Vec::new();
235 for _ in 0..*width {
236 block.push(next_msb as u8);
237 next_msb -= 1;
238 }
239 blocks.push(block);
240 }
241
242 blocks
243}
244
245fn combinations<F>(n: usize, k: usize, start: usize, chosen: &mut Vec<usize>, on_combo: &mut F)
246where
247 F: FnMut(&[usize]),
248{
249 if chosen.len() == k {
250 on_combo(chosen);
251 return;
252 }
253
254 for idx in start..n {
255 chosen.push(idx);
256 combinations(n, k, idx + 1, chosen, on_combo);
257 chosen.pop();
258 }
259}
260
261fn leading_prefix(value: u64, prefix_bits: u8) -> u64 {
262 if prefix_bits == 0 {
263 0
264 } else {
265 value >> (64 - prefix_bits)
266 }
267}
268
269fn prefix_range(permuted: u64, prefix_bits: u8) -> (u64, u64) {
270 if prefix_bits == 0 {
271 return (0, u64::MAX);
272 }
273 if prefix_bits == 64 {
274 return (permuted, permuted);
275 }
276
277 let prefix = leading_prefix(permuted, prefix_bits);
278 let suffix_bits = 64 - prefix_bits;
279 let start = prefix << suffix_bits;
280 let end = start | ((1u64 << suffix_bits) - 1);
281 (start, end)
282}
283
284fn lower_bound(entries: &[TableEntry], value: u64) -> usize {
285 let mut left = 0usize;
286 let mut right = entries.len();
287 while left < right {
288 let mid = left + (right - left) / 2;
289 if entries[mid].permuted < value {
290 left = mid + 1;
291 } else {
292 right = mid;
293 }
294 }
295 left
296}
297
298fn upper_bound(entries: &[TableEntry], value: u64) -> usize {
299 let mut left = 0usize;
300 let mut right = entries.len();
301 while left < right {
302 let mid = left + (right - left) / 2;
303 if entries[mid].permuted <= value {
304 left = mid + 1;
305 } else {
306 right = mid;
307 }
308 }
309 left
310}
311
312#[cfg(test)]
313mod tests {
314 use super::*;
315
316 #[test]
317 fn google_config_has_twenty_tables() {
318 let config = IndexConfig::google64_k3();
319 assert_eq!(config.max_distance, 3);
320 assert_eq!(config.tables.len(), 20);
321 assert!(config
322 .tables
323 .iter()
324 .all(|table| matches!(table.prefix_bits, 31 | 32 | 33)));
325 }
326
327 #[test]
328 fn permutation_moves_selected_block_to_prefix() {
329 let config = IndexConfig::from_block_widths(&[32, 32], 1, 1);
330 let first = &config.tables[0];
331 let second = &config.tables[1];
332 let hash = SimHash64(0xaaaa_aaaa_5555_5555);
333
334 assert_eq!(first.prefix(hash), 0xaaaa_aaaa);
335 assert_eq!(second.prefix(hash), 0x5555_5555);
336 }
337
338 #[test]
339 fn query_finds_items_within_threshold() {
340 let base = SimHash64(0x1234_5678_90ab_cdef);
341 let within = SimHash64(base.0 ^ 0b111);
342 let outside = SimHash64(base.0 ^ 0b1_1111);
343 let index = SimHashIndex::from_items(
344 IndexConfig::google64_k3(),
345 [("within", within), ("outside", outside)],
346 );
347
348 let matches = index.query(base);
349 assert_eq!(matches.len(), 1);
350 assert_eq!(matches[0].id, "within");
351 assert_eq!(matches[0].distance, 3);
352 }
353
354 #[test]
355 fn query_results_are_sorted_by_distance() {
356 let base = SimHash64(0xffff_0000_ffff_0000);
357 let one = SimHash64(base.0 ^ 0b1);
358 let three = SimHash64(base.0 ^ 0b111);
359 let index =
360 SimHashIndex::from_items(IndexConfig::google64_k3(), [("three", three), ("one", one)]);
361
362 let matches = index.query(base);
363 assert_eq!(matches[0].id, "one");
364 assert_eq!(matches[1].id, "three");
365 }
366
367 #[test]
368 fn empty_index_returns_no_matches() {
369 let index = SimHashIndex::new_google64_k3();
370 assert!(index.is_empty());
371 assert!(index.query(SimHash64(42)).is_empty());
372 }
373}