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#[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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
52pub enum IndexInsertOutcome {
53 Inserted,
54 Skipped,
55}
56
57#[derive(Clone, Copy, Debug, Eq, PartialEq)]
64pub enum IndexRemoveOutcome {
65 Removed,
66 Skipped,
67}
68
69#[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#[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#[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 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 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}