1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::Arc;
3
4use hashtree_core::{Cid, DirEntry, HashTree, HashTreeConfig, LinkType, Store};
5use hashtree_index::{BTree, BTreeOptions, SearchIndex};
6
7use crate::helpers::merge_search_index_options;
8use crate::manifest::write_collection_manifest_metadata;
9use crate::schema::normalize_typed_collection_item;
10use crate::{
11 create_empty_collection_state, default_search_prefix, load_collection_state,
12 CollectionDefinition, CollectionEntryContext, CollectionError, CollectionOptions,
13 CollectionState, CollectionWriteContext, COLLECTION_MANIFEST_METADATA_FILE, MANIFEST_BY_ID,
14};
15
16pub struct CollectionWriter<S: Store, T> {
17 tree: HashTree<S>,
18 index: BTree<S>,
19 search_indexes: BTreeMap<String, SearchIndex<S>>,
20 definition: CollectionDefinition<T>,
21 state: CollectionState,
22}
23
24impl<S: Store, T> CollectionWriter<S, T> {
25 pub fn new(store: Arc<S>, definition: CollectionDefinition<T>) -> Self {
26 Self::with_options(store, definition, CollectionOptions::default())
27 }
28
29 pub fn with_options(
30 store: Arc<S>,
31 definition: CollectionDefinition<T>,
32 options: CollectionOptions,
33 ) -> Self {
34 let state = create_empty_collection_state(&definition);
35 Self::with_state_and_options(store, definition, state, options)
36 }
37
38 pub fn with_state(
39 store: Arc<S>,
40 definition: CollectionDefinition<T>,
41 state: CollectionState,
42 ) -> Self {
43 Self::with_state_and_options(store, definition, state, CollectionOptions::default())
44 }
45
46 pub fn with_state_and_options(
47 store: Arc<S>,
48 definition: CollectionDefinition<T>,
49 state: CollectionState,
50 options: CollectionOptions,
51 ) -> Self {
52 let search_indexes = definition
53 .search_indexes()
54 .iter()
55 .map(|index| {
56 (
57 index.name().to_string(),
58 SearchIndex::new(
59 Arc::clone(&store),
60 merge_search_index_options(index.options().clone(), &options),
61 ),
62 )
63 })
64 .collect();
65
66 Self {
67 tree: HashTree::new(HashTreeConfig::new(Arc::clone(&store))),
68 index: BTree::new(
69 store,
70 BTreeOptions {
71 order: options.btree_order,
72 },
73 ),
74 search_indexes,
75 definition,
76 state,
77 }
78 }
79
80 pub async fn from_root(
81 store: Arc<S>,
82 definition: CollectionDefinition<T>,
83 root: Option<&Cid>,
84 options: CollectionOptions,
85 ) -> Result<Self, CollectionError> {
86 let state = load_collection_state(Arc::clone(&store), &definition, root).await?;
87 Ok(Self::with_state_and_options(
88 store, definition, state, options,
89 ))
90 }
91
92 pub fn snapshot(&self) -> CollectionState {
93 self.state.clone()
94 }
95
96 pub fn state(&self) -> &CollectionState {
97 &self.state
98 }
99
100 pub async fn write_root(&self) -> Result<Option<Cid>, CollectionError> {
101 let mut entries = Vec::new();
102 if let Some(cid) = self.state.by_id_root.as_ref() {
103 entries.push(DirEntry::from_cid(MANIFEST_BY_ID, cid).with_link_type(LinkType::Dir));
104 }
105 for index in self.definition.key_indexes() {
106 if let Some(cid) = self.state.key_root(index.name()) {
107 entries.push(DirEntry::from_cid(index.name(), cid).with_link_type(LinkType::Dir));
108 }
109 }
110 for index in self.definition.search_indexes() {
111 if let Some(cid) = self.state.search_root(index.name()) {
112 entries.push(DirEntry::from_cid(index.name(), cid).with_link_type(LinkType::Dir));
113 }
114 }
115 if let Some((cid, size)) =
116 write_collection_manifest_metadata(&self.tree, &self.definition).await?
117 {
118 entries.push(
119 DirEntry::from_cid(COLLECTION_MANIFEST_METADATA_FILE, &cid)
120 .with_size(size)
121 .with_link_type(LinkType::File),
122 );
123 }
124
125 if entries.is_empty() {
126 return Ok(None);
127 }
128
129 Ok(Some(self.tree.put_directory(entries).await?))
130 }
131
132 fn read_search_root_group(&self, root_name: &str) -> Option<Cid> {
133 for index in self.definition.search_indexes() {
134 if index.root_name().unwrap_or(index.name()) == root_name {
135 return self.state.search_root(index.name()).cloned();
136 }
137 }
138 None
139 }
140
141 fn assign_search_root_groups(&mut self, groups: &BTreeMap<String, Option<Cid>>) {
142 for index in self.definition.search_indexes() {
143 let root_name = index.root_name().unwrap_or(index.name());
144 if let Some(root) = groups.get(root_name) {
145 self.state
146 .search_roots
147 .insert(index.name().to_string(), root.clone());
148 }
149 }
150 }
151
152 fn has_derived_indexes(&self) -> bool {
153 !self.definition.key_indexes().is_empty() || !self.definition.search_indexes().is_empty()
154 }
155}
156
157impl<S: Store, T: Clone> CollectionWriter<S, T> {
158 pub fn normalize(&self, item: &T) -> Result<T, CollectionError> {
159 normalize_typed_collection_item(&self.definition, item)
160 }
161
162 pub async fn put(
163 &mut self,
164 item: &T,
165 cid: &Cid,
166 previous: Option<&T>,
167 ) -> Result<CollectionState, CollectionError> {
168 self.put_with_context(item, cid, previous, None, None).await
169 }
170
171 pub async fn replace(
172 &mut self,
173 item: &T,
174 cid: &Cid,
175 previous: &T,
176 ) -> Result<CollectionState, CollectionError> {
177 self.replace_with_context(item, cid, previous, None, None)
178 .await
179 }
180
181 pub async fn put_with_context(
182 &mut self,
183 item: &T,
184 cid: &Cid,
185 previous: Option<&T>,
186 context: Option<&CollectionWriteContext>,
187 previous_context: Option<&CollectionWriteContext>,
188 ) -> Result<CollectionState, CollectionError> {
189 let next_item = self.normalize(item)?;
190 let id = self.definition.item_id(&next_item)?;
191 let previous_item = previous
192 .map(|previous| self.normalize(previous))
193 .transpose()?;
194
195 if previous_item.is_none() && self.has_derived_indexes() {
196 let existing = if let Some(root) = self.state.by_id_root.as_ref() {
197 self.index.get_link(Some(root), &id).await?
198 } else {
199 None
200 };
201 if existing.is_some() {
202 return Err(CollectionError::MissingPreviousForOverwrite { id });
203 }
204 }
205
206 if let Some(previous_item) = previous_item.as_ref() {
207 self.delete_normalized_with_context(previous_item, previous_context.or(context))
208 .await?;
209 }
210
211 self.put_normalized_with_context(&next_item, cid, context)
212 .await
213 }
214
215 pub async fn replace_with_context(
216 &mut self,
217 item: &T,
218 cid: &Cid,
219 previous: &T,
220 context: Option<&CollectionWriteContext>,
221 previous_context: Option<&CollectionWriteContext>,
222 ) -> Result<CollectionState, CollectionError> {
223 self.put_with_context(item, cid, Some(previous), context, previous_context)
224 .await
225 }
226
227 pub async fn delete(&mut self, item: &T) -> Result<CollectionState, CollectionError> {
228 self.delete_with_context(item, None).await
229 }
230
231 pub async fn delete_with_context(
232 &mut self,
233 item: &T,
234 context: Option<&CollectionWriteContext>,
235 ) -> Result<CollectionState, CollectionError> {
236 let next_item = self.normalize(item)?;
237 self.delete_normalized_with_context(&next_item, context)
238 .await
239 }
240
241 pub async fn put_batch<I>(&mut self, entries: I) -> Result<CollectionState, CollectionError>
247 where
248 I: IntoIterator<Item = (T, Cid, Option<T>)>,
249 {
250 self.put_batch_with_known_absent_ids(entries, std::iter::empty())
251 .await
252 }
253
254 pub async fn put_batch_with_known_absent_ids<I, K>(
260 &mut self,
261 entries: I,
262 known_absent_ids: K,
263 ) -> Result<CollectionState, CollectionError>
264 where
265 I: IntoIterator<Item = (T, Cid, Option<T>)>,
266 K: IntoIterator<Item = String>,
267 {
268 if !self.definition.search_indexes().is_empty() {
269 return Err(CollectionError::Validation(
270 "batch collection updates do not support search indexes".to_string(),
271 ));
272 }
273
274 let known_absent_ids = known_absent_ids.into_iter().collect::<BTreeSet<_>>();
275 let mut normalized_entries = Vec::new();
276 let mut ids_without_previous = BTreeSet::new();
277 let mut seen_ids = BTreeSet::new();
278 for (item, cid, previous) in entries {
279 let item = self.normalize(&item)?;
280 let id = self.definition.item_id(&item)?;
281 if !seen_ids.insert(id.clone()) {
282 return Err(CollectionError::Validation(format!(
283 "batch collection update contains duplicate id `{id}`"
284 )));
285 }
286 let previous = previous
287 .map(|previous| {
288 let previous = self.normalize(&previous)?;
289 let previous_id = self.definition.item_id(&previous)?;
290 Ok::<_, CollectionError>((previous_id, previous))
291 })
292 .transpose()?;
293 if previous.is_none() && !known_absent_ids.contains(&id) {
294 ids_without_previous.insert(id.clone());
295 }
296 normalized_entries.push((id, item, cid, previous));
297 }
298 normalized_entries.sort_by(|left, right| left.0.cmp(&right.0));
299
300 let existing_without_previous = self
301 .index
302 .get_links(self.state.by_id_root.as_ref(), ids_without_previous)
303 .await?;
304 if let Some(id) = existing_without_previous.keys().next() {
305 return Err(CollectionError::MissingPreviousForOverwrite { id: id.clone() });
306 }
307
308 let mut by_id_changes = BTreeMap::<String, Option<Cid>>::new();
309 for (_, _, _, previous) in &normalized_entries {
310 if let Some((previous_id, _)) = previous {
311 by_id_changes.insert(previous_id.clone(), None);
312 }
313 }
314 for (id, _, cid, _) in &normalized_entries {
315 by_id_changes.insert(id.clone(), Some(cid.clone()));
316 }
317
318 self.state.by_id_root = self
319 .index
320 .update_links(self.state.by_id_root.as_ref(), by_id_changes)
321 .await?;
322 for index in self.definition.key_indexes() {
323 let mut changes = BTreeMap::<String, Option<Cid>>::new();
324 for (_, _, _, previous) in &normalized_entries {
325 if let Some((_, previous)) = previous {
326 for key in index.materialize_keys(previous) {
327 changes.insert(key, None);
328 }
329 }
330 }
331 for (_, item, cid, _) in &normalized_entries {
332 for key in index.materialize_keys(item) {
333 changes.insert(key, Some(cid.clone()));
334 }
335 }
336 let root = self
337 .index
338 .update_links(self.state.key_root(index.name()), changes)
339 .await?;
340 self.state.key_roots.insert(index.name().to_string(), root);
341 }
342
343 Ok(self.snapshot())
344 }
345
346 pub async fn rebuild<I>(&mut self, entries: I) -> Result<CollectionState, CollectionError>
347 where
348 I: IntoIterator<Item = (T, Cid)>,
349 {
350 let mut final_entries = BTreeMap::<String, (T, Cid, Option<CollectionWriteContext>)>::new();
351 for (item, cid) in entries {
352 let next_item = self.normalize(&item)?;
353 let id = self.definition.item_id(&next_item)?;
354 final_entries.insert(id, (next_item, cid, None));
355 }
356
357 self.rebuild_from_final_entries(final_entries).await
358 }
359
360 pub async fn reindex<I>(&mut self, entries: I) -> Result<CollectionState, CollectionError>
361 where
362 I: IntoIterator<Item = (T, Cid)>,
363 {
364 self.rebuild(entries).await
365 }
366
367 pub async fn rebuild_with_context<I>(
368 &mut self,
369 entries: I,
370 ) -> Result<CollectionState, CollectionError>
371 where
372 I: IntoIterator<Item = (T, Cid, Option<CollectionWriteContext>)>,
373 {
374 let mut final_entries = BTreeMap::<String, (T, Cid, Option<CollectionWriteContext>)>::new();
375 for (item, cid, context) in entries {
376 let next_item = self.normalize(&item)?;
377 let id = self.definition.item_id(&next_item)?;
378 final_entries.insert(id, (next_item, cid, context));
379 }
380
381 self.rebuild_from_final_entries(final_entries).await
382 }
383
384 pub async fn reindex_with_context<I>(
385 &mut self,
386 entries: I,
387 ) -> Result<CollectionState, CollectionError>
388 where
389 I: IntoIterator<Item = (T, Cid, Option<CollectionWriteContext>)>,
390 {
391 self.rebuild_with_context(entries).await
392 }
393
394 async fn put_normalized_with_context(
395 &mut self,
396 item: &T,
397 cid: &Cid,
398 context: Option<&CollectionWriteContext>,
399 ) -> Result<CollectionState, CollectionError> {
400 let id = self.definition.item_id(item)?;
401 self.state.by_id_root = Some(
402 self.index
403 .insert_link(self.state.by_id_root.as_ref(), &id, cid)
404 .await?,
405 );
406
407 for index in self.definition.key_indexes() {
408 let mut root = self.state.key_root(index.name()).cloned();
409 for key in index.materialize_keys(item) {
410 root = Some(self.index.insert_link(root.as_ref(), &key, cid).await?);
411 }
412 self.state.key_roots.insert(index.name().to_string(), root);
413 }
414
415 let mut search_root_groups = BTreeMap::<String, Option<Cid>>::new();
416 let entry_context = CollectionEntryContext {
417 id: &id,
418 cid: Some(cid),
419 write_context: context,
420 };
421 for index in self.definition.search_indexes() {
422 let Some(search_index) = self.search_indexes.get(index.name()) else {
423 continue;
424 };
425
426 let root_name = index.root_name().unwrap_or(index.name()).to_string();
427 let mut root = search_root_groups
428 .get(&root_name)
429 .cloned()
430 .unwrap_or_else(|| self.read_search_root_group(&root_name));
431
432 for entry in index.materialize_entries(item, &entry_context) {
433 let terms = search_index.parse_keywords(&entry.text);
434 if terms.is_empty() {
435 continue;
436 }
437 let entry_id = entry.id.as_deref().unwrap_or(&id);
438 let Some(entry_cid) = entry.cid.as_ref().or(Some(cid)) else {
439 continue;
440 };
441 let default_prefix = default_search_prefix(index.name());
442 root = Some(
443 search_index
444 .index_link(
445 root.as_ref(),
446 entry
447 .prefix
448 .as_deref()
449 .or_else(|| index.prefix())
450 .unwrap_or(&default_prefix),
451 &terms,
452 entry_id,
453 entry_cid,
454 )
455 .await?,
456 );
457 }
458
459 search_root_groups.insert(root_name, root);
460 }
461
462 if !search_root_groups.is_empty() {
463 self.assign_search_root_groups(&search_root_groups);
464 }
465
466 Ok(self.snapshot())
467 }
468
469 async fn delete_normalized_with_context(
470 &mut self,
471 item: &T,
472 context: Option<&CollectionWriteContext>,
473 ) -> Result<CollectionState, CollectionError> {
474 let id = self.definition.item_id(item)?;
475 if let Some(root) = self.state.by_id_root.as_ref() {
476 self.state.by_id_root = self.index.delete(root, &id).await?;
477 }
478
479 for index in self.definition.key_indexes() {
480 let mut root = self.state.key_root(index.name()).cloned();
481 for key in index.materialize_keys(item) {
482 let Some(active_root) = root.as_ref() else {
483 break;
484 };
485 root = self.index.delete(active_root, &key).await?;
486 }
487 self.state.key_roots.insert(index.name().to_string(), root);
488 }
489
490 let mut search_root_groups = BTreeMap::<String, Option<Cid>>::new();
491 let entry_context = CollectionEntryContext {
492 id: &id,
493 cid: None,
494 write_context: context,
495 };
496 for index in self.definition.search_indexes() {
497 let Some(search_index) = self.search_indexes.get(index.name()) else {
498 continue;
499 };
500
501 let root_name = index.root_name().unwrap_or(index.name()).to_string();
502 let mut root = search_root_groups
503 .get(&root_name)
504 .cloned()
505 .unwrap_or_else(|| self.read_search_root_group(&root_name));
506 let Some(existing_root) = root.clone() else {
507 continue;
508 };
509 root = Some(existing_root);
510
511 for entry in index.materialize_entries(item, &entry_context) {
512 let terms = search_index.parse_keywords(&entry.text);
513 if terms.is_empty() {
514 continue;
515 }
516 let entry_id = entry.id.as_deref().unwrap_or(&id);
517 let Some(active_root) = root.as_ref() else {
518 break;
519 };
520 let default_prefix = default_search_prefix(index.name());
521 root = search_index
522 .remove_link(
523 active_root,
524 entry
525 .prefix
526 .as_deref()
527 .or_else(|| index.prefix())
528 .unwrap_or(&default_prefix),
529 &terms,
530 entry_id,
531 )
532 .await?;
533 }
534
535 search_root_groups.insert(root_name, root);
536 }
537
538 if !search_root_groups.is_empty() {
539 self.assign_search_root_groups(&search_root_groups);
540 } else {
541 for index in self.definition.search_indexes() {
542 let root_name = index.root_name().unwrap_or(index.name()).to_string();
543 self.state.search_roots.insert(
544 index.name().to_string(),
545 self.read_search_root_group(&root_name),
546 );
547 }
548 }
549
550 Ok(self.snapshot())
551 }
552
553 async fn rebuild_without_search(
554 &mut self,
555 final_entries: BTreeMap<String, (T, Cid, Option<CollectionWriteContext>)>,
556 ) -> Result<CollectionState, CollectionError> {
557 let mut by_id = BTreeMap::<String, Cid>::new();
558 let mut key_roots = self
559 .definition
560 .key_indexes()
561 .iter()
562 .map(|index| (index.name().to_string(), BTreeMap::<String, Cid>::new()))
563 .collect::<BTreeMap<_, _>>();
564
565 for (id, (item, cid, _context)) in final_entries {
566 by_id.insert(id, cid.clone());
567 for index in self.definition.key_indexes() {
568 let root = key_roots
569 .get_mut(index.name())
570 .expect("collection key root must exist");
571 for key in index.materialize_keys(&item) {
572 root.insert(key, cid.clone());
573 }
574 }
575 }
576
577 self.state = create_empty_collection_state(&self.definition);
578 self.state.by_id_root = self.index.build_links(by_id).await?;
579 for index in self.definition.key_indexes() {
580 let root = self
581 .index
582 .build_links(key_roots.remove(index.name()).unwrap_or_default())
583 .await?;
584 self.state.key_roots.insert(index.name().to_string(), root);
585 }
586
587 Ok(self.snapshot())
588 }
589
590 async fn rebuild_from_final_entries(
591 &mut self,
592 final_entries: BTreeMap<String, (T, Cid, Option<CollectionWriteContext>)>,
593 ) -> Result<CollectionState, CollectionError> {
594 if self.definition.search_indexes().is_empty() {
595 return self.rebuild_without_search(final_entries).await;
596 }
597
598 let mut by_id = BTreeMap::<String, Cid>::new();
599 let mut key_roots = self
600 .definition
601 .key_indexes()
602 .iter()
603 .map(|index| (index.name().to_string(), BTreeMap::<String, Cid>::new()))
604 .collect::<BTreeMap<_, _>>();
605 let mut search_root_entries = BTreeMap::<String, BTreeMap<String, Cid>>::new();
606
607 for (id, (item, cid, context)) in final_entries {
608 by_id.insert(id.clone(), cid.clone());
609
610 for index in self.definition.key_indexes() {
611 let root = key_roots
612 .get_mut(index.name())
613 .expect("collection key root must exist");
614 for key in index.materialize_keys(&item) {
615 root.insert(key, cid.clone());
616 }
617 }
618
619 let entry_context = CollectionEntryContext {
620 id: &id,
621 cid: Some(&cid),
622 write_context: context.as_ref(),
623 };
624
625 for index in self.definition.search_indexes() {
626 let Some(search_index) = self.search_indexes.get(index.name()) else {
627 continue;
628 };
629
630 let root_name = index.root_name().unwrap_or(index.name()).to_string();
631 let root_entries = search_root_entries.entry(root_name).or_default();
632
633 for entry in index.materialize_entries(&item, &entry_context) {
634 let terms = search_index.parse_keywords(&entry.text);
635 if terms.is_empty() {
636 continue;
637 }
638
639 let entry_id = entry.id.as_deref().unwrap_or(&id);
640 let Some(entry_cid) = entry.cid.as_ref().or(Some(&cid)) else {
641 continue;
642 };
643 let default_prefix = default_search_prefix(index.name());
644 let prefix = entry
645 .prefix
646 .as_deref()
647 .or_else(|| index.prefix())
648 .unwrap_or(&default_prefix);
649
650 for term in terms {
651 root_entries
652 .insert(format!("{prefix}{term}:{entry_id}"), entry_cid.clone());
653 }
654 }
655 }
656 }
657
658 self.state = create_empty_collection_state(&self.definition);
659 self.state.by_id_root = self.index.build_links(by_id).await?;
660 for index in self.definition.key_indexes() {
661 let root = self
662 .index
663 .build_links(key_roots.remove(index.name()).unwrap_or_default())
664 .await?;
665 self.state.key_roots.insert(index.name().to_string(), root);
666 }
667
668 let mut search_root_groups = BTreeMap::<String, Option<Cid>>::new();
669 for (root_name, root_entries) in search_root_entries {
670 let Some(search_index) = self.definition.search_indexes().iter().find_map(|index| {
671 ((index.root_name().unwrap_or(index.name())) == root_name)
672 .then(|| self.search_indexes.get(index.name()))
673 .flatten()
674 }) else {
675 continue;
676 };
677
678 search_root_groups.insert(root_name, search_index.build_links(root_entries).await?);
679 }
680
681 if !search_root_groups.is_empty() {
682 self.assign_search_root_groups(&search_root_groups);
683 }
684
685 Ok(self.snapshot())
686 }
687}