icydb_core/db/store/
index.rs

1use crate::{
2    MAX_INDEX_FIELDS,
3    db::store::{DataKey, EntityName, IndexName, StoreRegistry},
4    prelude::*,
5    traits::Storable,
6};
7use candid::CandidType;
8use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, memory::VirtualMemory, storable::Bound};
9use canic_memory::impl_storable_unbounded;
10use derive_more::{Deref, DerefMut, Display};
11use serde::{Deserialize, Serialize};
12use std::{borrow::Cow, collections::HashSet};
13
14///
15/// IndexStoreRegistry
16///
17
18#[derive(Deref, DerefMut)]
19pub struct IndexStoreRegistry(StoreRegistry<IndexStore>);
20
21impl IndexStoreRegistry {
22    #[must_use]
23    pub fn new() -> Self {
24        Self(StoreRegistry::new())
25    }
26}
27
28impl Default for IndexStoreRegistry {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34///
35/// IndexInsertOutcome
36///
37
38#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum IndexInsertOutcome {
40    Inserted,
41    Skipped,
42}
43
44///
45/// IndexRemoveOutcome
46///
47
48#[derive(Clone, Copy, Debug, Eq, PartialEq)]
49pub enum IndexRemoveOutcome {
50    Removed,
51    Skipped,
52}
53
54///
55/// IndexInsertError
56///
57
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
59pub enum IndexInsertError {
60    UniqueViolation,
61}
62
63///
64/// IndexStore
65///
66
67#[derive(Deref, DerefMut)]
68pub struct IndexStore(BTreeMap<IndexKey, IndexEntry, VirtualMemory<DefaultMemoryImpl>>);
69
70impl IndexStore {
71    #[must_use]
72    pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
73        Self(BTreeMap::init(memory))
74    }
75
76    pub fn insert_index_entry<E: EntityKind>(
77        &mut self,
78        entity: &E,
79        index: &IndexModel,
80    ) -> Result<IndexInsertOutcome, IndexInsertError> {
81        let Some(index_key) = IndexKey::new(entity, index) else {
82            return Ok(IndexInsertOutcome::Skipped);
83        };
84        let key = entity.key();
85
86        if let Some(mut entry) = self.get(&index_key) {
87            if index.unique {
88                if entry.contains(&key) {
89                    return Ok(IndexInsertOutcome::Skipped);
90                }
91                if !entry.is_empty() {
92                    return Err(IndexInsertError::UniqueViolation);
93                }
94                self.insert(index_key, IndexEntry::new(key));
95            } else {
96                entry.insert_key(key);
97                self.insert(index_key, entry);
98            }
99        } else {
100            self.insert(index_key, IndexEntry::new(key));
101        }
102
103        Ok(IndexInsertOutcome::Inserted)
104    }
105
106    pub fn remove_index_entry<E: EntityKind>(
107        &mut self,
108        entity: &E,
109        index: &IndexModel,
110    ) -> IndexRemoveOutcome {
111        let Some(index_key) = IndexKey::new(entity, index) else {
112            return IndexRemoveOutcome::Skipped;
113        };
114
115        if let Some(mut entry) = self.get(&index_key) {
116            entry.remove_key(&entity.key());
117            if entry.is_empty() {
118                self.remove(&index_key);
119            } else {
120                self.insert(index_key, entry);
121            }
122            return IndexRemoveOutcome::Removed;
123        }
124
125        IndexRemoveOutcome::Skipped
126    }
127
128    pub fn resolve_data_values<E: EntityKind>(
129        &self,
130        index: &IndexModel,
131        prefix: &[Value],
132    ) -> Vec<DataKey> {
133        let mut out = Vec::new();
134
135        for (_, entry) in self.iter_with_prefix::<E>(index, prefix) {
136            out.extend(entry.keys.iter().map(|&k| DataKey::new::<E>(k)));
137        }
138
139        out
140    }
141
142    fn iter_with_prefix<E: EntityKind>(
143        &self,
144        index: &IndexModel,
145        prefix: &[Value],
146    ) -> Box<dyn Iterator<Item = (IndexKey, IndexEntry)> + '_> {
147        let index_id = IndexId::new::<E>(index);
148
149        let Some(fps) = prefix
150            .iter()
151            .map(Value::to_index_fingerprint)
152            .collect::<Option<Vec<_>>>()
153        else {
154            let empty = IndexKey::empty(index_id);
155            return Box::new(
156                self.range(empty.clone()..empty)
157                    .map(|e| (e.key().clone(), e.value())),
158            );
159        };
160
161        let (start, end) = IndexKey::bounds_for_prefix(index_id, index.fields.len(), &fps);
162
163        Box::new(
164            self.range(start..=end)
165                .map(|e| (e.key().clone(), e.value())),
166        )
167    }
168
169    /// Sum of bytes used by all index entries.
170    pub fn memory_bytes(&self) -> u64 {
171        self.iter()
172            .map(|entry| u64::from(IndexKey::STORED_SIZE) + entry.value().to_bytes().len() as u64)
173            .sum()
174    }
175}
176
177///
178/// IndexId
179///
180
181#[derive(Clone, Copy, Debug, Display, Eq, Hash, Ord, PartialEq, PartialOrd)]
182pub struct IndexId(pub IndexName);
183
184impl IndexId {
185    #[must_use]
186    pub fn new<E: EntityKind>(index: &IndexModel) -> Self {
187        let entity = EntityName::from_static(E::ENTITY_NAME);
188        Self(IndexName::from_parts(&entity, index.fields))
189    }
190
191    #[must_use]
192    pub const fn max_storable() -> Self {
193        Self(IndexName::max_storable())
194    }
195}
196
197///
198/// IndexKey
199/// (FIXED-SIZE, MANUAL)
200///
201
202#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
203pub struct IndexKey {
204    index_id: IndexId,
205    len: u8,
206    values: [[u8; 16]; MAX_INDEX_FIELDS],
207}
208
209impl IndexKey {
210    #[allow(clippy::cast_possible_truncation)]
211    pub const STORED_SIZE: u32 = IndexName::STORED_SIZE + 1 + (MAX_INDEX_FIELDS as u32 * 16);
212
213    pub fn new<E: EntityKind>(entity: &E, index: &IndexModel) -> Option<Self> {
214        if index.fields.len() > MAX_INDEX_FIELDS {
215            return None;
216        }
217
218        let mut values = [[0u8; 16]; MAX_INDEX_FIELDS];
219        let mut len = 0usize;
220
221        for field in index.fields {
222            let value = entity.get_value(field)?;
223            let fp = value.to_index_fingerprint()?;
224            values[len] = fp;
225            len += 1;
226        }
227
228        #[allow(clippy::cast_possible_truncation)]
229        Some(Self {
230            index_id: IndexId::new::<E>(index),
231            len: len as u8,
232            values,
233        })
234    }
235
236    #[must_use]
237    pub const fn empty(index_id: IndexId) -> Self {
238        Self {
239            index_id,
240            len: 0,
241            values: [[0u8; 16]; MAX_INDEX_FIELDS],
242        }
243    }
244
245    #[must_use]
246    #[allow(clippy::cast_possible_truncation)]
247    pub fn bounds_for_prefix(
248        index_id: IndexId,
249        index_len: usize,
250        prefix: &[[u8; 16]],
251    ) -> (Self, Self) {
252        let mut start = Self::empty(index_id);
253        let mut end = Self::empty(index_id);
254
255        for (i, fp) in prefix.iter().enumerate() {
256            start.values[i] = *fp;
257            end.values[i] = *fp;
258        }
259
260        start.len = index_len as u8;
261        end.len = start.len;
262
263        for value in end.values.iter_mut().take(index_len).skip(prefix.len()) {
264            *value = [0xFF; 16];
265        }
266
267        (start, end)
268    }
269}
270
271impl Storable for IndexKey {
272    fn to_bytes(&self) -> Cow<'_, [u8]> {
273        let mut buf = Vec::with_capacity(Self::STORED_SIZE as usize);
274
275        // IMPORTANT: fixed-size IndexName
276        buf.extend_from_slice(&self.index_id.0.to_bytes());
277
278        // len
279        buf.push(self.len);
280
281        // ALL value slots (fixed-size)
282        for value in &self.values {
283            buf.extend_from_slice(value);
284        }
285
286        debug_assert_eq!(buf.len(), Self::STORED_SIZE as usize);
287        buf.into()
288    }
289
290    fn from_bytes(bytes: Cow<'_, [u8]>) -> Self {
291        let bytes = bytes.as_ref();
292        assert_eq!(
293            bytes.len(),
294            Self::STORED_SIZE as usize,
295            "corrupted IndexKey: invalid size"
296        );
297        let mut offset = 0;
298
299        let index_name =
300            IndexName::from_bytes(&bytes[offset..offset + IndexName::STORED_SIZE as usize])
301                .expect("corrupted IndexKey: invalid IndexName");
302        offset += IndexName::STORED_SIZE as usize;
303
304        let len = bytes[offset];
305        offset += 1;
306
307        assert!(
308            len as usize <= MAX_INDEX_FIELDS,
309            "corrupted IndexKey: len exceeds MAX_INDEX_FIELDS"
310        );
311
312        let mut values = [[0u8; 16]; MAX_INDEX_FIELDS];
313        for value in &mut values {
314            value.copy_from_slice(&bytes[offset..offset + 16]);
315            offset += 16;
316        }
317
318        Self {
319            index_id: IndexId(index_name),
320            len,
321            values,
322        }
323    }
324
325    fn into_bytes(self) -> Vec<u8> {
326        self.to_bytes().into_owned()
327    }
328
329    const BOUND: Bound = Bound::Bounded {
330        max_size: Self::STORED_SIZE,
331        is_fixed_size: true,
332    };
333}
334
335///
336/// IndexEntry (VALUE, UNBOUNDED)
337///
338
339#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
340pub struct IndexEntry {
341    keys: HashSet<Key>,
342}
343
344impl IndexEntry {
345    #[must_use]
346    pub fn new(key: Key) -> Self {
347        let mut keys = HashSet::with_capacity(1);
348        keys.insert(key);
349
350        Self { keys }
351    }
352
353    pub fn insert_key(&mut self, key: Key) {
354        self.keys.insert(key);
355    }
356
357    pub fn remove_key(&mut self, key: &Key) {
358        self.keys.remove(key);
359    }
360
361    #[must_use]
362    pub fn contains(&self, key: &Key) -> bool {
363        self.keys.contains(key)
364    }
365
366    #[must_use]
367    pub fn is_empty(&self) -> bool {
368        self.keys.is_empty()
369    }
370
371    #[must_use]
372    pub fn len(&self) -> usize {
373        self.keys.len()
374    }
375
376    #[must_use]
377    pub fn single_key(&self) -> Option<Key> {
378        if self.keys.len() == 1 {
379            self.keys.iter().copied().next()
380        } else {
381            None
382        }
383    }
384}
385
386impl_storable_unbounded!(IndexEntry);
387
388///
389/// TESTS
390///
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395    use crate::traits::Storable;
396
397    #[test]
398    #[should_panic(expected = "corrupted IndexKey: invalid size")]
399    fn index_key_rejects_undersized_bytes() {
400        let buf = vec![0u8; IndexKey::STORED_SIZE as usize - 1];
401        let _ = IndexKey::from_bytes(buf.into());
402    }
403
404    #[test]
405    #[should_panic(expected = "corrupted IndexKey: invalid size")]
406    fn index_key_rejects_oversized_bytes() {
407        let buf = vec![0u8; IndexKey::STORED_SIZE as usize + 1];
408        let _ = IndexKey::from_bytes(buf.into());
409    }
410
411    #[test]
412    #[allow(clippy::cast_possible_truncation)]
413    #[should_panic(expected = "corrupted IndexKey: len exceeds MAX_INDEX_FIELDS")]
414    fn index_key_rejects_len_over_max() {
415        let key = IndexKey::empty(IndexId::max_storable());
416        let mut bytes = key.to_bytes().into_owned();
417        let len_offset = IndexName::STORED_SIZE as usize;
418        bytes[len_offset] = (MAX_INDEX_FIELDS as u8) + 1;
419        let _ = IndexKey::from_bytes(bytes.into());
420    }
421
422    #[test]
423    #[should_panic(expected = "corrupted IndexKey: invalid IndexName")]
424    fn index_key_rejects_invalid_index_name() {
425        let key = IndexKey::empty(IndexId::max_storable());
426        let mut bytes = key.to_bytes().into_owned();
427        bytes[0] = 0;
428        bytes[1] = 0;
429        let _ = IndexKey::from_bytes(bytes.into());
430    }
431
432    #[test]
433    #[expect(clippy::large_types_passed_by_value)]
434    fn index_key_ordering_matches_bytes() {
435        fn make_key(index_id: IndexId, len: u8, first: u8, second: u8) -> IndexKey {
436            let mut values = [[0u8; 16]; MAX_INDEX_FIELDS];
437            values[0] = [first; 16];
438            values[1] = [second; 16];
439            IndexKey {
440                index_id,
441                len,
442                values,
443            }
444        }
445
446        let entity = EntityName::from_static("entity");
447        let idx_a = IndexId(IndexName::from_parts(&entity, &["a"]));
448        let idx_b = IndexId(IndexName::from_parts(&entity, &["b"]));
449
450        let keys = vec![
451            make_key(idx_a, 1, 1, 0),
452            make_key(idx_a, 2, 1, 2),
453            make_key(idx_a, 1, 2, 0),
454            make_key(idx_b, 1, 0, 0),
455        ];
456
457        let mut sorted_by_ord = keys.clone();
458        sorted_by_ord.sort();
459
460        let mut sorted_by_bytes = keys;
461        sorted_by_bytes.sort_by(|a, b| a.to_bytes().cmp(&b.to_bytes()));
462
463        assert_eq!(
464            sorted_by_ord, sorted_by_bytes,
465            "IndexKey Ord and byte ordering diverged"
466        );
467    }
468}
icydb_core/db/store/index.rs

icydb_core/db/store/
index.rs
