Skip to main content

icydb_core/db/index/
store.rs

1use crate::{
2    db::{
3        index::{
4            entry::{
5                IndexEntry, IndexEntryCorruption, IndexEntryEncodeError, MAX_INDEX_ENTRY_KEYS,
6                RawIndexEntry,
7            },
8            fingerprint,
9            key::{IndexId, IndexKey, RawIndexKey},
10        },
11        store::{DataKey, StoreRegistry},
12    },
13    error::{ErrorClass, ErrorOrigin, InternalError},
14    key::KeyEncodeError,
15    prelude::{EntityKind, IndexModel, Value},
16};
17use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, memory::VirtualMemory};
18use derive_more::{Deref, DerefMut};
19use thiserror::Error as ThisError;
20
21///
22/// IndexStoreRegistry
23///
24/// Registry managing all index stores for the database.
25/// Provides lifecycle and lookup for per-index stable-memory stores.
26///
27
28#[derive(Deref, DerefMut)]
29pub struct IndexStoreRegistry(StoreRegistry<IndexStore>);
30
31impl IndexStoreRegistry {
32    #[must_use]
33    pub fn new() -> Self {
34        Self(StoreRegistry::new())
35    }
36}
37
38impl Default for IndexStoreRegistry {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44///
45/// IndexInsertOutcome
46///
47/// Result of attempting to insert an entity into an index.
48/// Distinguishes between a successful mutation and a no-op.
49///
50
51#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub enum IndexInsertOutcome {
53    Inserted,
54    Skipped,
55}
56
57///
58/// IndexRemoveOutcome
59///
60/// Result of attempting to remove an entity from an index.
61///
62
63#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum IndexRemoveOutcome {
65    Removed,
66    Skipped,
67}
68
69///
70/// IndexRemoveError
71///
72/// Errors returned when removing an entry from an index.
73///
74
75#[derive(Debug, ThisError)]
76pub enum IndexRemoveError {
77    #[error("index entry corrupted: {0}")]
78    Corruption(#[from] IndexEntryCorruption),
79    #[error("index entry key encoding failed: {0}")]
80    KeyEncoding(#[from] KeyEncodeError),
81    #[error("index key construction failed: {0}")]
82    KeyConstruction(#[from] InternalError),
83}
84
85///
86/// IndexInsertError
87///
88/// Errors that may occur while inserting into an index.
89/// Represents both logical constraint violations and corruption.
90///
91
92#[derive(Debug, ThisError)]
93pub enum IndexInsertError {
94    #[error("unique index violation")]
95    UniqueViolation,
96    #[error("index entry corrupted: {0}")]
97    CorruptedEntry(#[from] IndexEntryCorruption),
98    #[error("index entry exceeds max keys: {keys} (limit {MAX_INDEX_ENTRY_KEYS})")]
99    EntryTooLarge { keys: usize },
100    #[error("index entry key encoding failed: {0}")]
101    KeyEncoding(#[from] KeyEncodeError),
102    #[error("index key construction failed: {0}")]
103    KeyConstruction(#[from] InternalError),
104}
105
106///
107/// IndexStore
108///
109/// Stable-memory backed secondary index.
110/// Maps composite index keys to one or more entity primary keys.
111///
112
113#[derive(Deref, DerefMut)]
114pub struct IndexStore(BTreeMap<RawIndexKey, RawIndexEntry, VirtualMemory<DefaultMemoryImpl>>);
115
116impl IndexStore {
117    #[must_use]
118    pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
119        Self(BTreeMap::init(memory))
120    }
121
122    pub fn insert_index_entry<E: EntityKind>(
123        &mut self,
124        entity: &E,
125        index: &IndexModel,
126    ) -> Result<IndexInsertOutcome, IndexInsertError> {
127        let Some(index_key) = IndexKey::new(entity, index)? else {
128            return Ok(IndexInsertOutcome::Skipped);
129        };
130        let raw_key = index_key.to_raw();
131        let key = entity.key();
132
133        if let Some(raw_entry) = self.get(&raw_key) {
134            let mut entry = raw_entry.try_decode()?;
135            if index.unique {
136                if entry.len() > 1 {
137                    return Err(IndexEntryCorruption::NonUniqueEntry { keys: entry.len() }.into());
138                }
139                if entry.contains(&key) {
140                    return Ok(IndexInsertOutcome::Skipped);
141                }
142                if !entry.is_empty() {
143                    return Err(IndexInsertError::UniqueViolation);
144                }
145                let entry = IndexEntry::new(key);
146                let raw_entry = RawIndexEntry::try_from_entry(&entry).map_err(|err| match err {
147                    IndexEntryEncodeError::TooManyKeys { keys } => {
148                        IndexInsertError::EntryTooLarge { keys }
149                    }
150                    IndexEntryEncodeError::KeyEncoding(err) => IndexInsertError::KeyEncoding(err),
151                })?;
152                self.insert(raw_key, raw_entry);
153            } else {
154                entry.insert_key(key);
155                let raw_entry = RawIndexEntry::try_from_entry(&entry).map_err(|err| match err {
156                    IndexEntryEncodeError::TooManyKeys { keys } => {
157                        IndexInsertError::EntryTooLarge { keys }
158                    }
159                    IndexEntryEncodeError::KeyEncoding(err) => IndexInsertError::KeyEncoding(err),
160                })?;
161                self.insert(raw_key, raw_entry);
162            }
163        } else {
164            let entry = IndexEntry::new(key);
165            let raw_entry = RawIndexEntry::try_from_entry(&entry).map_err(|err| match err {
166                IndexEntryEncodeError::TooManyKeys { keys } => {
167                    IndexInsertError::EntryTooLarge { keys }
168                }
169                IndexEntryEncodeError::KeyEncoding(err) => IndexInsertError::KeyEncoding(err),
170            })?;
171            self.insert(raw_key, raw_entry);
172        }
173
174        Ok(IndexInsertOutcome::Inserted)
175    }
176
177    pub fn remove_index_entry<E: EntityKind>(
178        &mut self,
179        entity: &E,
180        index: &IndexModel,
181    ) -> Result<IndexRemoveOutcome, IndexRemoveError> {
182        let Some(index_key) = IndexKey::new(entity, index)? else {
183            return Ok(IndexRemoveOutcome::Skipped);
184        };
185        let raw_key = index_key.to_raw();
186
187        if let Some(raw_entry) = self.get(&raw_key) {
188            let mut entry = raw_entry.try_decode()?;
189            let entity_key = entity.key();
190            // Treat missing membership as index/data divergence.
191            if !entry.contains(&entity_key) {
192                return Err(IndexRemoveError::Corruption(
193                    IndexEntryCorruption::missing_key(raw_key, entity_key),
194                ));
195            }
196            entry.remove_key(&entity_key);
197            if entry.is_empty() {
198                self.remove(&raw_key);
199            } else {
200                let raw_entry = RawIndexEntry::try_from_entry(&entry).map_err(|err| match err {
201                    IndexEntryEncodeError::TooManyKeys { keys } => {
202                        IndexEntryCorruption::TooManyKeys { count: keys }.into()
203                    }
204                    IndexEntryEncodeError::KeyEncoding(err) => IndexRemoveError::KeyEncoding(err),
205                })?;
206                self.insert(raw_key, raw_entry);
207            }
208            return Ok(IndexRemoveOutcome::Removed);
209        }
210
211        Ok(IndexRemoveOutcome::Skipped)
212    }
213
214    pub fn resolve_data_values<E: EntityKind>(
215        &self,
216        index: &IndexModel,
217        prefix: &[Value],
218    ) -> Result<Vec<DataKey>, InternalError> {
219        let mut out = Vec::new();
220        if prefix.len() > index.fields.len() {
221            return Err(InternalError::new(
222                ErrorClass::Unsupported,
223                ErrorOrigin::Index,
224                format!(
225                    "index prefix length {} exceeds field count {}",
226                    prefix.len(),
227                    index.fields.len()
228                ),
229            ));
230        }
231        let index_id = IndexId::new_unchecked::<E>(index);
232
233        let mut fps = Vec::with_capacity(prefix.len());
234        for value in prefix {
235            let Some(fp) = fingerprint::to_index_fingerprint(value)? else {
236                return Err(InternalError::new(
237                    ErrorClass::Unsupported,
238                    ErrorOrigin::Index,
239                    "index prefix value is not indexable",
240                ));
241            };
242            fps.push(fp);
243        }
244
245        let (start, end) = IndexKey::bounds_for_prefix(index_id, index.fields.len(), &fps);
246        let start_raw = start.to_raw();
247        let end_raw = end.to_raw();
248
249        for entry in self.range(start_raw..=end_raw) {
250            let _ = IndexKey::try_from_raw(entry.key()).map_err(|err| {
251                InternalError::new(
252                    ErrorClass::Corruption,
253                    ErrorOrigin::Index,
254                    format!("index key corrupted during resolve: {err}"),
255                )
256            })?;
257            let decoded = entry.value().try_decode().map_err(|err| {
258                InternalError::new(ErrorClass::Corruption, ErrorOrigin::Index, err.to_string())
259            })?;
260            if index.unique && decoded.len() != 1 {
261                return Err(InternalError::new(
262                    ErrorClass::Corruption,
263                    ErrorOrigin::Index,
264                    "unique index entry contains an unexpected number of keys",
265                ));
266            }
267            out.extend(decoded.iter_keys().map(|k| DataKey::new::<E>(k)));
268        }
269
270        Ok(out)
271    }
272
273    /// Sum of bytes used by all index entries.
274    pub fn memory_bytes(&self) -> u64 {
275        self.iter()
276            .map(|entry| u64::from(IndexKey::STORED_SIZE) + entry.value().len() as u64)
277            .sum()
278    }
279}