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};
17
18pub const FIRST_USER_ID: u64 = 1_000;
20
21#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
23pub enum DbView {
24 #[default]
26 Current,
27 AsOf(u64),
29 Since(u64),
31 History,
33}
34
35#[derive(Clone, Debug, Default)]
37pub struct Idents {
38 by_keyword: BTreeMap<Keyword, EntityId>,
39 by_id: BTreeMap<EntityId, Keyword>,
40}
41
42impl Idents {
43 pub fn insert(&mut self, keyword: Keyword, id: EntityId) {
45 self.by_id.insert(id, keyword.clone());
46 self.by_keyword.insert(keyword, id);
47 }
48
49 #[must_use]
51 pub fn entid(&self, keyword: &Keyword) -> Option<EntityId> {
52 self.by_keyword.get(keyword).copied()
53 }
54
55 #[must_use]
57 pub fn ident(&self, id: EntityId) -> Option<&Keyword> {
58 self.by_id.get(&id)
59 }
60
61 pub fn iter(&self) -> impl Iterator<Item = (&Keyword, &EntityId)> {
63 self.by_keyword.iter()
64 }
65}
66
67#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
69pub struct AttrStats {
70 pub count: usize,
72 pub distinct_values: usize,
74 pub distinct_entities: usize,
76}
77
78#[derive(Clone, Debug, Default)]
80pub struct PlannerStats {
81 pub per_attr: BTreeMap<AttrId, AttrStats>,
83 pub total_datoms: usize,
85 pub entity_count: usize,
87}
88
89impl PlannerStats {
90 #[must_use]
92 pub fn estimate(&self, e_bound: bool, a: Option<AttrId>, v_bound: bool) -> usize {
93 let attr = a.and_then(|a| self.per_attr.get(&a));
94 match (e_bound, attr) {
95 (true, Some(stats)) => (stats.count / stats.distinct_entities.max(1)).max(1),
97 (true, None) => (self.total_datoms / self.entity_count.max(1)).max(1),
98 (false, Some(stats)) if v_bound => (stats.count / stats.distinct_values.max(1)).max(1),
99 (false, Some(stats)) => stats.count.max(1),
100 (false, None) if a.is_some() => 1,
102 (false, None) => self.total_datoms.max(1),
103 }
104 }
105}
106
107type Index = BTreeMap<Vec<u8>, Datom>;
108
109const ORDERS: [IndexOrder; 4] = [
110 IndexOrder::Eavt,
111 IndexOrder::Aevt,
112 IndexOrder::Avet,
113 IndexOrder::Vaet,
114];
115
116const fn slot(order: IndexOrder) -> usize {
117 match order {
118 IndexOrder::Eavt => 0,
119 IndexOrder::Aevt => 1,
120 IndexOrder::Avet => 2,
121 IndexOrder::Vaet => 3,
122 }
123}
124
125#[must_use]
130pub fn key_prefix(
131 order: IndexOrder,
132 e: Option<EntityId>,
133 a: Option<AttrId>,
134 v: Option<&Value>,
135) -> Vec<u8> {
136 enum C {
137 E,
138 A,
139 V,
140 }
141 let components = match order {
142 IndexOrder::Eavt => [C::E, C::A, C::V],
143 IndexOrder::Aevt => [C::A, C::E, C::V],
144 IndexOrder::Avet => [C::A, C::V, C::E],
145 IndexOrder::Vaet => [C::V, C::A, C::E],
146 };
147 let mut out = Vec::new();
148 for component in components {
149 match component {
150 C::E => match e {
151 Some(e) => e.encode_into(&mut out),
152 None => break,
153 },
154 C::A => match a {
155 Some(a) => a.encode_into(&mut out),
156 None => break,
157 },
158 C::V => match v {
159 Some(v) => v.encode_into(&mut out),
160 None => break,
161 },
162 }
163 }
164 out
165}
166
167#[derive(Clone, Debug, Default)]
169pub struct Db {
170 basis_t: u64,
171 schema: Schema,
172 recorded: Arc<Vec<Datom>>,
173 idents: Arc<Idents>,
174 interner: Arc<KeywordInterner>,
175 view: DbView,
176 indexes: Arc<OnceLock<[Index; 4]>>,
177 stats: Arc<OnceLock<PlannerStats>>,
178}
179
180impl Db {
181 #[must_use]
183 pub fn new(schema: Schema) -> Self {
184 Self {
185 schema,
186 ..Self::default()
187 }
188 }
189
190 #[must_use]
192 pub fn with_naming(mut self, idents: Idents, interner: KeywordInterner) -> Self {
193 self.idents = Arc::new(idents);
194 self.interner = Arc::new(interner);
195 self
196 }
197
198 #[must_use]
200 pub const fn basis_t(&self) -> u64 {
201 self.basis_t
202 }
203
204 #[must_use]
206 pub const fn schema(&self) -> &Schema {
207 &self.schema
208 }
209
210 #[must_use]
212 pub fn idents(&self) -> &Idents {
213 &self.idents
214 }
215
216 #[must_use]
218 pub fn interner(&self) -> &KeywordInterner {
219 &self.interner
220 }
221
222 #[must_use]
224 pub const fn view(&self) -> DbView {
225 self.view
226 }
227
228 #[must_use]
230 pub fn recorded_datoms(&self) -> &[Datom] {
231 &self.recorded
232 }
233
234 #[must_use]
236 pub fn as_of(&self, t: u64) -> Self {
237 self.with_view(DbView::AsOf(t))
238 }
239
240 #[must_use]
242 pub fn since(&self, t: u64) -> Self {
243 self.with_view(DbView::Since(t))
244 }
245
246 #[must_use]
248 pub fn history(&self) -> Self {
249 self.with_view(DbView::History)
250 }
251
252 fn with_view(&self, view: DbView) -> Self {
253 if view == self.view {
254 return self.clone();
255 }
256 Self {
257 view,
258 indexes: Arc::new(OnceLock::new()),
259 stats: Arc::new(OnceLock::new()),
260 ..self.clone()
261 }
262 }
263
264 #[must_use]
266 pub fn tx_range(&self, start: u64, end: Option<u64>) -> Vec<(u64, Vec<Datom>)> {
267 let mut by_t: BTreeMap<u64, Vec<Datom>> = BTreeMap::new();
268 for datom in self.recorded.iter() {
269 let t = datom.tx.sequence();
270 if t >= start && end.is_none_or(|end| t < end) {
271 by_t.entry(t).or_default().push(datom.clone());
272 }
273 }
274 by_t.into_iter().collect()
275 }
276
277 #[must_use]
279 pub fn datoms(&self) -> Vec<Datom> {
280 self.datoms_at(IndexOrder::Eavt).cloned().collect()
281 }
282
283 pub fn datoms_at(&self, order: IndexOrder) -> impl Iterator<Item = &Datom> {
288 self.indexes()[slot(order)].values()
289 }
290
291 pub fn datoms_for_attribute(&self, a: AttrId) -> impl Iterator<Item = &Datom> {
297 let prefix = key_prefix(IndexOrder::Aevt, None, Some(a), None);
298 self.indexes()[slot(IndexOrder::Aevt)]
299 .range(prefix.clone()..)
300 .take_while(move |(key, _)| key.starts_with(&prefix))
301 .map(|(_, datom)| datom)
302 }
303
304 pub fn datoms_prefix<'a>(
306 &'a self,
307 order: IndexOrder,
308 prefix: &'a [u8],
309 ) -> impl Iterator<Item = &'a Datom> {
310 self.indexes()[slot(order)]
311 .range(prefix.to_vec()..)
312 .take_while(move |(key, _)| key.starts_with(prefix))
313 .map(|(_, datom)| datom)
314 }
315
316 pub fn seek_datoms<'a>(
318 &'a self,
319 order: IndexOrder,
320 start: &[u8],
321 ) -> impl Iterator<Item = &'a Datom> {
322 self.indexes()[slot(order)]
323 .range(start.to_vec()..)
324 .map(|(_, datom)| datom)
325 }
326
327 pub fn index_range<'a>(
331 &'a self,
332 a: AttrId,
333 start: Option<&Value>,
334 end: Option<&'a Value>,
335 ) -> impl Iterator<Item = &'a Datom> {
336 let a_prefix = key_prefix(IndexOrder::Avet, None, Some(a), None);
337 let start_key = key_prefix(IndexOrder::Avet, None, Some(a), start);
338 self.indexes()[slot(IndexOrder::Avet)]
339 .range(start_key..)
340 .take_while(move |(key, _)| key.starts_with(&a_prefix))
341 .map(|(_, datom)| datom)
342 .take_while(move |datom| end.is_none_or(|end| datom.v < *end))
343 }
344
345 #[must_use]
347 pub fn values(&self, e: EntityId, a: AttrId) -> Vec<Value> {
348 let prefix = key_prefix(IndexOrder::Eavt, Some(e), Some(a), None);
349 self.datoms_prefix(IndexOrder::Eavt, &prefix)
350 .map(|datom| datom.v.clone())
351 .collect()
352 }
353
354 #[must_use]
356 pub fn lookup(&self, a: AttrId, v: &Value) -> Option<EntityId> {
357 if avet_covered(&self.schema, a) {
358 let prefix = key_prefix(IndexOrder::Avet, None, Some(a), Some(v));
359 self.datoms_prefix(IndexOrder::Avet, &prefix)
360 .next()
361 .map(|datom| datom.e)
362 } else {
363 let prefix = key_prefix(IndexOrder::Aevt, None, Some(a), None);
364 self.datoms_prefix(IndexOrder::Aevt, &prefix)
365 .find(|datom| datom.v == *v)
366 .map(|datom| datom.e)
367 }
368 }
369
370 #[must_use]
374 pub fn with_transaction(&self, t: u64, datoms: &[Datom]) -> Self {
375 debug_assert!(
376 self.view == DbView::Current,
377 "with_transaction applies only to the current view"
378 );
379 let mut next = self.clone();
380 next.basis_t = t;
381 Arc::make_mut(&mut next.recorded).extend_from_slice(datoms);
382 next.indexes = Arc::new(OnceLock::new());
383 next.stats = Arc::new(OnceLock::new());
384 if let Some(parent) = self.indexes.get() {
387 let mut derived = parent.clone();
388 apply_current(&mut derived, datoms.iter(), &self.schema);
389 let _ = next.indexes.set(derived);
390 }
391 next
392 }
393
394 #[must_use]
396 pub fn stats(&self) -> DbStats {
397 let planner = self.planner_stats();
398 DbStats {
399 datoms: planner.total_datoms,
400 entities: planner.entity_count,
401 attributes: planner.per_attr.len(),
402 }
403 }
404
405 #[must_use]
407 pub fn planner_stats(&self) -> &PlannerStats {
408 self.stats.get_or_init(|| {
409 let mut stats = PlannerStats::default();
410 let mut values: BTreeMap<AttrId, BTreeSet<&Value>> = BTreeMap::new();
411 let mut attr_entities: BTreeMap<AttrId, BTreeSet<EntityId>> = BTreeMap::new();
412 let mut entities: BTreeSet<EntityId> = BTreeSet::new();
413 for datom in self.datoms_at(IndexOrder::Eavt) {
414 stats.total_datoms += 1;
415 stats.per_attr.entry(datom.a).or_default().count += 1;
416 values.entry(datom.a).or_default().insert(&datom.v);
417 attr_entities.entry(datom.a).or_default().insert(datom.e);
418 entities.insert(datom.e);
419 }
420 for (a, entry) in &mut stats.per_attr {
421 entry.distinct_values = values.get(a).map_or(0, BTreeSet::len);
422 entry.distinct_entities = attr_entities.get(a).map_or(0, BTreeSet::len);
423 }
424 stats.entity_count = entities.len();
425 stats
426 })
427 }
428
429 fn indexes(&self) -> &[Index; 4] {
430 self.indexes.get_or_init(|| {
431 let mut indexes: [Index; 4] = Default::default();
432 match self.view {
433 DbView::History => {
434 for datom in self.recorded.iter() {
435 if self.schema.get(datom.a).is_some_and(|a| a.no_history) {
436 continue;
437 }
438 insert_datom(&mut indexes, datom, &self.schema, true);
439 }
440 }
441 DbView::Current | DbView::AsOf(_) | DbView::Since(_) => {
442 let cutoff = match self.view {
443 DbView::AsOf(t) => Some(t),
444 _ => None,
445 };
446 let filtered = self
447 .recorded
448 .iter()
449 .filter(|d| cutoff.is_none_or(|t| d.tx.sequence() <= t));
450 apply_current(&mut indexes, filtered, &self.schema);
451 if let DbView::Since(t) = self.view {
452 for index in &mut indexes {
453 index.retain(|_, datom| datom.tx.sequence() > t);
454 }
455 }
456 }
457 }
458 indexes
459 })
460 }
461}
462
463fn apply_current<'a>(
469 indexes: &mut [Index; 4],
470 datoms: impl Iterator<Item = &'a Datom>,
471 schema: &Schema,
472) {
473 for datom in datoms {
474 if datom.added {
475 insert_datom(indexes, datom, schema, false);
476 } else {
477 for order in ORDERS {
478 if covered(schema, order, datom) {
479 let key = key_prefix(order, Some(datom.e), Some(datom.a), Some(&datom.v));
480 indexes[slot(order)].remove(&key);
481 }
482 }
483 }
484 }
485}
486
487fn insert_datom(indexes: &mut [Index; 4], datom: &Datom, schema: &Schema, with_tx: bool) {
488 for order in ORDERS {
489 if covered(schema, order, datom) {
490 let key = if with_tx {
491 datom.key(order)
492 } else {
493 key_prefix(order, Some(datom.e), Some(datom.a), Some(&datom.v))
494 };
495 indexes[slot(order)].insert(key, datom.clone());
496 }
497 }
498}
499
500fn covered(schema: &Schema, order: IndexOrder, datom: &Datom) -> bool {
501 match order {
502 IndexOrder::Eavt | IndexOrder::Aevt => true,
503 IndexOrder::Avet => avet_covered(schema, datom.a),
504 IndexOrder::Vaet => matches!(datom.v, Value::Ref(_)),
505 }
506}
507
508#[must_use]
510pub fn avet_covered(schema: &Schema, a: AttrId) -> bool {
511 schema
512 .get(a)
513 .is_some_and(|attr| attr.indexed || attr.unique.is_some())
514}
515
516#[derive(Clone, Copy, Debug, Eq, PartialEq)]
518pub struct DbStats {
519 pub datoms: usize,
521 pub entities: usize,
523 pub attributes: usize,
525}
526
527#[must_use]
529pub const fn attribute(
530 id: u64,
531 value_type: corium_core::ValueType,
532 cardinality: Cardinality,
533 unique: Option<Unique>,
534) -> corium_core::Attribute {
535 corium_core::Attribute {
536 id: EntityId::new(Partition::Db as u32, id),
537 value_type,
538 cardinality,
539 unique,
540 is_component: false,
541 indexed: unique.is_some(),
542 no_history: false,
543 }
544}
545
546#[cfg(test)]
547mod tests {
548 use super::*;
549 use corium_core::ValueType;
550
551 fn schema() -> Schema {
552 let mut schema = Schema::default();
553 schema.insert(attribute(1, ValueType::Str, Cardinality::One, None));
554 schema.insert(attribute(
555 2,
556 ValueType::Long,
557 Cardinality::One,
558 Some(Unique::Identity),
559 ));
560 schema.insert(attribute(3, ValueType::Ref, Cardinality::Many, None));
561 schema
562 }
563
564 fn attr(id: u64) -> AttrId {
565 EntityId::new(Partition::Db as u32, id)
566 }
567
568 fn entity(id: u64) -> EntityId {
569 EntityId::new(Partition::User as u32, id)
570 }
571
572 fn tx_entity(t: u64) -> EntityId {
573 EntityId::new(Partition::Tx as u32, t)
574 }
575
576 fn datom(e: u64, a: u64, v: Value, t: u64, added: bool) -> Datom {
577 Datom {
578 e: entity(e),
579 a: attr(a),
580 v,
581 tx: tx_entity(t),
582 added,
583 }
584 }
585
586 fn sample() -> Db {
587 Db::new(schema())
588 .with_transaction(
589 1,
590 &[
591 datom(1, 1, Value::Str("alice".into()), 1, true),
592 datom(1, 2, Value::Long(7), 1, true),
593 ],
594 )
595 .with_transaction(
596 2,
597 &[
598 datom(1, 1, Value::Str("alice".into()), 2, false),
599 datom(1, 1, Value::Str("alicia".into()), 2, true),
600 datom(2, 3, Value::Ref(entity(1)), 2, true),
601 ],
602 )
603 }
604
605 #[test]
606 fn current_view_folds_retractions() {
607 let db = sample();
608 assert_eq!(
609 db.values(entity(1), attr(1)),
610 vec![Value::Str("alicia".into())]
611 );
612 assert_eq!(db.stats().datoms, 3);
613 }
614
615 #[test]
616 fn attribute_scan_uses_complete_aevt_coverage() {
617 let db = sample();
618 let datoms = db.datoms_for_attribute(attr(1)).collect::<Vec<_>>();
619 assert_eq!(datoms.len(), 1);
620 assert_eq!(datoms[0].v, Value::Str("alicia".into()));
621 assert!(datoms.iter().all(|datom| datom.a == attr(1)));
622 }
623
624 #[test]
625 fn as_of_reconstructs_past_basis() {
626 let db = sample().as_of(1);
627 assert_eq!(
628 db.values(entity(1), attr(1)),
629 vec![Value::Str("alice".into())]
630 );
631 assert_eq!(db.stats().datoms, 2);
632 }
633
634 #[test]
635 fn since_excludes_older_live_facts() {
636 let db = sample().since(1);
637 assert_eq!(db.values(entity(1), attr(2)), Vec::<Value>::new());
639 assert_eq!(
640 db.values(entity(1), attr(1)),
641 vec![Value::Str("alicia".into())]
642 );
643 }
644
645 #[test]
646 fn history_exposes_assertions_and_retractions() {
647 let db = sample().history();
648 let names: Vec<_> = db
649 .datoms_prefix(
650 IndexOrder::Eavt,
651 &key_prefix(IndexOrder::Eavt, Some(entity(1)), Some(attr(1)), None),
652 )
653 .map(|d| (d.v.clone(), d.added))
654 .collect();
655 assert_eq!(names.len(), 3);
656 assert!(names.contains(&(Value::Str("alice".into()), false)));
657 }
658
659 #[test]
660 fn avet_only_covers_indexed_attributes() {
661 let db = sample();
662 assert_eq!(db.datoms_at(IndexOrder::Avet).count(), 1);
663 assert_eq!(db.datoms_at(IndexOrder::Vaet).count(), 1);
664 }
665
666 #[test]
667 fn index_range_scans_value_bounds() {
668 let db = Db::new(schema()).with_transaction(
669 1,
670 &[
671 datom(1, 2, Value::Long(1), 1, true),
672 datom(2, 2, Value::Long(5), 1, true),
673 datom(3, 2, Value::Long(9), 1, true),
674 ],
675 );
676 let hits: Vec<_> = db
677 .index_range(attr(2), Some(&Value::Long(2)), Some(&Value::Long(9)))
678 .map(|d| d.e)
679 .collect();
680 assert_eq!(hits, vec![entity(2)]);
681 }
682
683 #[test]
684 fn tx_range_groups_by_transaction() {
685 let ranged = sample().tx_range(2, None);
686 assert_eq!(ranged.len(), 1);
687 assert_eq!(ranged[0].0, 2);
688 assert_eq!(ranged[0].1.len(), 3);
689 }
690
691 #[test]
692 fn incremental_indexes_match_rebuilt_indexes() {
693 let base = Db::new(schema());
694 let tx1 = [datom(1, 1, Value::Str("a".into()), 1, true)];
695 let tx2 = [
696 datom(1, 1, Value::Str("a".into()), 2, false),
697 datom(1, 1, Value::Str("b".into()), 2, true),
698 ];
699 let warm = base.with_transaction(1, &tx1);
701 let _ = warm.datoms();
702 let incremental = warm.with_transaction(2, &tx2);
703 let cold = base.with_transaction(1, &tx1).with_transaction(2, &tx2);
704 assert_eq!(incremental.datoms(), cold.datoms());
705 }
706
707 #[test]
708 fn planner_stats_count_attributes() {
709 let stats_owner = sample();
710 let stats = stats_owner.planner_stats();
711 assert_eq!(stats.total_datoms, 3);
712 assert_eq!(stats.per_attr[&attr(1)].count, 1);
713 assert!(stats.estimate(false, Some(attr(1)), false) <= stats.total_datoms);
714 }
715}