icydb_core/db/store/
index.rs

1use crate::{
2    Error, IndexSpec, Key, MAX_INDEX_FIELDS, Value,
3    db::{
4        executor::ExecutorError,
5        store::{DataKey, StoreRegistry},
6    },
7    obs::metrics,
8    traits::EntityKind,
9};
10use candid::CandidType;
11use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, memory::VirtualMemory};
12use canic_memory::{impl_storable_bounded, impl_storable_unbounded};
13use canic_utils::hash::hash_u64;
14use derive_more::{Deref, DerefMut, Display};
15use serde::{Deserialize, Serialize};
16use std::{
17    collections::HashSet,
18    fmt::{self, Display},
19};
20
21///
22/// IndexStoreRegistry
23///
24
25#[derive(Deref, DerefMut)]
26pub struct IndexStoreRegistry(StoreRegistry<IndexStore>);
27
28impl IndexStoreRegistry {
29    #[must_use]
30    #[allow(clippy::new_without_default)]
31    /// Create an empty index store registry.
32    pub fn new() -> Self {
33        Self(StoreRegistry::new())
34    }
35}
36
37///
38/// IndexStore
39///
40
41#[derive(Deref, DerefMut)]
42pub struct IndexStore(BTreeMap<IndexKey, IndexEntry, VirtualMemory<DefaultMemoryImpl>>);
43
44impl IndexStore {
45    #[must_use]
46    /// Initialize an index store with the provided backing memory.
47    pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
48        Self(BTreeMap::init(memory))
49    }
50
51    /// Inserts the given entity into the index defined by `I`.
52    /// - If `I::UNIQUE`, insertion will fail if a conflicting entry already exists.
53    /// - If the entity is missing required fields for this index, insertion is skipped.
54    /// - Insert an entity's index entry, enforcing uniqueness where required.
55    pub fn insert_index_entry<E: EntityKind>(
56        &mut self,
57        entity: &E,
58        index: &IndexSpec,
59    ) -> Result<(), Error> {
60        // Skip if index key can't be built (e.g. optional fields missing)
61        let Some(index_key) = IndexKey::new(entity, index) else {
62            return Ok(());
63        };
64        let key = entity.key();
65
66        if let Some(mut existing) = self.get(&index_key) {
67            if index.unique {
68                // Already present → skip redundant write
69                if existing.contains(&key) {
70                    return Ok(());
71                }
72
73                // Any different existing key violates UNIQUE
74                if !existing.is_empty() {
75                    metrics::with_state_mut(|m| metrics::record_unique_violation_for::<E>(m));
76
77                    return Err(ExecutorError::index_violation(E::PATH, index.fields))?;
78                }
79
80                self.insert(index_key.clone(), IndexEntry::new(index.fields, key));
81            } else {
82                existing.insert_key(key); // <-- add to the set
83                self.insert(index_key.clone(), existing);
84            }
85        } else {
86            self.insert(index_key, IndexEntry::new(index.fields, key));
87        }
88        metrics::with_state_mut(|m| {
89            m.ops.index_inserts += 1;
90            let entry = m.entities.entry(E::PATH.to_string()).or_default();
91            entry.index_inserts = entry.index_inserts.saturating_add(1);
92        });
93
94        Ok(())
95    }
96
97    // remove_index_entry
98    /// Remove an entity's entry for the given index if present.
99    pub fn remove_index_entry<E: EntityKind>(&mut self, entity: &E, index: &IndexSpec) {
100        // Skip if index key can't be built (e.g. optional fields missing)
101        let Some(index_key) = IndexKey::new(entity, index) else {
102            return;
103        };
104
105        if let Some(mut existing) = self.get(&index_key) {
106            existing.remove_key(&entity.key()); // remove from the set
107
108            if existing.is_empty() {
109                self.remove(&index_key);
110            } else {
111                // Move the updated entry back without cloning
112                self.insert(index_key, existing);
113            }
114            metrics::with_state_mut(|m| {
115                m.ops.index_removes += 1;
116                let entry = m.entities.entry(E::PATH.to_string()).or_default();
117                entry.index_removes = entry.index_removes.saturating_add(1);
118            });
119        }
120    }
121
122    #[must_use]
123    /// Resolve data keys for a given index prefix.
124    pub fn resolve_data_values<E: EntityKind>(
125        &self,
126        index: &IndexSpec,
127        prefix: &[Value],
128    ) -> Vec<DataKey> {
129        let mut out = Vec::new();
130
131        for (_, entry) in self.iter_with_hashed_prefix::<E>(index, prefix) {
132            out.extend(entry.keys.iter().map(|&k| DataKey::new::<E>(k)));
133        }
134
135        out
136    }
137
138    /// Sum of bytes used by all index entries.
139    pub fn memory_bytes(&self) -> u64 {
140        self.iter()
141            .map(|entry| u64::from(IndexKey::STORABLE_MAX_SIZE) + entry.value().len() as u64)
142            .sum()
143    }
144
145    /// Internal: iterate entries for this index whose `hashed_values` start with the hashed `prefix`.
146    /// Uses a bounded range for efficient scanning.
147    fn iter_with_hashed_prefix<E: EntityKind>(
148        &self,
149        index: &IndexSpec,
150        prefix: &[Value],
151    ) -> impl Iterator<Item = (IndexKey, IndexEntry)> {
152        let index_id = IndexId::new::<E>(index);
153        let hashed_prefix_opt = Self::index_fingerprints(prefix); // Option<Vec<[u8;16]>>
154
155        // Compute start..end bounds. If the prefix isn't indexable, construct an empty range
156        // (same iterator type) by using identical start==end under the same index_id.
157        let (start_key, end_key) = if let Some(hp) = hashed_prefix_opt {
158            IndexKey::bounds_for_prefix(index_id, hp)
159        } else {
160            (
161                IndexKey {
162                    index_id,
163                    hashed_values: Vec::new(),
164                },
165                IndexKey {
166                    index_id,
167                    hashed_values: Vec::new(),
168                },
169            )
170        };
171
172        self.range(start_key..end_key)
173            .map(|entry| (entry.key().clone(), entry.value()))
174    }
175
176    fn index_fingerprints(values: &[Value]) -> Option<Vec<[u8; 16]>> {
177        // collects to Option<Vec<_>>: None if any element was non-indexable
178        values.iter().map(Value::to_index_fingerprint).collect()
179    }
180}
181
182///
183/// IndexId
184///
185
186#[derive(
187    Clone,
188    Debug,
189    Display,
190    Copy,
191    PartialEq,
192    Eq,
193    Hash,
194    PartialOrd,
195    Ord,
196    CandidType,
197    Serialize,
198    Deserialize,
199)]
200pub struct IndexId(u64);
201
202impl IndexId {
203    #[must_use]
204    /// Deterministic index identifier derived from entity path and field list.
205    pub fn new<E: EntityKind>(index: &IndexSpec) -> Self {
206        Self::from_path_and_fields(E::PATH, index.fields)
207    }
208
209    fn from_path_and_fields(path: &str, fields: &[&str]) -> Self {
210        let cap = path.len() + fields.iter().map(|f| f.len() + 1).sum::<usize>();
211        let mut buffer = Vec::with_capacity(cap);
212
213        // much more efficient than format
214        buffer.extend_from_slice(path.as_bytes());
215        for field in fields {
216            buffer.extend_from_slice(field.as_bytes());
217            buffer.extend_from_slice(b"|");
218        }
219
220        Self(hash_u64(&buffer))
221    }
222
223    #[must_use]
224    /// Worst-case index identifier used for sizing tests.
225    pub fn max_storable() -> Self {
226        Self::from_path_and_fields(
227            "path::to::long::entity::name::Entity",
228            &[
229                "long_field_one",
230                "long_field_two",
231                "long_field_three",
232                "long_field_four",
233            ],
234        )
235    }
236}
237
238///
239/// IndexKey
240///
241
242#[derive(Clone, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)]
243pub struct IndexKey {
244    pub index_id: IndexId,
245    pub hashed_values: Vec<[u8; 16]>,
246}
247
248impl IndexKey {
249    // Sized with headroom for worst‑case hashed key payload
250    pub const STORABLE_MAX_SIZE: u32 = 180;
251
252    #[must_use]
253    /// Build an index key from an entity and spec, returning None if a component is missing/non-indexable.
254    pub fn new<E: EntityKind>(entity: &E, index: &IndexSpec) -> Option<Self> {
255        let mut hashed_values = Vec::<[u8; 16]>::with_capacity(index.fields.len());
256
257        // get each value and convert to key
258        for field in index.fields {
259            let value = entity.get_value(field)?;
260            let fp = value.to_index_fingerprint()?; // bail if any component is non-indexable
261
262            hashed_values.push(fp);
263        }
264
265        Some(Self {
266            index_id: IndexId::new::<E>(index),
267            hashed_values,
268        })
269    }
270
271    // max_storable
272    #[must_use]
273    /// Largest representable index key (for sizing).
274    pub fn max_storable() -> Self {
275        Self {
276            index_id: IndexId::max_storable(),
277            hashed_values: (0..MAX_INDEX_FIELDS).map(|_| [u8::MAX; 16]).collect(),
278        }
279    }
280
281    /// Compute the bounded start..end keys for a given hashed prefix under an index id.
282    /// End is exclusive and created by appending a single 0xFF..0xFF block to the prefix.
283    #[must_use]
284    /// The returned range can be used with `BTreeMap::range`.
285    pub fn bounds_for_prefix(index_id: IndexId, mut prefix: Vec<[u8; 16]>) -> (Self, Self) {
286        let start = Self {
287            index_id,
288            hashed_values: prefix.clone(),
289        };
290        prefix.push([0xFF; 16]);
291        let end = Self {
292            index_id,
293            hashed_values: prefix,
294        };
295        (start, end)
296    }
297}
298
299impl Display for IndexKey {
300    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
301        write!(
302            f,
303            "id: {}, values: {}",
304            self.index_id,
305            self.hashed_values.len()
306        )
307    }
308}
309
310impl_storable_bounded!(IndexKey, IndexKey::STORABLE_MAX_SIZE, false);
311
312///
313/// IndexEntry
314///
315
316#[derive(CandidType, Clone, Debug, Deserialize, Serialize)]
317pub struct IndexEntry {
318    fields: Vec<String>,
319    keys: HashSet<Key>,
320}
321
322impl IndexEntry {
323    #[must_use]
324    /// Create an index entry with the initial key.
325    pub fn new(fields: &[&str], key: Key) -> Self {
326        let mut key_set = HashSet::with_capacity(1);
327        key_set.insert(key);
328
329        Self {
330            fields: fields.iter().map(ToString::to_string).collect(),
331            keys: key_set,
332        }
333    }
334
335    /// Insert a key into the entry's key set.
336    pub fn insert_key(&mut self, key: Key) {
337        let _ = self.keys.insert(key);
338    }
339
340    /// Removes the key from the set.
341    pub fn remove_key(&mut self, key: &Key) {
342        let _ = self.keys.remove(key);
343    }
344
345    /// Checks if the set contains the key.
346    #[must_use]
347    pub fn contains(&self, key: &Key) -> bool {
348        self.keys.contains(key)
349    }
350
351    /// Returns `true` if the set is empty.
352    #[must_use]
353    pub fn is_empty(&self) -> bool {
354        self.keys.is_empty()
355    }
356
357    /// Returns number of keys in the entry.
358    #[must_use]
359    pub fn len(&self) -> usize {
360        self.keys.len()
361    }
362
363    /// Returns the only key when exactly one exists.
364    #[must_use]
365    pub fn single_key(&self) -> Option<Key> {
366        if self.keys.len() == 1 {
367            self.keys.iter().copied().next()
368        } else {
369            None
370        }
371    }
372
373    /// Returns keys as a sorted `Vec<Key>` (useful for serialization/debug).
374    #[must_use]
375    pub fn to_sorted_vec(&self) -> Vec<Key> {
376        let mut keys: Vec<_> = self.keys.iter().copied().collect();
377        keys.sort_unstable(); // uses Ord, fast
378        keys
379    }
380}
381
382impl_storable_unbounded!(IndexEntry);
383
384///
385/// TESTS
386///
387
388#[cfg(test)]
389mod tests {
390    use super::*;
391    use crate::traits::*;
392
393    #[test]
394    fn index_key_max_size_is_bounded() {
395        let index_key = IndexKey::max_storable();
396        let size = Storable::to_bytes(&index_key).len();
397
398        assert!(
399            size <= IndexKey::STORABLE_MAX_SIZE as usize,
400            "serialized IndexKey too large: got {size} bytes (limit {})",
401            IndexKey::STORABLE_MAX_SIZE
402        );
403    }
404
405    #[test]
406    fn index_entry_round_trip() {
407        let original = IndexEntry::new(&["a", "b"], Key::from(1u64));
408        let encoded = Storable::to_bytes(&original);
409        let decoded = IndexEntry::from_bytes(encoded);
410
411        assert_eq!(original.fields, decoded.fields);
412        assert_eq!(original.to_sorted_vec(), decoded.to_sorted_vec());
413    }
414}
icydb_core/db/store/index.rs

icydb_core/db/store/
index.rs
