1use std::collections::{BTreeMap, BTreeSet};
11use std::sync::{Arc, OnceLock};
12
13use corium_core::{
14 AttrId, Cardinality, Datom, EntityId, IndexOrder, Keyword, KeywordInterner, Partition, Schema,
15 Unique, Value, encoding::Encodable,
16};
17use rpds::{RedBlackTreeMapSync, VectorSync};
18
19pub const FIRST_USER_ID: u64 = 1_000;
21
22#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
24pub enum DbView {
25 #[default]
27 Current,
28 AsOf(u64),
30 Since(u64),
32 History,
34}
35
36#[derive(Clone, Debug, Default)]
38pub struct Idents {
39 by_keyword: BTreeMap<Keyword, EntityId>,
40 by_id: BTreeMap<EntityId, Keyword>,
41}
42
43impl Idents {
44 pub fn insert(&mut self, keyword: Keyword, id: EntityId) {
46 self.by_id.insert(id, keyword.clone());
47 self.by_keyword.insert(keyword, id);
48 }
49
50 #[must_use]
52 pub fn entid(&self, keyword: &Keyword) -> Option<EntityId> {
53 self.by_keyword.get(keyword).copied()
54 }
55
56 #[must_use]
58 pub fn ident(&self, id: EntityId) -> Option<&Keyword> {
59 self.by_id.get(&id)
60 }
61
62 pub fn iter(&self) -> impl Iterator<Item = (&Keyword, &EntityId)> {
64 self.by_keyword.iter()
65 }
66}
67
68#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
70pub struct AttrStats {
71 pub count: usize,
73 pub distinct_values: usize,
75 pub distinct_entities: usize,
77}
78
79#[derive(Clone, Debug, Default)]
81pub struct PlannerStats {
82 pub per_attr: BTreeMap<AttrId, AttrStats>,
84 pub total_datoms: usize,
86 pub entity_count: usize,
88}
89
90impl PlannerStats {
91 #[must_use]
93 pub fn estimate(&self, e_bound: bool, a: Option<AttrId>, v_bound: bool) -> usize {
94 let attr = a.and_then(|a| self.per_attr.get(&a));
95 match (e_bound, attr) {
96 (true, Some(stats)) => (stats.count / stats.distinct_entities.max(1)).max(1),
98 (true, None) => (self.total_datoms / self.entity_count.max(1)).max(1),
99 (false, Some(stats)) if v_bound => (stats.count / stats.distinct_values.max(1)).max(1),
100 (false, Some(stats)) => stats.count.max(1),
101 (false, None) if a.is_some() => 1,
103 (false, None) => self.total_datoms.max(1),
104 }
105 }
106}
107
108type Index = RedBlackTreeMapSync<Vec<u8>, Datom>;
117
118const ORDERS: [IndexOrder; 4] = [
119 IndexOrder::Eavt,
120 IndexOrder::Aevt,
121 IndexOrder::Avet,
122 IndexOrder::Vaet,
123];
124
125const fn slot(order: IndexOrder) -> usize {
126 match order {
127 IndexOrder::Eavt => 0,
128 IndexOrder::Aevt => 1,
129 IndexOrder::Avet => 2,
130 IndexOrder::Vaet => 3,
131 }
132}
133
134#[must_use]
139pub fn key_prefix(
140 order: IndexOrder,
141 e: Option<EntityId>,
142 a: Option<AttrId>,
143 v: Option<&Value>,
144) -> Vec<u8> {
145 enum C {
146 E,
147 A,
148 V,
149 }
150 let components = match order {
151 IndexOrder::Eavt => [C::E, C::A, C::V],
152 IndexOrder::Aevt => [C::A, C::E, C::V],
153 IndexOrder::Avet => [C::A, C::V, C::E],
154 IndexOrder::Vaet => [C::V, C::A, C::E],
155 };
156 let mut out = Vec::new();
157 for component in components {
158 match component {
159 C::E => match e {
160 Some(e) => e.encode_into(&mut out),
161 None => break,
162 },
163 C::A => match a {
164 Some(a) => a.encode_into(&mut out),
165 None => break,
166 },
167 C::V => match v {
168 Some(v) => v.encode_into(&mut out),
169 None => break,
170 },
171 }
172 }
173 out
174}
175
176#[derive(Clone, Debug, Default)]
178pub struct Db {
179 basis_t: u64,
180 schema: Schema,
181 recorded: VectorSync<Datom>,
182 idents: Arc<Idents>,
183 interner: Arc<KeywordInterner>,
184 view: DbView,
185 indexes: Arc<OnceLock<[Index; 4]>>,
186 stats: Arc<OnceLock<PlannerStats>>,
187}
188
189impl Db {
190 #[must_use]
192 pub fn new(schema: Schema) -> Self {
193 Self {
194 schema,
195 ..Self::default()
196 }
197 }
198
199 #[must_use]
207 pub fn from_current_snapshot(
208 basis_t: u64,
209 schema: Schema,
210 idents: Idents,
211 interner: KeywordInterner,
212 datoms: Vec<Datom>,
213 ) -> Self {
214 Self {
215 basis_t,
216 schema,
217 recorded: datoms.into_iter().collect(),
218 idents: Arc::new(idents),
219 interner: Arc::new(interner),
220 view: DbView::Current,
221 indexes: Arc::new(OnceLock::new()),
222 stats: Arc::new(OnceLock::new()),
223 }
224 }
225
226 #[must_use]
228 pub fn with_naming(mut self, idents: Idents, interner: KeywordInterner) -> Self {
229 self.idents = Arc::new(idents);
230 self.interner = Arc::new(interner);
231 self
232 }
233
234 #[must_use]
236 pub const fn basis_t(&self) -> u64 {
237 self.basis_t
238 }
239
240 #[must_use]
242 pub const fn schema(&self) -> &Schema {
243 &self.schema
244 }
245
246 #[must_use]
248 pub fn idents(&self) -> &Idents {
249 &self.idents
250 }
251
252 #[must_use]
254 pub fn interner(&self) -> &KeywordInterner {
255 &self.interner
256 }
257
258 #[must_use]
260 pub const fn view(&self) -> DbView {
261 self.view
262 }
263
264 pub fn recorded_datoms(&self) -> impl Iterator<Item = &Datom> {
266 self.recorded.iter()
267 }
268
269 #[must_use]
271 pub fn recorded_len(&self) -> usize {
272 self.recorded.len()
273 }
274
275 #[must_use]
277 pub fn as_of(&self, t: u64) -> Self {
278 self.with_view(DbView::AsOf(t))
279 }
280
281 #[must_use]
283 pub fn since(&self, t: u64) -> Self {
284 self.with_view(DbView::Since(t))
285 }
286
287 #[must_use]
289 pub fn history(&self) -> Self {
290 self.with_view(DbView::History)
291 }
292
293 fn with_view(&self, view: DbView) -> Self {
294 if view == self.view {
295 return self.clone();
296 }
297 Self {
298 view,
299 indexes: Arc::new(OnceLock::new()),
300 stats: Arc::new(OnceLock::new()),
301 ..self.clone()
302 }
303 }
304
305 #[must_use]
307 pub fn tx_range(&self, start: u64, end: Option<u64>) -> Vec<(u64, Vec<Datom>)> {
308 let mut by_t: BTreeMap<u64, Vec<Datom>> = BTreeMap::new();
309 for datom in &self.recorded {
310 let t = datom.tx.sequence();
311 if t >= start && end.is_none_or(|end| t < end) {
312 by_t.entry(t).or_default().push(datom.clone());
313 }
314 }
315 by_t.into_iter().collect()
316 }
317
318 #[must_use]
320 pub fn datoms(&self) -> Vec<Datom> {
321 self.datoms_at(IndexOrder::Eavt).cloned().collect()
322 }
323
324 pub fn datoms_at(&self, order: IndexOrder) -> impl Iterator<Item = &Datom> {
329 self.indexes()[slot(order)].values()
330 }
331
332 pub fn datoms_for_attribute(&self, a: AttrId) -> impl Iterator<Item = &Datom> {
338 let prefix = key_prefix(IndexOrder::Aevt, None, Some(a), None);
339 self.indexes()[slot(IndexOrder::Aevt)]
340 .range(prefix.clone()..)
341 .take_while(move |(key, _)| key.starts_with(&prefix))
342 .map(|(_, datom)| datom)
343 }
344
345 pub fn datoms_prefix<'a>(
347 &'a self,
348 order: IndexOrder,
349 prefix: &'a [u8],
350 ) -> impl Iterator<Item = &'a Datom> {
351 self.indexes()[slot(order)]
352 .range(prefix.to_vec()..)
353 .take_while(move |(key, _)| key.starts_with(prefix))
354 .map(|(_, datom)| datom)
355 }
356
357 pub fn seek_datoms<'a>(
359 &'a self,
360 order: IndexOrder,
361 start: &[u8],
362 ) -> impl Iterator<Item = &'a Datom> {
363 self.indexes()[slot(order)]
364 .range(start.to_vec()..)
365 .map(|(_, datom)| datom)
366 }
367
368 pub fn index_range<'a>(
372 &'a self,
373 a: AttrId,
374 start: Option<&Value>,
375 end: Option<&'a Value>,
376 ) -> impl Iterator<Item = &'a Datom> {
377 let a_prefix = key_prefix(IndexOrder::Avet, None, Some(a), None);
378 let start_key = key_prefix(IndexOrder::Avet, None, Some(a), start);
379 self.indexes()[slot(IndexOrder::Avet)]
380 .range(start_key..)
381 .take_while(move |(key, _)| key.starts_with(&a_prefix))
382 .map(|(_, datom)| datom)
383 .take_while(move |datom| end.is_none_or(|end| datom.v < *end))
384 }
385
386 #[must_use]
388 pub fn values(&self, e: EntityId, a: AttrId) -> Vec<Value> {
389 let prefix = key_prefix(IndexOrder::Eavt, Some(e), Some(a), None);
390 self.datoms_prefix(IndexOrder::Eavt, &prefix)
391 .map(|datom| datom.v.clone())
392 .collect()
393 }
394
395 #[must_use]
397 pub fn lookup(&self, a: AttrId, v: &Value) -> Option<EntityId> {
398 if avet_covered(&self.schema, a) {
399 let prefix = key_prefix(IndexOrder::Avet, None, Some(a), Some(v));
400 self.datoms_prefix(IndexOrder::Avet, &prefix)
401 .next()
402 .map(|datom| datom.e)
403 } else {
404 let prefix = key_prefix(IndexOrder::Aevt, None, Some(a), None);
405 self.datoms_prefix(IndexOrder::Aevt, &prefix)
406 .find(|datom| datom.v == *v)
407 .map(|datom| datom.e)
408 }
409 }
410
411 #[must_use]
415 pub fn with_transaction(&self, t: u64, datoms: &[Datom]) -> Self {
416 debug_assert!(
417 self.view == DbView::Current,
418 "with_transaction applies only to the current view"
419 );
420 let mut next = self.clone();
421 next.basis_t = t;
422 for datom in datoms {
427 next.recorded.push_back_mut(datom.clone());
428 }
429 next.indexes = Arc::new(OnceLock::new());
430 next.stats = Arc::new(OnceLock::new());
431 if let Some(parent) = self.indexes.get() {
434 let mut derived = parent.clone();
435 apply_current(&mut derived, datoms.iter(), &self.schema);
436 let _ = next.indexes.set(derived);
437 }
438 next
439 }
440
441 #[must_use]
443 pub fn stats(&self) -> DbStats {
444 let planner = self.planner_stats();
445 DbStats {
446 datoms: planner.total_datoms,
447 entities: planner.entity_count,
448 attributes: planner.per_attr.len(),
449 }
450 }
451
452 #[must_use]
454 pub fn planner_stats(&self) -> &PlannerStats {
455 self.stats.get_or_init(|| {
456 let mut stats = PlannerStats::default();
457 let mut values: BTreeMap<AttrId, BTreeSet<&Value>> = BTreeMap::new();
458 let mut attr_entities: BTreeMap<AttrId, BTreeSet<EntityId>> = BTreeMap::new();
459 let mut entities: BTreeSet<EntityId> = BTreeSet::new();
460 for datom in self.datoms_at(IndexOrder::Eavt) {
461 stats.total_datoms += 1;
462 stats.per_attr.entry(datom.a).or_default().count += 1;
463 values.entry(datom.a).or_default().insert(&datom.v);
464 attr_entities.entry(datom.a).or_default().insert(datom.e);
465 entities.insert(datom.e);
466 }
467 for (a, entry) in &mut stats.per_attr {
468 entry.distinct_values = values.get(a).map_or(0, BTreeSet::len);
469 entry.distinct_entities = attr_entities.get(a).map_or(0, BTreeSet::len);
470 }
471 stats.entity_count = entities.len();
472 stats
473 })
474 }
475
476 fn indexes(&self) -> &[Index; 4] {
477 self.indexes.get_or_init(|| {
478 let mut indexes: [Index; 4] = Default::default();
479 match self.view {
480 DbView::History => {
481 for datom in &self.recorded {
482 if self.schema.get(datom.a).is_some_and(|a| a.no_history) {
483 continue;
484 }
485 insert_datom(&mut indexes, datom, &self.schema, true);
486 }
487 }
488 DbView::Current | DbView::AsOf(_) | DbView::Since(_) => {
489 let cutoff = match self.view {
490 DbView::AsOf(t) => Some(t),
491 _ => None,
492 };
493 let filtered = self
494 .recorded
495 .iter()
496 .filter(|d| cutoff.is_none_or(|t| d.tx.sequence() <= t));
497 apply_current(&mut indexes, filtered, &self.schema);
498 if let DbView::Since(t) = self.view {
499 for index in &mut indexes {
500 *index = index
501 .iter()
502 .filter(|(_, datom)| datom.tx.sequence() > t)
503 .map(|(key, datom)| (key.clone(), datom.clone()))
504 .collect();
505 }
506 }
507 }
508 }
509 indexes
510 })
511 }
512}
513
514fn apply_current<'a>(
520 indexes: &mut [Index; 4],
521 datoms: impl Iterator<Item = &'a Datom>,
522 schema: &Schema,
523) {
524 for datom in datoms {
525 if datom.added {
526 insert_datom(indexes, datom, schema, false);
527 } else {
528 for order in ORDERS {
529 if covered(schema, order, datom) {
530 let key = key_prefix(order, Some(datom.e), Some(datom.a), Some(&datom.v));
531 indexes[slot(order)].remove_mut(&key);
532 }
533 }
534 }
535 }
536}
537
538fn insert_datom(indexes: &mut [Index; 4], datom: &Datom, schema: &Schema, with_tx: bool) {
539 for order in ORDERS {
540 if covered(schema, order, datom) {
541 let key = if with_tx {
542 datom.key(order)
543 } else {
544 key_prefix(order, Some(datom.e), Some(datom.a), Some(&datom.v))
545 };
546 indexes[slot(order)].insert_mut(key, datom.clone());
547 }
548 }
549}
550
551fn covered(schema: &Schema, order: IndexOrder, datom: &Datom) -> bool {
552 match order {
553 IndexOrder::Eavt | IndexOrder::Aevt => true,
554 IndexOrder::Avet => avet_covered(schema, datom.a),
555 IndexOrder::Vaet => matches!(datom.v, Value::Ref(_)),
556 }
557}
558
559#[must_use]
561pub fn avet_covered(schema: &Schema, a: AttrId) -> bool {
562 schema
563 .get(a)
564 .is_some_and(|attr| attr.indexed || attr.unique.is_some())
565}
566
567#[derive(Clone, Copy, Debug, Eq, PartialEq)]
569pub struct DbStats {
570 pub datoms: usize,
572 pub entities: usize,
574 pub attributes: usize,
576}
577
578#[must_use]
580pub const fn attribute(
581 id: u64,
582 value_type: corium_core::ValueType,
583 cardinality: Cardinality,
584 unique: Option<Unique>,
585) -> corium_core::Attribute {
586 corium_core::Attribute {
587 id: EntityId::new(Partition::Db as u32, id),
588 value_type,
589 cardinality,
590 unique,
591 is_component: false,
592 indexed: unique.is_some(),
593 no_history: false,
594 }
595}
596
597#[cfg(test)]
598mod tests {
599 use super::*;
600 use corium_core::ValueType;
601
602 fn schema() -> Schema {
603 let mut schema = Schema::default();
604 schema.insert(attribute(1, ValueType::Str, Cardinality::One, None));
605 schema.insert(attribute(
606 2,
607 ValueType::Long,
608 Cardinality::One,
609 Some(Unique::Identity),
610 ));
611 schema.insert(attribute(3, ValueType::Ref, Cardinality::Many, None));
612 schema
613 }
614
615 fn attr(id: u64) -> AttrId {
616 EntityId::new(Partition::Db as u32, id)
617 }
618
619 fn entity(id: u64) -> EntityId {
620 EntityId::new(Partition::User as u32, id)
621 }
622
623 fn tx_entity(t: u64) -> EntityId {
624 EntityId::new(Partition::Tx as u32, t)
625 }
626
627 fn datom(e: u64, a: u64, v: Value, t: u64, added: bool) -> Datom {
628 Datom {
629 e: entity(e),
630 a: attr(a),
631 v,
632 tx: tx_entity(t),
633 added,
634 }
635 }
636
637 fn sample() -> Db {
638 Db::new(schema())
639 .with_transaction(
640 1,
641 &[
642 datom(1, 1, Value::Str("alice".into()), 1, true),
643 datom(1, 2, Value::Long(7), 1, true),
644 ],
645 )
646 .with_transaction(
647 2,
648 &[
649 datom(1, 1, Value::Str("alice".into()), 2, false),
650 datom(1, 1, Value::Str("alicia".into()), 2, true),
651 datom(2, 3, Value::Ref(entity(1)), 2, true),
652 ],
653 )
654 }
655
656 #[test]
657 fn current_view_folds_retractions() {
658 let db = sample();
659 assert_eq!(
660 db.values(entity(1), attr(1)),
661 vec![Value::Str("alicia".into())]
662 );
663 assert_eq!(db.stats().datoms, 3);
664 }
665
666 #[test]
667 fn attribute_scan_uses_complete_aevt_coverage() {
668 let db = sample();
669 let datoms = db.datoms_for_attribute(attr(1)).collect::<Vec<_>>();
670 assert_eq!(datoms.len(), 1);
671 assert_eq!(datoms[0].v, Value::Str("alicia".into()));
672 assert!(datoms.iter().all(|datom| datom.a == attr(1)));
673 }
674
675 #[test]
676 fn as_of_reconstructs_past_basis() {
677 let db = sample().as_of(1);
678 assert_eq!(
679 db.values(entity(1), attr(1)),
680 vec![Value::Str("alice".into())]
681 );
682 assert_eq!(db.stats().datoms, 2);
683 }
684
685 #[test]
686 fn since_excludes_older_live_facts() {
687 let db = sample().since(1);
688 assert_eq!(db.values(entity(1), attr(2)), Vec::<Value>::new());
690 assert_eq!(
691 db.values(entity(1), attr(1)),
692 vec![Value::Str("alicia".into())]
693 );
694 }
695
696 #[test]
697 fn history_exposes_assertions_and_retractions() {
698 let db = sample().history();
699 let names: Vec<_> = db
700 .datoms_prefix(
701 IndexOrder::Eavt,
702 &key_prefix(IndexOrder::Eavt, Some(entity(1)), Some(attr(1)), None),
703 )
704 .map(|d| (d.v.clone(), d.added))
705 .collect();
706 assert_eq!(names.len(), 3);
707 assert!(names.contains(&(Value::Str("alice".into()), false)));
708 }
709
710 #[test]
711 fn avet_only_covers_indexed_attributes() {
712 let db = sample();
713 assert_eq!(db.datoms_at(IndexOrder::Avet).count(), 1);
714 assert_eq!(db.datoms_at(IndexOrder::Vaet).count(), 1);
715 }
716
717 #[test]
718 fn index_range_scans_value_bounds() {
719 let db = Db::new(schema()).with_transaction(
720 1,
721 &[
722 datom(1, 2, Value::Long(1), 1, true),
723 datom(2, 2, Value::Long(5), 1, true),
724 datom(3, 2, Value::Long(9), 1, true),
725 ],
726 );
727 let hits: Vec<_> = db
728 .index_range(attr(2), Some(&Value::Long(2)), Some(&Value::Long(9)))
729 .map(|d| d.e)
730 .collect();
731 assert_eq!(hits, vec![entity(2)]);
732 }
733
734 #[test]
735 fn tx_range_groups_by_transaction() {
736 let ranged = sample().tx_range(2, None);
737 assert_eq!(ranged.len(), 1);
738 assert_eq!(ranged[0].0, 2);
739 assert_eq!(ranged[0].1.len(), 3);
740 }
741
742 #[test]
743 fn incremental_indexes_match_rebuilt_indexes() {
744 let base = Db::new(schema());
745 let tx1 = [datom(1, 1, Value::Str("a".into()), 1, true)];
746 let tx2 = [
747 datom(1, 1, Value::Str("a".into()), 2, false),
748 datom(1, 1, Value::Str("b".into()), 2, true),
749 ];
750 let warm = base.with_transaction(1, &tx1);
752 let _ = warm.datoms();
753 let incremental = warm.with_transaction(2, &tx2);
754 let cold = base.with_transaction(1, &tx1).with_transaction(2, &tx2);
755 assert_eq!(incremental.datoms(), cold.datoms());
756 }
757
758 #[test]
759 fn with_transaction_leaves_parent_value_unchanged() {
760 let parent = Db::new(schema())
764 .with_transaction(1, &[datom(1, 1, Value::Str("alice".into()), 1, true)]);
765 let parent_datoms = parent.datoms();
767 let parent_recorded = parent.recorded_len();
768
769 let child = parent.with_transaction(
770 2,
771 &[
772 datom(1, 1, Value::Str("alice".into()), 2, false),
773 datom(1, 1, Value::Str("alicia".into()), 2, true),
774 ],
775 );
776
777 assert_eq!(parent.basis_t(), 1);
779 assert_eq!(parent.recorded_len(), parent_recorded);
780 assert_eq!(parent.datoms(), parent_datoms);
781 assert_eq!(
782 parent.values(entity(1), attr(1)),
783 vec![Value::Str("alice".into())]
784 );
785 assert_eq!(child.basis_t(), 2);
787 assert_eq!(
788 child.values(entity(1), attr(1)),
789 vec![Value::Str("alicia".into())]
790 );
791 }
792
793 #[test]
794 fn planner_stats_count_attributes() {
795 let stats_owner = sample();
796 let stats = stats_owner.planner_stats();
797 assert_eq!(stats.total_datoms, 3);
798 assert_eq!(stats.per_attr[&attr(1)].count, 1);
799 assert!(stats.estimate(false, Some(attr(1)), false) <= stats.total_datoms);
800 }
801}