ironcore_search_helpers/
lib.rs

1use Result::{Err, Ok};
2use itertools::*;
3use lazy_static::*;
4use rand::RngExt;
5use rand::distr::Uniform;
6use rand::{CryptoRng, Rng};
7use sha2::digest::Update;
8use sha2::{Digest, Sha256};
9use std::collections::HashSet;
10use std::ops::DerefMut;
11use std::sync::{Mutex, MutexGuard};
12use unicode_segmentation::UnicodeSegmentation;
13use unidecode::unidecode_char;
14
15const FILTERED_CHARS: [char; 31] = [
16    '!', '@', '#', '$', '%', '^', '&', '*', '(', ')', '{', '}', '_', '<', '>', ':', ';', ',', '.',
17    '"', '\'', '`', '|', '+', '=', '/', '~', '[', ']', '\\', '-',
18];
19
20///True if we should keep the character in the string.
21fn should_keep_char(c: &char) -> bool {
22    !FILTERED_CHARS.contains(c)
23}
24lazy_static! {
25    ///Special chars that should be filtered out.
26    static ref ALL_U32: Uniform<u32> = Uniform::new_inclusive(0u32, u32::MAX).unwrap(); //safe because low < high
27    //We use this so we don't have to generate the floating numbers and do comparisons on them. It allows us to do 1/2 percent level scaling.
28    static ref ONE_TO_TWO_HUNDRED: Uniform<u8> = Uniform::new_inclusive(1, 200).unwrap(); // safe because low < high
29}
30
31///Something over 200 chars isn't really suitable for this approach, so we won't accept it.
32const MAX_STRING_LEN: usize = 200;
33
34/// Make an index, for the string s considering all tri-grams.
35/// The string will be latinised, lowercased and stripped of special chars before being broken into tri-grams.
36/// The values will be prefixed with partition_id and salt before being hashed.
37/// Each entry in the HashSet will be truncated to 32 bits and will be encoded as a big endian number.
38/// This function will also add some random entries to the HashSet to not expose how many tri-grams were actually found.
39pub fn generate_hashes_for_string_with_padding<R: Rng + CryptoRng>(
40    s: &str,
41    partition_id: Option<&str>,
42    salt: &[u8],
43    rng: &Mutex<R>,
44) -> Result<HashSet<u32>, String> {
45    let mut hashes = generate_hashes_for_string(s, partition_id, salt)?;
46
47    let prob = take_lock(rng).deref_mut().sample(*ONE_TO_TWO_HUNDRED);
48    let to_add: u8 = {
49        //Just take the lock once because we need it in all cases and it makes the code look better.
50        let r = &mut *take_lock(rng);
51        if prob <= 1 {
52            r.random_range(1..200)
53        } else if prob <= 5 {
54            r.random_range(1..30)
55        } else if prob <= 50 {
56            r.random_range(1..10)
57        } else {
58            r.random_range(1..5)
59        }
60    };
61    //This will never be negative because generate_hashes_for_string would error if hashes was going to be larger than and will never be larger than MAX_STRING_LEN.
62    //This also ensures we're able to pad by at least 2 since the maximum trigram length is always 2 less than the max string length.
63    let pad_len = std::cmp::min(MAX_STRING_LEN - hashes.len(), to_add as usize);
64    hashes.extend(
65        take_lock(rng)
66            .deref_mut()
67            .sample_iter(*ALL_U32)
68            .take(pad_len),
69    );
70    Ok(hashes)
71}
72
73/// Make an index, for the string s considering all tri-grams.
74/// The string will be latinised, lowercased and stripped of special chars before being broken into tri-grams.
75/// The values will be prefixed with partition_id and salt before being hashed.
76/// Each entry in the HasheSet will be truncated to 32 bits and will be encoded as a big endian number.
77/// If the string is longer than 200 characters, this will return an error.
78pub fn generate_hashes_for_string(
79    s: &str,
80    partition_id: Option<&str>,
81    salt: &[u8],
82) -> Result<HashSet<u32>, String> {
83    if s.len() > MAX_STRING_LEN {
84        Err(format!(
85            "The input string is too long. This function only supports strings that are no longer than {} chars.",
86            MAX_STRING_LEN
87        ))
88    } else {
89        //Compute a partial sha256 with the partition_id and the salt - We can reuse this for each word
90        let partial_sha256 = partition_id
91            .map(|k| k.as_bytes())
92            .iter()
93            .chain([salt].iter())
94            .fold(Sha256::new(), |hasher, k| hasher.chain(k));
95
96        let short_hash = |word: &[u8]| -> u32 {
97            let sha256_hash = partial_sha256.clone().chain(word);
98            as_u32_be(&sha256_hash.finalize().into())
99        };
100
101        let result: HashSet<_> = make_tri_grams(s)
102            .iter()
103            .map(|tri_gram| short_hash(tri_gram.as_bytes()))
104            .collect();
105        Ok(result)
106    }
107}
108
109/// Generate a version of the input string where each character has been latinized using the
110/// same function as our tokenization routines.
111pub fn transliterate_string(s: &str) -> String {
112    s.chars()
113        .filter(should_keep_char)
114        .map(char_to_trans)
115        .collect()
116}
117
118/// If s is empty, the resulting set will also be empty.
119/// If s is shorter than 3, '-' padding will be added to the end.
120/// All Strings inside of the resulting set will always be of size 3.
121fn make_tri_grams(s: &str) -> HashSet<String> {
122    let converted_string = transliterate_string(s);
123    converted_string
124        .unicode_words()
125        .map(|short_word| {
126            let short_word_len = short_word.chars().count();
127            if short_word_len < 3 {
128                //Pad the short_word with
129                format!("{:-<3}", short_word)
130            } else {
131                short_word.to_string()
132            }
133        })
134        .flat_map(|word| word_to_trigrams(&word))
135        .collect()
136}
137
138fn word_to_trigrams(s: &str) -> HashSet<String> {
139    s.chars()
140        .tuple_windows()
141        .map(|(c1, c2, c3)| format!("{}{}{}", c1, c2, c3))
142        .collect()
143}
144
145///Convert the char if we can, if we can't just create a string out of the character.
146fn char_to_trans(c: char) -> String {
147    let trans_string = unidecode_char(c);
148    if trans_string.is_empty() {
149        format!("{}", c)
150    } else {
151        trans_string.to_lowercase()
152    }
153}
154
155///Interpret the most significant 4 bytes as a bigendian u32
156#[inline]
157fn as_u32_be(slice: &[u8; 32]) -> u32 {
158    ((slice[0] as u32) << 24)
159        + ((slice[1] as u32) << 16)
160        + ((slice[2] as u32) << 8)
161        + (slice[3] as u32)
162}
163
164/// Acquire mutex in a blocking fashion. If the Mutex is or becomes poisoned, panic.
165///
166/// The lock is released when the returned MutexGuard falls out of scope.
167///
168/// # Usage:
169/// single statement (mut)
170/// `let result = take_lock(&t).deref_mut().call_method_on_t();`
171///
172/// multi-statement (mut)
173///
174/// ```ignore
175/// let t = T {};
176/// let result = {
177///     let g = &mut *take_lock(&t);
178///     g.call_method_on_t()
179/// }; // lock released here
180/// ```
181///
182fn take_lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
183    m.lock().unwrap_or_else(|e| {
184        let error = format!("Error when acquiring lock: {}", e);
185        panic!("{}", error);
186    })
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use rand::{distr::Alphanumeric, rngs::ThreadRng};
193
194    fn make_set(array: &[&str]) -> HashSet<String> {
195        array
196            .into_iter()
197            .map(|&s| From::from(s))
198            .collect::<HashSet<_>>()
199    }
200
201    #[test]
202    fn as_u32_be_known_result() {
203        let known_result = 16909060u32; //16777216 + 131072 + 768 + 4
204        let mut input = [0u8; 32];
205        input[0] = 1;
206        input[1] = 2;
207        input[2] = 3;
208        input[3] = 4;
209        let result = as_u32_be(&input);
210        assert_eq!(result, known_result);
211    }
212
213    #[test]
214    fn string_transliterated() {
215        assert_eq!(transliterate_string("Gumby, dammit!"), "gumby dammit");
216        assert_eq!(transliterate_string("北亰"), "bei jing ");
217        assert_eq!(transliterate_string("Æneid"), "aeneid");
218    }
219
220    #[test]
221    fn word_to_trigrams_known() {
222        let result = word_to_trigrams("five");
223        assert_eq!(result, make_set(&["fiv", "ive"]));
224    }
225
226    #[test]
227    fn make_tri_grams_works_multi_word() {
228        assert_eq!(
229            make_tri_grams("123 José  Núñez 812-111-7654"),
230            make_set(&[
231                "123", "jos", "ose", "nun", "une", "nez", "812", "121", "211", "111", "117", "176",
232                "765", "654",
233            ])
234        );
235    }
236
237    #[test]
238    fn make_tri_grams_works_non_ascii() {
239        assert_eq!(
240            make_tri_grams("TİRYAKİ"),
241            make_set(&["tir", "iry", "rya", "yak", "aki"])
242        );
243    }
244
245    #[test]
246    fn make_tri_grams_eliminates_duplicates() {
247        assert_eq!(
248            make_tri_grams("TİRYAKİ TİRYAKİ"),
249            make_set(&["tir", "iry", "rya", "yak", "aki"])
250        );
251    }
252
253    #[test]
254    fn make_tri_grams_works_short_non_ascii() {
255        assert_eq!(make_tri_grams("Tİ"), make_set(&["ti-"]));
256    }
257
258    #[test]
259    fn make_tri_grams_works_multichar_translate() {
260        assert_eq!(
261            make_tri_grams("志    豪 İ"),
262            make_set(&["zhi", "hao", "i--"])
263        );
264    }
265
266    #[test]
267    fn make_tri_grams_works_arabic() {
268        assert_eq!(
269            make_tri_grams("شريط فو"),
270            make_set(&["shr", "hry", "ryt", "fw-"])
271        );
272    }
273    #[test]
274    fn make_tri_grams_works_short_multibyte() {
275        assert_eq!(
276            make_tri_grams("\u{102AE}\u{102AF}"),
277            make_set(&["\u{102AE}\u{102AF}-"])
278        );
279    }
280
281    #[test]
282    fn char_to_trans_latinizable() {
283        assert_eq!(char_to_trans('İ'), "i")
284    }
285
286    #[test]
287    fn char_to_trans_not_latinizable() {
288        let c = "\u{102AE}".chars().nth(0).unwrap();
289        assert_eq!(char_to_trans(c), "\u{102AE}")
290    }
291    #[test]
292    fn generate_hashes_for_string_compute_known_value() -> Result<(), String> {
293        let result = generate_hashes_for_string("123", Some("foo"), &[0u8; 1])?;
294        //We compute this to catch cases where this computation might change.
295        let expected_result = {
296            let mut hasher = Sha256::new();
297            sha2::Digest::update(&mut hasher, "foo".as_bytes());
298            sha2::Digest::update(&mut hasher, [0u8; 1]);
299            sha2::Digest::update(&mut hasher, "123".as_bytes());
300            as_u32_be(&(hasher.finalize().into()))
301        };
302        assert_eq!(result, [expected_result].iter().map(|x| *x).collect());
303        Ok(())
304    }
305
306    #[test]
307    fn generate_hashes_for_string_with_padding_adds_at_least_one() -> Result<(), String> {
308        let rng = Mutex::new(ThreadRng::default());
309        let result = generate_hashes_for_string_with_padding("123", Some("foo"), &[0u8; 1], &rng)?;
310        assert!(result.len() > 1);
311        Ok(())
312    }
313
314    #[test]
315    fn generate_hashes_for_string_with_padding_empty_string() -> Result<(), String> {
316        let rng = Mutex::new(ThreadRng::default());
317        let result = generate_hashes_for_string_with_padding("", Some("foo"), &[0u8; 1], &rng)?;
318        assert!(result.len() >= 1);
319        Ok(())
320    }
321
322    #[test]
323    fn generate_hashes_for_string_too_long_errors() -> Result<(), String> {
324        let rng = ThreadRng::default();
325        let input: Vec<u8> = rng.sample_iter(Alphanumeric).take(201).collect();
326        generate_hashes_for_string(std::str::from_utf8(&input).unwrap(), Some("foo"), &[0u8; 1])
327            .unwrap_err();
328        Ok(())
329    }
330}
ironcore_search_helpers/lib.rs

ironcore_search_helpers/
lib.rs
