ixa/people/
index.rs

1use super::methods::Methods;
2use crate::people::external_api::ContextPeopleExtCrate;
3use crate::{Context, ContextPeopleExt, PersonId, PersonProperty};
4use crate::{HashMap, HashMapExt, HashSet, HashSetExt};
5use bincode::serialize;
6use serde::ser::{Serialize, Serializer};
7use std::any::TypeId;
8use std::cell::RefCell;
9use std::collections::hash_map::DefaultHasher;
10use std::hash::{Hash, Hasher};
11use std::sync::Arc;
12use std::sync::LazyLock;
13use std::sync::Mutex;
14
15#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
16// The lookup key for entries in the index. This is a serialized version of the value.
17// If that serialization fits in 128 bits, we store it in `IndexValue::Fixed`
18// Otherwise, we use a 64 bit hash.
19#[doc(hidden)]
20pub enum IndexValue {
21    Fixed(u128),
22    Hashed(u64),
23}
24
25impl Serialize for IndexValue {
26    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
27    where
28        S: Serializer,
29    {
30        match self {
31            IndexValue::Fixed(value) => serializer.serialize_u128(*value),
32            IndexValue::Hashed(value) => serializer.serialize_u64(*value),
33        }
34    }
35}
36
37impl IndexValue {
38    pub fn compute<T: Serialize>(val: &T) -> IndexValue {
39        // Serialize `val` to a `Vec<u8>` using `bincode`
40        let serialized_data = serialize(val).expect("Failed to serialize value");
41
42        // If serialized data fits within 16 bytes...
43        if serialized_data.len() <= 16 {
44            // ...store it as `IndexValue::Fixed`
45            let mut tmp: [u8; 16] = [0; 16];
46            tmp[..serialized_data.len()].copy_from_slice(&serialized_data[..]);
47
48            IndexValue::Fixed(u128::from_le_bytes(tmp))
49        } else {
50            // Otherwise, hash the data and store it as `IndexValue::Hashed`
51            let mut hasher = DefaultHasher::new();
52            serialized_data.hash(&mut hasher);
53            let hash = hasher.finish(); // Produces a 64-bit hash
54            IndexValue::Hashed(hash)
55        }
56    }
57}
58
59// An index for a single property.
60pub struct Index {
61    // Primarily for debugging purposes
62    #[allow(dead_code)]
63    pub(super) name: &'static str,
64
65    // The hash of the property value maps to a list of PersonIds
66    // or None if we're not indexing
67    pub(super) lookup: Option<HashMap<IndexValue, (String, HashSet<PersonId>)>>,
68
69    // The largest person ID that has been indexed. Used so that we
70    // can lazily index when a person is added.
71    pub(super) max_indexed: usize,
72}
73
74impl Index {
75    pub(super) fn new<T: PersonProperty>(_context: &Context, _property: T) -> Self {
76        Self {
77            name: std::any::type_name::<T>(),
78            lookup: None,
79            max_indexed: 0,
80        }
81    }
82
83    pub(super) fn add_person(&mut self, context: &Context, methods: &Methods, person_id: PersonId) {
84        let hash = (methods.indexer)(context, person_id);
85        self.lookup
86            .as_mut()
87            .unwrap()
88            .entry(hash)
89            .or_insert_with(|| ((methods.get_display)(context, person_id), HashSet::new()))
90            .1
91            .insert(person_id);
92    }
93
94    pub(super) fn remove_person(
95        &mut self,
96        context: &Context,
97        methods: &Methods,
98        person_id: PersonId,
99    ) {
100        let hash = (methods.indexer)(context, person_id);
101        if let Some(entry) = self.lookup.as_mut().unwrap().get_mut(&hash) {
102            entry.1.remove(&person_id);
103            // Clean up the entry if there are no people
104            if entry.0.is_empty() {
105                self.lookup.as_mut().unwrap().remove(&hash);
106            }
107        }
108    }
109
110    pub(super) fn index_unindexed_people(&mut self, context: &Context, methods: &Methods) {
111        if self.lookup.is_none() {
112            return;
113        }
114        let current_pop = context.get_current_population();
115        for id in self.max_indexed..current_pop {
116            let person_id = PersonId(id);
117            self.add_person(context, methods, person_id);
118        }
119        self.max_indexed = current_pop;
120    }
121}
122
123/// The common callback used by multiple `Context` methods for future events
124type IndexCallback = dyn Fn(&Context) + Send + Sync;
125
126pub struct MultiIndex {
127    register: Box<IndexCallback>,
128    type_id: TypeId,
129}
130
131// A static map of multi-property indices. This is used to register multi-property property indices
132// when they are first created. The map is keyed by the property IDs of the properties that are
133// indexed. The values are the `MultiIndex` objects that contain the registration function and
134// the type ID of the index.
135#[doc(hidden)]
136#[allow(clippy::type_complexity)]
137pub static MULTI_PROPERTY_INDEX_MAP: LazyLock<
138    Mutex<RefCell<HashMap<Vec<TypeId>, Arc<MultiIndex>>>>,
139> = LazyLock::new(|| Mutex::new(RefCell::new(HashMap::new())));
140
141#[allow(dead_code)]
142pub fn add_multi_property_index<T: PersonProperty>(property_ids: &[TypeId], index_type: TypeId) {
143    let current_map = MULTI_PROPERTY_INDEX_MAP.lock().unwrap();
144    let mut map = current_map.borrow_mut();
145    let mut ordered_property_ids = property_ids.to_owned();
146    ordered_property_ids.sort();
147    map.entry(ordered_property_ids)
148        .or_insert(Arc::new(MultiIndex {
149            register: Box::new(|context| {
150                context.register_property::<T>();
151                context.index_property_by_id(TypeId::of::<T>());
152            }),
153            type_id: index_type,
154        }));
155}
156
157pub fn get_and_register_multi_property_index(
158    query: &[(TypeId, IndexValue)],
159    context: &Context,
160) -> Option<TypeId> {
161    let map = MULTI_PROPERTY_INDEX_MAP.lock().unwrap();
162    let map = map.borrow();
163    let mut sorted_query = query.to_owned();
164    sorted_query.sort_by(|a, b| a.0.cmp(&b.0));
165    let items = query.iter().map(|(t, _)| *t).collect::<Vec<_>>();
166    if let Some(multi_index) = map.get(&items) {
167        (multi_index.register)(context);
168        return Some(multi_index.type_id);
169    }
170    None
171}
172
173pub fn get_multi_property_hash(query: &[(TypeId, IndexValue)]) -> IndexValue {
174    let mut sorted_query = query.to_owned();
175    sorted_query.sort_by(|a, b| a.0.cmp(&b.0));
176    let items = query.iter().map(|(_, i)| *i).collect::<Vec<_>>();
177    IndexValue::compute(&IndexValue::compute(&items))
178}
179
180pub fn process_indices(
181    context: &Context,
182    remaining_indices: &[&Index],
183    property_names: &mut Vec<String>,
184    current_matches: &HashSet<PersonId>,
185    print_fn: &dyn Fn(&Context, &[String], usize),
186) {
187    if remaining_indices.is_empty() {
188        print_fn(context, property_names, current_matches.len());
189        return;
190    }
191
192    let (next_index, rest_indices) = remaining_indices.split_first().unwrap();
193    let lookup = next_index.lookup.as_ref().unwrap();
194
195    // If there is nothing in the index, we don't need to process it
196    if lookup.is_empty() {
197        return;
198    }
199
200    for (display, people) in lookup.values() {
201        let intersect = !property_names.is_empty();
202        property_names.push(display.clone());
203
204        let matches = if intersect {
205            &current_matches.intersection(people).copied().collect()
206        } else {
207            people
208        };
209
210        process_indices(context, rest_indices, property_names, matches, print_fn);
211        property_names.pop();
212    }
213}
214
215#[cfg(test)]
216mod test {
217    // Tests in `src/people/query.rs` also exercise indexing code.
218
219    use crate::people::index::{Index, IndexValue};
220    use crate::{define_person_property, Context, ContextPeopleExt};
221    use std::any::TypeId;
222
223    define_person_property!(Age, u8);
224
225    #[test]
226    fn index_name() {
227        let context = Context::new();
228        let index = Index::new(&context, Age);
229        assert!(index.name.contains("Age"));
230    }
231
232    #[test]
233    fn test_index_value_hasher_finish2_short() {
234        let value = 42;
235        let index = IndexValue::compute(&value);
236        assert!(matches!(index, IndexValue::Fixed(_)));
237    }
238
239    #[test]
240    fn test_index_value_hasher_finish2_long() {
241        let value = "this is a longer string that exceeds 16 bytes";
242        let index = IndexValue::compute(&value);
243        assert!(matches!(index, IndexValue::Hashed(_)));
244    }
245
246    #[test]
247    fn test_index_value_compute_same_values() {
248        let value = "test value";
249        let value2 = "test value";
250        assert_eq!(IndexValue::compute(&value), IndexValue::compute(&value2));
251    }
252
253    #[test]
254    fn test_index_value_compute_different_values() {
255        let value1 = 42;
256        let value2 = 43;
257        assert_ne!(IndexValue::compute(&value1), IndexValue::compute(&value2));
258    }
259
260    #[test]
261    fn test_multi_property_index_map_add_get_and_hash() {
262        let mut context = Context::new();
263        let _person = context.add_person((Age, 64)).unwrap();
264        let property_ids = vec![TypeId::of::<Age>()];
265        let index_type = TypeId::of::<IndexValue>();
266        super::add_multi_property_index::<Age>(&property_ids, index_type);
267        let query = vec![(TypeId::of::<Age>(), IndexValue::Fixed(42))];
268        let _ = super::get_multi_property_hash(&query);
269        let registered_index = super::get_and_register_multi_property_index(&query, &context);
270        assert!(registered_index.is_some());
271    }
272}