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 bytes");
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: invalid index length"
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        let len_usize = len as usize;
319        for value in values.iter().skip(len_usize) {
320            assert!(
321                value.iter().all(|&b| b == 0),
322                "corrupted IndexKey: non-zero fingerprint padding"
323            );
324        }
325
326        Self {
327            index_id: IndexId(index_name),
328            len,
329            values,
330        }
331    }
332
333    fn into_bytes(self) -> Vec<u8> {
334        self.to_bytes().into_owned()
335    }
336
337    const BOUND: Bound = Bound::Bounded {
338        max_size: Self::STORED_SIZE,
339        is_fixed_size: true,
340    };
341}
342
343///
344/// IndexEntry (VALUE, UNBOUNDED)
345///
346
347#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
348pub struct IndexEntry {
349    keys: HashSet<Key>,
350}
351
352impl IndexEntry {
353    #[must_use]
354    pub fn new(key: Key) -> Self {
355        let mut keys = HashSet::with_capacity(1);
356        keys.insert(key);
357
358        Self { keys }
359    }
360
361    pub fn insert_key(&mut self, key: Key) {
362        self.keys.insert(key);
363    }
364
365    pub fn remove_key(&mut self, key: &Key) {
366        self.keys.remove(key);
367    }
368
369    #[must_use]
370    pub fn contains(&self, key: &Key) -> bool {
371        self.keys.contains(key)
372    }
373
374    #[must_use]
375    pub fn is_empty(&self) -> bool {
376        self.keys.is_empty()
377    }
378
379    #[must_use]
380    pub fn len(&self) -> usize {
381        self.keys.len()
382    }
383
384    #[must_use]
385    pub fn single_key(&self) -> Option<Key> {
386        if self.keys.len() == 1 {
387            self.keys.iter().copied().next()
388        } else {
389            None
390        }
391    }
392}
393
394impl_storable_unbounded!(IndexEntry);
395
396///
397/// TESTS
398///
399
400#[cfg(test)]
401mod tests {
402    use super::*;
403    use crate::traits::Storable;
404
405    #[test]
406    #[should_panic(expected = "corrupted IndexKey: invalid size")]
407    fn index_key_rejects_undersized_bytes() {
408        let buf = vec![0u8; IndexKey::STORED_SIZE as usize - 1];
409        let _ = IndexKey::from_bytes(buf.into());
410    }
411
412    #[test]
413    #[should_panic(expected = "corrupted IndexKey: invalid size")]
414    fn index_key_rejects_oversized_bytes() {
415        let buf = vec![0u8; IndexKey::STORED_SIZE as usize + 1];
416        let _ = IndexKey::from_bytes(buf.into());
417    }
418
419    #[test]
420    #[allow(clippy::cast_possible_truncation)]
421    #[should_panic(expected = "corrupted IndexKey: invalid index length")]
422    fn index_key_rejects_len_over_max() {
423        let key = IndexKey::empty(IndexId::max_storable());
424        let mut bytes = key.to_bytes().into_owned();
425        let len_offset = IndexName::STORED_SIZE as usize;
426        bytes[len_offset] = (MAX_INDEX_FIELDS as u8) + 1;
427        let _ = IndexKey::from_bytes(bytes.into());
428    }
429
430    #[test]
431    #[should_panic(expected = "corrupted IndexKey: invalid IndexName bytes")]
432    fn index_key_rejects_invalid_index_name() {
433        let key = IndexKey::empty(IndexId::max_storable());
434        let mut bytes = key.to_bytes().into_owned();
435        bytes[0] = 0;
436        bytes[1] = 0;
437        let _ = IndexKey::from_bytes(bytes.into());
438    }
439
440    #[test]
441    #[should_panic(expected = "corrupted IndexKey: non-zero fingerprint padding")]
442    fn index_key_rejects_fingerprint_padding() {
443        let key = IndexKey::empty(IndexId::max_storable());
444        let mut bytes = key.to_bytes().into_owned();
445        let values_offset = IndexName::STORED_SIZE as usize + 1;
446        bytes[values_offset] = 1;
447        let _ = IndexKey::from_bytes(bytes.into());
448    }
449
450    #[test]
451    #[expect(clippy::large_types_passed_by_value)]
452    fn index_key_ordering_matches_bytes() {
453        fn make_key(index_id: IndexId, len: u8, first: u8, second: u8) -> IndexKey {
454            let mut values = [[0u8; 16]; MAX_INDEX_FIELDS];
455            values[0] = [first; 16];
456            values[1] = [second; 16];
457            IndexKey {
458                index_id,
459                len,
460                values,
461            }
462        }
463
464        let entity = EntityName::from_static("entity");
465        let idx_a = IndexId(IndexName::from_parts(&entity, &["a"]));
466        let idx_b = IndexId(IndexName::from_parts(&entity, &["b"]));
467
468        let keys = vec![
469            make_key(idx_a, 1, 1, 0),
470            make_key(idx_a, 2, 1, 2),
471            make_key(idx_a, 1, 2, 0),
472            make_key(idx_b, 1, 0, 0),
473        ];
474
475        let mut sorted_by_ord = keys.clone();
476        sorted_by_ord.sort();
477
478        let mut sorted_by_bytes = keys;
479        sorted_by_bytes.sort_by(|a, b| a.to_bytes().cmp(&b.to_bytes()));
480
481        assert_eq!(
482            sorted_by_ord, sorted_by_bytes,
483            "IndexKey Ord and byte ordering diverged"
484        );
485    }
486}
icydb_core/db/store/index.rs

icydb_core/db/store/
index.rs
