icydb_core/db/store/
index.rs

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

icydb_core/db/store/
index.rs
