1use super::Immutable;
4use crate::{
5 Context,
6 journal::{authenticated, contiguous::Mutable},
7 merkle::{Family, Location, Proof},
8 qmdb::{
9 Error,
10 any::{ValueEncoding, batch::lookup_sorted},
11 chain::{self, Bounds, Commitment, Compatible},
12 immutable::operation::Operation,
13 operation::Key,
14 },
15 translator::Translator,
16};
17use commonware_codec::EncodeShared;
18use commonware_cryptography::{Digest, Hasher};
19use commonware_parallel::Strategy;
20use commonware_utils::iter::zip_eq;
21use std::{
22 collections::BTreeMap,
23 iter,
24 sync::{Arc, Weak},
25};
26
27type DiffVec<K, F, V> = Vec<(K, DiffEntry<F, V>)>;
28
29#[derive(Clone)]
31pub(crate) struct DiffEntry<F: Family, V> {
32 pub(crate) value: V,
33 pub(crate) loc: Location<F>,
34}
35
36fn lookup_active<'a, K: Ord, F: Family, V>(
41 diff: &'a [(K, DiffEntry<F, V>)],
42 key: &K,
43 floor: Location<F>,
44) -> Option<&'a DiffEntry<F, V>> {
45 lookup_sorted(diff, key).filter(|entry| entry.loc >= floor)
46}
47
48type MerkleizeResult<F, D, K, V, S> = Result<Arc<MerkleizedBatch<F, D, K, V, S>>, Error<F>>;
50
51#[allow(clippy::type_complexity)]
57pub struct UnmerkleizedBatch<F, H, K, V, S: Strategy>
58where
59 F: Family,
60 K: Key,
61 V: ValueEncoding,
62 H: Hasher,
63{
64 journal_batch: authenticated::UnmerkleizedBatch<F, H, Operation<F, K, V>, S>,
66
67 mutations: BTreeMap<K, V::Value>,
69
70 parent: Option<Arc<MerkleizedBatch<F, H::Digest, K, V, S>>>,
72
73 base: Commitment<F, H::Digest>,
76}
77
78type JournalBatch<F, D, K, V, S> = Arc<authenticated::MerkleizedBatch<F, D, Operation<F, K, V>, S>>;
80
81#[derive(Clone)]
86pub struct MerkleizedBatch<F: Family, D: Digest, K: Key, V: ValueEncoding, S: Strategy> {
87 pub(super) journal_batch: JournalBatch<F, D, K, V, S>,
89
90 pub(super) diff: Arc<DiffVec<K, F, V::Value>>,
93
94 pub(super) parent: Option<Weak<Self>>,
96
97 pub(super) ancestor_diffs: Vec<Arc<DiffVec<K, F, V::Value>>>,
101
102 pub(super) bounds: chain::Bounds<F, D>,
104}
105
106impl<F, H, K, V, S: Strategy> UnmerkleizedBatch<F, H, K, V, S>
107where
108 F: Family,
109 K: Key,
110 V: ValueEncoding,
111 H: Hasher,
112 Operation<F, K, V>: EncodeShared,
113{
114 pub(super) fn new<E, C, T>(
116 immutable: &Immutable<F, E, K, V, C, H, T, S>,
117 base: Commitment<F, H::Digest>,
118 ) -> Self
119 where
120 E: Context,
121 C: Mutable<Item = Operation<F, K, V>>,
122 C::Item: EncodeShared,
123 T: Translator,
124 {
125 Self {
126 journal_batch: immutable.journal.new_batch(),
127 mutations: BTreeMap::new(),
128 parent: None,
129 base,
130 }
131 }
132
133 fn db(&self) -> Commitment<F, H::Digest> {
137 self.parent
138 .as_ref()
139 .map_or(self.base, |parent| parent.bounds.db)
140 }
141
142 fn floor<E, C, T>(&self, db: &Immutable<F, E, K, V, C, H, T, S>) -> Location<F>
145 where
146 E: Context,
147 C: Mutable<Item = Operation<F, K, V>>,
148 C::Item: EncodeShared,
149 T: Translator,
150 {
151 self.parent.as_ref().map_or_else(
152 || db.inactivity_floor_loc(),
153 |parent| parent.bounds.inactivity_floor,
154 )
155 }
156
157 #[allow(clippy::type_complexity)]
160 fn compatible<'a, E, C, T>(
161 &self,
162 db: &'a Immutable<F, E, K, V, C, H, T, S>,
163 ) -> Result<Compatible<'a, Immutable<F, E, K, V, C, H, T, S>>, Error<F>>
164 where
165 E: Context,
166 C: Mutable<Item = Operation<F, K, V>>,
167 C::Item: EncodeShared,
168 T: Translator,
169 {
170 self.parent.as_ref().map_or_else(
171 || self.base.compatible(db, db.commitment()),
172 |parent| parent.bounds.compatible(db, db.commitment()),
173 )
174 }
175
176 pub fn set(mut self, key: K, value: V::Value) -> Self {
180 self.mutations.insert(key, value);
181 self
182 }
183
184 pub async fn get<E, C, T>(
193 &self,
194 key: &K,
195 db: &Immutable<F, E, K, V, C, H, T, S>,
196 ) -> Result<Option<V::Value>, Error<F>>
197 where
198 E: Context,
199 C: Mutable<Item = Operation<F, K, V>>,
200 C::Item: EncodeShared,
201 T: Translator,
202 {
203 let db = self.compatible(db)?;
204 if let Some(value) = self.mutations.get(key) {
206 return Ok(Some(value.clone()));
207 }
208 let floor = self.floor(&*db);
209 if let Some(parent) = self.parent.as_ref() {
210 for diff in parent.read_diffs(db.bounds().end) {
211 if let Some(entry) = lookup_active(diff, key, floor) {
212 return Ok(Some(entry.value.clone()));
213 }
214 }
215 }
216 db.get_from(key, floor).await
218 }
219
220 pub async fn get_many<E, C, T>(
228 &self,
229 keys: &[&K],
230 db: &Immutable<F, E, K, V, C, H, T, S>,
231 ) -> Result<Vec<Option<V::Value>>, Error<F>>
232 where
233 E: Context,
234 C: Mutable<Item = Operation<F, K, V>>,
235 C::Item: EncodeShared,
236 T: Translator,
237 {
238 let db = self.compatible(db)?;
239 if keys.is_empty() {
240 return Ok(Vec::new());
241 }
242 let floor = self.floor(&*db);
243
244 let mut results: Vec<Option<V::Value>> = Vec::with_capacity(keys.len());
245 let mut db_indices = Vec::new();
246 let mut db_keys = Vec::new();
247
248 for (i, key) in keys.iter().enumerate() {
249 if let Some(value) = self.mutations.get(*key) {
251 results.push(Some(value.clone()));
252 continue;
253 }
254
255 if let Some(entry) = self.parent.as_ref().and_then(|parent| {
257 parent
258 .read_diffs(db.bounds().end)
259 .find_map(|diff| lookup_active(diff, *key, floor))
260 }) {
261 results.push(Some(entry.value.clone()));
262 continue;
263 }
264
265 db_indices.push(i);
267 db_keys.push(*key);
268 results.push(None);
269 }
270
271 if !db_keys.is_empty() {
272 let db_results = db.get_many_from(&db_keys, floor).await?;
273 for (slot, value) in zip_eq(db_indices, db_results) {
274 results[slot] = value;
275 }
276 }
277
278 Ok(results)
279 }
280
281 #[tracing::instrument(name = "qmdb.immutable.batch.merkleize", level = "info", skip_all)]
294 pub async fn merkleize<E, C, T>(
295 self,
296 db: &Immutable<F, E, K, V, C, H, T, S>,
297 metadata: Option<V::Value>,
298 inactivity_floor: Location<F>,
299 ) -> MerkleizeResult<F, H::Digest, K, V, S>
300 where
301 E: Context,
302 C: Mutable<Item = Operation<F, K, V>>,
303 C::Item: EncodeShared,
304 T: Translator,
305 {
306 let base = self.base.size;
307
308 let live_ancestors: Vec<_> =
309 chain::parent_and_ancestors(self.parent.as_ref(), |parent| parent.ancestors())
310 .collect();
311 let boundary = chain::effective_boundary(
312 self.db(),
313 live_ancestors.last().map(|oldest| oldest.bounds.base),
314 );
315
316 let mut ancestor_diffs = Vec::new();
318 let mut ancestors = Vec::new();
319 for batch in &live_ancestors {
320 ancestor_diffs.push(Arc::clone(&batch.diff));
321 ancestors.push(batch.commitment());
322 }
323 let db = chain::merkleizable(db, db.commitment(), boundary, ancestors.iter().copied())?;
324 let start_floor = self.floor(&*db);
325
326 let mut ops: Vec<Operation<F, K, V>> = Vec::with_capacity(self.mutations.len() + 1);
329 let mut diff: DiffVec<K, F, V::Value> = Vec::with_capacity(self.mutations.len());
330
331 for (key, value) in self.mutations {
332 let loc = base + ops.len() as u64;
333 ops.push(Operation::Set(key.clone(), value.clone()));
334 diff.push((key, DiffEntry { value, loc }));
335 }
336
337 ops.push(Operation::Commit(metadata, inactivity_floor));
338
339 let total_size = base + ops.len() as u64;
340 chain::validate_merkleize_floor::<F>(start_floor, inactivity_floor, total_size - 1)?;
341 let inactive_peaks = F::inactive_peaks(total_size, inactivity_floor);
342
343 let (journal, root) = db
346 .journal
347 .merkleize(self.journal_batch, ops, inactive_peaks)
348 .await?;
349
350 drop(live_ancestors);
352
353 Ok(Arc::new(MerkleizedBatch {
354 journal_batch: journal,
355 diff: Arc::new(diff),
356 parent: self.parent.as_ref().map(Arc::downgrade),
357 ancestor_diffs,
358 bounds: chain::Bounds {
359 base: self.base,
360 db: boundary,
361 tip: Commitment::new(total_size, root),
362 ancestors,
363 inactivity_floor,
364 },
365 }))
366 }
367}
368
369impl<F: Family, D: Digest, K: Key, V: ValueEncoding, S: Strategy> MerkleizedBatch<F, D, K, V, S>
370where
371 Operation<F, K, V>: EncodeShared,
372{
373 pub const fn root(&self) -> D {
375 self.bounds.tip.root
376 }
377
378 pub const fn bounds(&self) -> &Bounds<F, D> {
380 &self.bounds
381 }
382
383 #[allow(clippy::type_complexity)]
385 pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, K, V>>>) {
386 (
387 self.bounds.base.size,
388 Arc::clone(self.journal_batch.items()),
389 )
390 }
391
392 pub fn proof<E, C, H, T>(
409 &self,
410 db: &Immutable<F, E, K, V, C, H, T, S>,
411 ) -> Result<Proof<F, D>, Error<F>>
412 where
413 E: Context,
414 C: Mutable<Item = Operation<F, K, V>>,
415 H: Hasher<Digest = D>,
416 T: Translator,
417 {
418 let db = self.bounds.compatible(db, db.commitment())?;
419 let inactive_peaks = F::inactive_peaks(self.bounds.tip.size, self.bounds.inactivity_floor);
420 db.journal
421 .speculative_proof(&self.journal_batch, inactive_peaks)
422 .map_err(Into::into)
423 }
424
425 pub fn pinned_nodes<E, C, H, T>(
441 &self,
442 db: &Immutable<F, E, K, V, C, H, T, S>,
443 ) -> Result<Vec<D>, Error<F>>
444 where
445 E: Context,
446 C: Mutable<Item = Operation<F, K, V>>,
447 H: Hasher<Digest = D>,
448 T: Translator,
449 {
450 let db = self.bounds.compatible(db, db.commitment())?;
451 db.journal
452 .speculative_pinned_nodes(&self.journal_batch)
453 .map_err(Into::into)
454 }
455
456 fn read_diffs(
464 &self,
465 db_size: Location<F>,
466 ) -> impl Iterator<Item = &[(K, DiffEntry<F, V::Value>)]> {
467 let mut reachable = iter::once_with(|| self.ancestors()).flatten();
468 iter::once(self.diff.as_slice()).chain(
469 self.ancestor_diffs
470 .iter()
471 .zip(&self.bounds.ancestors)
472 .take_while(move |(_, state)| reachable.next().is_some() || state.size > db_size)
473 .map(|(diff, _)| diff.as_slice()),
474 )
475 }
476
477 pub(super) fn ancestors(&self) -> impl Iterator<Item = Arc<Self>> + use<F, D, K, V, S> {
479 chain::ancestors(self.parent.clone(), |batch| batch.parent.as_ref())
480 }
481
482 pub(super) const fn commitment(&self) -> Commitment<F, D> {
484 self.bounds.tip
485 }
486
487 pub async fn get<E, C, H, T>(
496 &self,
497 key: &K,
498 db: &Immutable<F, E, K, V, C, H, T, S>,
499 ) -> Result<Option<V::Value>, Error<F>>
500 where
501 E: Context,
502 C: Mutable<Item = Operation<F, K, V>>,
503 C::Item: EncodeShared,
504 H: Hasher<Digest = D>,
505 T: Translator,
506 {
507 let db = self.bounds.compatible(db, db.commitment())?;
508 let floor = self.bounds.inactivity_floor;
509 for diff in self.read_diffs(db.bounds().end) {
510 if let Some(entry) = lookup_active(diff, key, floor) {
511 return Ok(Some(entry.value.clone()));
512 }
513 }
514 db.get_from(key, floor).await
515 }
516
517 pub async fn get_many<E, C, H, T>(
525 &self,
526 keys: &[&K],
527 db: &Immutable<F, E, K, V, C, H, T, S>,
528 ) -> Result<Vec<Option<V::Value>>, Error<F>>
529 where
530 E: Context,
531 C: Mutable<Item = Operation<F, K, V>>,
532 C::Item: EncodeShared,
533 H: Hasher<Digest = D>,
534 T: Translator,
535 {
536 let db = self.bounds.compatible(db, db.commitment())?;
537 if keys.is_empty() {
538 return Ok(Vec::new());
539 }
540 let floor = self.bounds.inactivity_floor;
541
542 let mut results: Vec<Option<V::Value>> = Vec::with_capacity(keys.len());
543 let mut db_indices = Vec::new();
544 let mut db_keys = Vec::new();
545
546 for (i, key) in keys.iter().enumerate() {
547 if let Some(entry) = self
548 .read_diffs(db.bounds().end)
549 .find_map(|diff| lookup_active(diff, *key, floor))
550 {
551 results.push(Some(entry.value.clone()));
552 continue;
553 }
554
555 db_indices.push(i);
557 db_keys.push(*key);
558 results.push(None);
559 }
560
561 if !db_keys.is_empty() {
562 let db_results = db.get_many_from(&db_keys, floor).await?;
563 for (slot, value) in zip_eq(db_indices, db_results) {
564 results[slot] = value;
565 }
566 }
567
568 Ok(results)
569 }
570
571 pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, K, V, S>
580 where
581 H: Hasher<Digest = D>,
582 {
583 UnmerkleizedBatch {
584 journal_batch: self.journal_batch.new_batch::<H>(),
585 mutations: BTreeMap::new(),
586 parent: Some(Arc::clone(self)),
587 base: self.commitment(),
588 }
589 }
590}
591
592impl<F, E, K, V, C, H, T, S> Immutable<F, E, K, V, C, H, T, S>
593where
594 F: Family,
595 E: Context,
596 K: Key,
597 V: ValueEncoding,
598 C: Mutable<Item = Operation<F, K, V>>,
599 C::Item: EncodeShared,
600 H: Hasher,
601 T: Translator,
602 S: Strategy,
603{
604 pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, K, V, S>> {
606 Arc::new(MerkleizedBatch {
607 journal_batch: self.journal.to_merkleized_batch(),
608 diff: Arc::new(Vec::new()),
609 parent: None,
610 ancestor_diffs: Vec::new(),
611 bounds: chain::Bounds::from_db(self.commitment(), self.inactivity_floor_loc),
612 })
613 }
614}