1#![cfg_attr(not(feature = "std"), no_std)]
18
19#[cfg(not(feature = "std"))]
20extern crate alloc;
21
22mod malloc_size_of;
23pub use malloc_size_of::*;
24
25use tetsy_hash_db::{HashDB, HashDBRef, PlainDB, PlainDBRef, Hasher as KeyHasher,
26 AsHashDB, AsPlainDB, Prefix};
27use tetsy_util_mem::{MallocSizeOf, MallocSizeOfOps, MallocShallowSizeOf};
28#[cfg(feature = "std")]
29use std::{
30 collections::hash_map::Entry,
31 collections::HashMap,
32 hash,
33 mem,
34 marker::PhantomData,
35 cmp::Eq,
36 borrow::Borrow,
37};
38
39#[cfg(not(feature = "std"))]
40use hashbrown::{
41 HashMap,
42 hash_map::Entry,
43};
44
45#[cfg(not(feature = "std"))]
46use core::{
47 hash,
48 mem,
49 marker::PhantomData,
50 cmp::Eq,
51 borrow::Borrow,
52};
53
54#[cfg(not(feature = "std"))]
55use alloc::vec::Vec;
56
57#[cfg(feature = "std")]
58pub trait MaybeDebug: std::fmt::Debug {}
59#[cfg(feature = "std")]
60impl<T: std::fmt::Debug> MaybeDebug for T {}
61#[cfg(not(feature = "std"))]
62pub trait MaybeDebug {}
63#[cfg(not(feature = "std"))]
64impl<T> MaybeDebug for T {}
65
66#[cfg(feature = "std")]
68pub type DefaultMemTracker<T> = MemCounter<T>;
69#[cfg(not(feature = "std"))]
71pub type DefaultMemTracker<T> = NoopTracker<T>;
72
73pub struct MemoryDB<H, KF, T, M = DefaultMemTracker<T>>
119where
120 H: KeyHasher,
121 KF: KeyFunction<H>,
122 M: MemTracker<T>,
123{
124 data: HashMap<KF::Key, (T, i32)>,
125 malloc_tracker: M,
128 hashed_null_node: H::Out,
129 null_node_data: T,
130 _kf: PhantomData<KF>,
131}
132
133impl<H, KF, T, M> Clone for MemoryDB<H, KF, T, M>
134where
135 H: KeyHasher,
136 KF: KeyFunction<H>,
137 T: Clone,
138 M: MemTracker<T> + Copy,
139{
140 fn clone(&self) -> Self {
141 Self {
142 data: self.data.clone(),
143 hashed_null_node: self.hashed_null_node,
144 null_node_data: self.null_node_data.clone(),
145 malloc_tracker: self.malloc_tracker,
146 _kf: Default::default(),
147 }
148 }
149}
150
151impl<H, KF, T, M> PartialEq<MemoryDB<H, KF, T, M>> for MemoryDB<H, KF, T, M>
152 where
153 H: KeyHasher,
154 KF: KeyFunction<H>,
155 <KF as KeyFunction<H>>::Key: Eq + MaybeDebug,
156 T: Eq + MaybeDebug,
157 M: MemTracker<T> + PartialEq,
158{
159 fn eq(&self, other: &MemoryDB<H, KF, T, M>) -> bool {
160 for a in self.data.iter() {
161 match other.data.get(&a.0) {
162 Some(v) if v != a.1 => return false,
163 None => return false,
164 _ => (),
165 }
166 }
167 true
168 }
169}
170
171impl<H, KF, T, M> Eq for MemoryDB<H, KF, T, M>
172 where
173 H: KeyHasher,
174 KF: KeyFunction<H>,
175 <KF as KeyFunction<H>>::Key: Eq + MaybeDebug,
176 T: Eq + MaybeDebug,
177 M: MemTracker<T> + Eq,
178{}
179
180pub trait KeyFunction<H: KeyHasher> {
181 type Key: Send + Sync + Clone + hash::Hash + Eq;
182
183 fn key(hash: &H::Out, prefix: Prefix) -> Self::Key;
184}
185
186pub struct HashKey<H>(PhantomData<H>);
188
189impl<H> Clone for HashKey<H> {
190 fn clone(&self) -> Self {
191 Self(Default::default())
192 }
193}
194
195impl<H> core::fmt::Debug for HashKey<H> {
196 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
197 core::write!(f, "HashKey")
198 }
199}
200
201impl<H: KeyHasher> KeyFunction<H> for HashKey<H> {
202 type Key = H::Out;
203
204 fn key(hash: &H::Out, prefix: Prefix) -> H::Out {
205 hash_key::<H>(hash, prefix)
206 }
207}
208
209pub fn hash_key<H: KeyHasher>(key: &H::Out, _prefix: Prefix) -> H::Out {
211 *key
212}
213
214pub struct PrefixedKey<H>(PhantomData<H>);
216
217impl<H> Clone for PrefixedKey<H> {
218 fn clone(&self) -> Self {
219 Self(Default::default())
220 }
221}
222
223impl<H> core::fmt::Debug for PrefixedKey<H> {
224 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
225 core::write!(f, "PrefixedKey")
226 }
227}
228
229impl<H: KeyHasher> KeyFunction<H> for PrefixedKey<H> {
230 type Key = Vec<u8>;
231
232 fn key(hash: &H::Out, prefix: Prefix) -> Vec<u8> {
233 prefixed_key::<H>(hash, prefix)
234 }
235}
236
237pub fn prefixed_key<H: KeyHasher>(key: &H::Out, prefix: Prefix) -> Vec<u8> {
239 let mut prefixed_key = Vec::with_capacity(key.as_ref().len() + prefix.0.len() + 1);
240 prefixed_key.extend_from_slice(prefix.0);
241 if let Some(last) = prefix.1 {
242 prefixed_key.push(last);
243 }
244 prefixed_key.extend_from_slice(key.as_ref());
245 prefixed_key
246}
247
248#[derive(Clone, Debug)]
253#[deprecated(since="0.22.0")]
254pub struct LegacyPrefixedKey<H: KeyHasher>(PhantomData<H>);
255
256#[allow(deprecated)]
257impl<H: KeyHasher> KeyFunction<H> for LegacyPrefixedKey<H> {
258 type Key = Vec<u8>;
259
260 fn key(hash: &H::Out, prefix: Prefix) -> Vec<u8> {
261 legacy_prefixed_key::<H>(hash, prefix)
262 }
263}
264
265#[deprecated(since="0.22.0")]
268pub fn legacy_prefixed_key<H: KeyHasher>(key: &H::Out, prefix: Prefix) -> Vec<u8> {
269 let mut prefixed_key = Vec::with_capacity(key.as_ref().len() + prefix.0.len() + 1);
270 if let Some(last) = prefix.1 {
271 let mut prev = 0x01u8;
272 for i in prefix.0.iter() {
273 prefixed_key.push((prev << 4) + (*i >> 4));
274 prev = *i;
275 }
276 prefixed_key.push((prev << 4) + (last >> 4));
277 } else {
278 prefixed_key.push(0);
279 prefixed_key.extend_from_slice(prefix.0);
280 }
281 prefixed_key.extend_from_slice(key.as_ref());
282 prefixed_key
283}
284
285impl<'a, H, KF, T, M> Default for MemoryDB<H, KF, T, M>
286where
287 H: KeyHasher,
288 T: From<&'a [u8]>,
289 KF: KeyFunction<H>,
290 M: MemTracker<T> + Default,
291{
292 fn default() -> Self {
293 Self::from_null_node(&[0u8][..], [0u8][..].into())
294 }
295}
296
297impl<H, KF, T, M> MemoryDB<H, KF, T, M>
299where
300 H: KeyHasher,
301 T: Default,
302 KF: KeyFunction<H>,
303 M: MemTracker<T>,
304{
305 pub fn remove_and_purge(&mut self, key: &<H as KeyHasher>::Out, prefix: Prefix) -> Option<T> {
308 if key == &self.hashed_null_node {
309 return None;
310 }
311 let key = KF::key(key, prefix);
312 match self.data.entry(key) {
313 Entry::Occupied(mut entry) =>
314 if entry.get().1 == 1 {
315 let (value, _) = entry.remove();
316 self.malloc_tracker.on_remove(&value);
317 Some(value)
318 } else {
319 entry.get_mut().1 -= 1;
320 None
321 },
322 Entry::Vacant(entry) => {
323 let value = T::default();
324 self.malloc_tracker.on_insert(&value);
325 entry.insert((value, -1));
326 None
327 }
328 }
329 }
330
331 #[inline]
335 pub fn shrink_to_fit(&mut self) {
336 self.data.shrink_to_fit();
337 }
338}
339
340impl<'a, H, KF, T, M> MemoryDB<H, KF, T, M>
341where
342 H: KeyHasher,
343 T: From<&'a [u8]>,
344 KF: KeyFunction<H>,
345 M: MemTracker<T> + Default,
346{
347 pub fn from_null_node(null_key: &'a [u8], null_node_data: T) -> Self {
349 MemoryDB {
350 data: HashMap::default(),
351 hashed_null_node: H::hash(null_key),
352 null_node_data,
353 malloc_tracker: M::default(),
354 _kf: Default::default(),
355 }
356 }
357
358 pub fn new(data: &'a [u8]) -> Self {
360 Self::from_null_node(data, data.into())
361 }
362
363 pub fn default_with_root() -> (Self, H::Out) {
365 let db = Self::default();
366 let root = db.hashed_null_node;
367
368 (db, root)
369 }
370
371 pub fn clear(&mut self) {
393 self.malloc_tracker.on_clear();
394 self.data.clear();
395 }
396
397 pub fn purge(&mut self) {
399 let malloc_tracker = &mut self.malloc_tracker;
400 self.data.retain(|_, (v, rc)| {
401 let keep = *rc != 0;
402 if !keep {
403 malloc_tracker.on_remove(v);
404 }
405 keep
406 });
407 }
408
409 pub fn drain(&mut self) -> HashMap<KF::Key, (T, i32)> {
411 self.malloc_tracker.on_clear();
412 mem::take(&mut self.data)
413 }
414
415 pub fn raw(&self, key: &<H as KeyHasher>::Out, prefix: Prefix) -> Option<(&T, i32)> {
421 if key == &self.hashed_null_node {
422 return Some((&self.null_node_data, 1));
423 }
424 self.data.get(&KF::key(key, prefix)).map(|(value, count)| (value, *count))
425 }
426
427 pub fn consolidate(&mut self, mut other: Self) {
429 for (key, (value, rc)) in other.drain() {
430 match self.data.entry(key) {
431 Entry::Occupied(mut entry) => {
432 if entry.get().1 < 0 {
433 self.malloc_tracker.on_insert(&value);
434 self.malloc_tracker.on_remove(&entry.get().0);
435 entry.get_mut().0 = value;
436 }
437
438 entry.get_mut().1 += rc;
439 }
440 Entry::Vacant(entry) => {
441 self.malloc_tracker.on_insert(&value);
442 entry.insert((value, rc));
443 }
444 }
445 }
446 }
447
448 pub fn keys(&self) -> HashMap<KF::Key, i32> {
450 self.data.iter()
451 .filter_map(|(k, v)| if v.1 != 0 {
452 Some((k.clone(), v.1))
453 } else {
454 None
455 })
456 .collect()
457 }
458}
459
460impl<H, KF, T, M> MallocSizeOf for MemoryDB<H, KF, T, M>
461where
462 H: KeyHasher,
463 H::Out: MallocSizeOf,
464 T: MallocSizeOf,
465 KF: KeyFunction<H>,
466 KF::Key: MallocSizeOf,
467 M: MemTracker<T>,
468{
469 fn size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
470 self.data.shallow_size_of(ops)
471 + self.malloc_tracker.get_size()
472 + self.null_node_data.size_of(ops)
473 + self.hashed_null_node.size_of(ops)
474 }
475}
476
477impl<H, KF, T, M> PlainDB<H::Out, T> for MemoryDB<H, KF, T, M>
478where
479 H: KeyHasher,
480 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
481 KF: Send + Sync + KeyFunction<H>,
482 KF::Key: Borrow<[u8]> + for <'a> From<&'a [u8]>,
483 M: MemTracker<T> + Send + Sync,
484{
485 fn get(&self, key: &H::Out) -> Option<T> {
486 match self.data.get(key.as_ref()) {
487 Some(&(ref d, rc)) if rc > 0 => Some(d.clone()),
488 _ => None
489 }
490 }
491
492 fn contains(&self, key: &H::Out) -> bool {
493 match self.data.get(key.as_ref()) {
494 Some(&(_, x)) if x > 0 => true,
495 _ => false
496 }
497 }
498
499 fn emplace(&mut self, key: H::Out, value: T) {
500 match self.data.entry(key.as_ref().into()) {
501 Entry::Occupied(mut entry) => {
502 let &mut (ref mut old_value, ref mut rc) = entry.get_mut();
503 if *rc <= 0 {
504 self.malloc_tracker.on_insert(&value);
505 self.malloc_tracker.on_remove(old_value);
506 *old_value = value;
507 }
508 *rc += 1;
509 },
510 Entry::Vacant(entry) => {
511 self.malloc_tracker.on_insert(&value);
512 entry.insert((value, 1));
513 },
514 }
515 }
516
517 fn remove(&mut self, key: &H::Out) {
518 match self.data.entry(key.as_ref().into()) {
519 Entry::Occupied(mut entry) => {
520 let &mut (_, ref mut rc) = entry.get_mut();
521 *rc -= 1;
522 },
523 Entry::Vacant(entry) => {
524 let value = T::default();
525 self.malloc_tracker.on_insert(&value);
526 entry.insert((value, -1));
527 },
528 }
529 }
530}
531
532impl<H, KF, T, M> PlainDBRef<H::Out, T> for MemoryDB<H, KF, T, M>
533where
534 H: KeyHasher,
535 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
536 KF: Send + Sync + KeyFunction<H>,
537 KF::Key: Borrow<[u8]> + for <'a> From<&'a [u8]>,
538 M: MemTracker<T> + Send + Sync,
539{
540 fn get(&self, key: &H::Out) -> Option<T> { PlainDB::get(self, key) }
541 fn contains(&self, key: &H::Out) -> bool { PlainDB::contains(self, key) }
542}
543
544impl<H, KF, T, M> HashDB<H, T> for MemoryDB<H, KF, T, M>
545where
546 H: KeyHasher,
547 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
548 KF: Send + Sync + KeyFunction<H>,
549 M: MemTracker<T> + Send + Sync,
550{
551 fn get(&self, key: &H::Out, prefix: Prefix) -> Option<T> {
552 if key == &self.hashed_null_node {
553 return Some(self.null_node_data.clone());
554 }
555
556 let key = KF::key(key, prefix);
557 match self.data.get(&key) {
558 Some(&(ref d, rc)) if rc > 0 => Some(d.clone()),
559 _ => None
560 }
561 }
562
563 fn contains(&self, key: &H::Out, prefix: Prefix) -> bool {
564 if key == &self.hashed_null_node {
565 return true;
566 }
567
568 let key = KF::key(key, prefix);
569 match self.data.get(&key) {
570 Some(&(_, x)) if x > 0 => true,
571 _ => false
572 }
573 }
574
575 fn emplace(&mut self, key: H::Out, prefix: Prefix, value: T) {
576 if value == self.null_node_data {
577 return;
578 }
579
580 let key = KF::key(&key, prefix);
581 match self.data.entry(key) {
582 Entry::Occupied(mut entry) => {
583 let &mut (ref mut old_value, ref mut rc) = entry.get_mut();
584 if *rc <= 0 {
585 self.malloc_tracker.on_insert(&value);
586 self.malloc_tracker.on_remove(old_value);
587 *old_value = value;
588 }
589 *rc += 1;
590 },
591 Entry::Vacant(entry) => {
592 self.malloc_tracker.on_insert(&value);
593 entry.insert((value, 1));
594 },
595 }
596 }
597
598 fn insert(&mut self, prefix: Prefix, value: &[u8]) -> H::Out {
599 if T::from(value) == self.null_node_data {
600 return self.hashed_null_node;
601 }
602
603 let key = H::hash(value);
604 HashDB::emplace(self, key, prefix, value.into());
605 key
606 }
607
608 fn remove(&mut self, key: &H::Out, prefix: Prefix) {
609 if key == &self.hashed_null_node {
610 return;
611 }
612
613 let key = KF::key(key, prefix);
614 match self.data.entry(key) {
615 Entry::Occupied(mut entry) => {
616 let &mut (_, ref mut rc) = entry.get_mut();
617 *rc -= 1;
618 },
619 Entry::Vacant(entry) => {
620 let value = T::default();
621 self.malloc_tracker.on_insert(&value);
622 entry.insert((value, -1));
623 },
624 }
625 }
626}
627
628impl<H, KF, T, M> HashDBRef<H, T> for MemoryDB<H, KF, T, M>
629where
630 H: KeyHasher,
631 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
632 KF: KeyFunction<H> + Send + Sync,
633 M: MemTracker<T> + Send + Sync,
634{
635 fn get(&self, key: &H::Out, prefix: Prefix) -> Option<T> { HashDB::get(self, key, prefix) }
636 fn contains(&self, key: &H::Out, prefix: Prefix) -> bool { HashDB::contains(self, key, prefix) }
637}
638
639impl<H, KF, T, M> AsPlainDB<H::Out, T> for MemoryDB<H, KF, T, M>
640where
641 H: KeyHasher,
642 T: Default + PartialEq<T> + for<'a> From<&'a[u8]> + Clone + Send + Sync,
643 KF: KeyFunction<H> + Send + Sync,
644 KF::Key: Borrow<[u8]> + for <'a> From<&'a [u8]>,
645 M: MemTracker<T> + Send + Sync,
646{
647 fn as_plain_db(&self) -> &dyn PlainDB<H::Out, T> { self }
648 fn as_plain_db_mut(&mut self) -> &mut dyn PlainDB<H::Out, T> { self }
649}
650
651impl<H, KF, T, M> AsHashDB<H, T> for MemoryDB<H, KF, T, M>
652where
653 H: KeyHasher,
654 T: Default + PartialEq<T> + for<'a> From<&'a[u8]> + Clone + Send + Sync,
655 KF: KeyFunction<H> + Send + Sync,
656 M: MemTracker<T> + Send + Sync,
657{
658 fn as_hash_db(&self) -> &dyn HashDB<H, T> { self }
659 fn as_hash_db_mut(&mut self) -> &mut dyn HashDB<H, T> { self }
660}
661
662#[cfg(test)]
663mod tests {
664 use super::{MemoryDB, HashDB, KeyHasher, HashKey};
665 use tetsy_util_mem::malloc_size;
666 use tetsy_hash_db::EMPTY_PREFIX;
667 use tetsy_keccak_hasher::KeccakHasher;
668
669 #[test]
670 fn memorydb_remove_and_purge() {
671 let hello_bytes = b"Hello world!";
672 let hello_key = KeccakHasher::hash(hello_bytes);
673
674 let mut m = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
675 m.remove(&hello_key, EMPTY_PREFIX);
676 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1);
677 m.purge();
678 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1);
679 m.insert(EMPTY_PREFIX, hello_bytes);
680 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, 0);
681 m.purge();
682 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX), None);
683
684 let mut m = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
685 assert!(m.remove_and_purge(&hello_key, EMPTY_PREFIX).is_none());
686 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1);
687 m.insert(EMPTY_PREFIX, hello_bytes);
688 m.insert(EMPTY_PREFIX, hello_bytes);
689 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, 1);
690 assert_eq!(&*m.remove_and_purge(&hello_key, EMPTY_PREFIX).unwrap(), hello_bytes);
691 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX), None);
692 assert!(m.remove_and_purge(&hello_key, EMPTY_PREFIX).is_none());
693 }
694
695 #[test]
696 fn consolidate() {
697 let mut main = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
698 let mut other = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
699 let remove_key = other.insert(EMPTY_PREFIX, b"doggo");
700 main.remove(&remove_key, EMPTY_PREFIX);
701
702 let insert_key = other.insert(EMPTY_PREFIX, b"arf");
703 main.emplace(insert_key, EMPTY_PREFIX, "arf".as_bytes().to_vec());
704
705 let negative_remove_key = other.insert(EMPTY_PREFIX, b"negative");
706 other.remove(&negative_remove_key, EMPTY_PREFIX); other.remove(&negative_remove_key, EMPTY_PREFIX); main.remove(&negative_remove_key, EMPTY_PREFIX); main.consolidate(other);
711
712 assert_eq!(main.raw(&remove_key, EMPTY_PREFIX).unwrap(), (&"doggo".as_bytes().to_vec(), 0));
713 assert_eq!(main.raw(&insert_key, EMPTY_PREFIX).unwrap(), (&"arf".as_bytes().to_vec(), 2));
714 assert_eq!(
715 main.raw(&negative_remove_key, EMPTY_PREFIX).unwrap(),
716 (&"negative".as_bytes().to_vec(), -2),
717 );
718 }
719
720 #[test]
721 fn default_works() {
722 let mut db = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
723 let hashed_null_node = KeccakHasher::hash(&[0u8][..]);
724 assert_eq!(db.insert(EMPTY_PREFIX, &[0u8][..]), hashed_null_node);
725
726 let (db2, root) = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default_with_root();
727 assert!(db2.contains(&root, EMPTY_PREFIX));
728 assert!(db.contains(&root, EMPTY_PREFIX));
729 }
730
731 #[test]
732 fn malloc_size_of() {
733 let mut db = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
734 for i in 0u32..1024 {
735 let bytes = i.to_be_bytes();
736 let prefix = (bytes.as_ref(), None);
737 db.insert(prefix, &bytes);
738 }
739 assert_eq!(
740 malloc_size(&db),
741 malloc_size(&db.data) + malloc_size(&db.null_node_data) + malloc_size(&db.hashed_null_node)
742 );
743 }
744}