1#![cfg_attr(not(feature = "std"), no_std)]
18
19#[cfg(not(feature = "std"))]
20extern crate alloc;
21
22use core::hash::BuildHasher;
23use hash_db::{
24 AsHashDB, AsPlainDB, HashDB, HashDBRef, Hasher as KeyHasher, MaybeDebug, PlainDB, PlainDBRef,
25 Prefix,
26};
27#[cfg(feature = "std")]
28use std::{
29 borrow::Borrow, cmp::Eq, collections::hash_map::Entry, collections::HashMap as Map, hash,
30 hash::RandomState, marker::PhantomData, mem,
31};
32
33#[cfg(not(feature = "std"))]
34use hashbrown::{hash_map::Entry, HashMap as Map};
35
36#[cfg(not(feature = "std"))]
37use foldhash::quality::RandomState;
38
39#[cfg(not(feature = "std"))]
40use core::{borrow::Borrow, cmp::Eq, hash, marker::PhantomData, mem};
41
42#[cfg(not(feature = "std"))]
43use alloc::vec::Vec;
44
45pub struct MemoryDB<H, KF, T, S = RandomState>
88where
89 H: KeyHasher,
90 KF: KeyFunction<H>,
91{
92 data: Map<KF::Key, (T, i32), S>,
93 hashed_null_node: H::Out,
94 null_node_data: T,
95 _kf: PhantomData<KF>,
96}
97
98impl<H, KF, T, S> Clone for MemoryDB<H, KF, T, S>
99where
100 H: KeyHasher,
101 KF: KeyFunction<H>,
102 T: Clone,
103 S: Clone,
104{
105 fn clone(&self) -> Self {
106 Self {
107 data: self.data.clone(),
108 hashed_null_node: self.hashed_null_node,
109 null_node_data: self.null_node_data.clone(),
110 _kf: Default::default(),
111 }
112 }
113}
114
115impl<H, KF, T, S> PartialEq<MemoryDB<H, KF, T, S>> for MemoryDB<H, KF, T, S>
116where
117 H: KeyHasher,
118 KF: KeyFunction<H>,
119 T: Eq + MaybeDebug,
120 S: BuildHasher,
121{
122 fn eq(&self, other: &MemoryDB<H, KF, T, S>) -> bool {
123 for a in self.data.iter() {
124 match other.data.get(a.0) {
125 Some(v) if v != a.1 => return false,
126 None => return false,
127 _ => (),
128 }
129 }
130 true
131 }
132}
133
134impl<H, KF, T, S> Eq for MemoryDB<H, KF, T, S>
135where
136 H: KeyHasher,
137 KF: KeyFunction<H>,
138 T: Eq + MaybeDebug,
139 S: BuildHasher,
140{
141}
142
143pub trait KeyFunction<H: KeyHasher> {
144 type Key: Send + Sync + Clone + hash::Hash + Eq + MaybeDebug + core::cmp::Ord;
145
146 fn key(hash: &H::Out, prefix: Prefix) -> Self::Key;
147}
148
149pub struct HashKey<H>(PhantomData<H>);
151
152impl<H> Clone for HashKey<H> {
153 fn clone(&self) -> Self {
154 Self(Default::default())
155 }
156}
157
158impl<H> core::fmt::Debug for HashKey<H> {
159 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
160 core::write!(f, "HashKey")
161 }
162}
163
164impl<H: KeyHasher> KeyFunction<H> for HashKey<H> {
165 type Key = H::Out;
166
167 fn key(hash: &H::Out, prefix: Prefix) -> H::Out {
168 hash_key::<H>(hash, prefix)
169 }
170}
171
172pub fn hash_key<H: KeyHasher>(key: &H::Out, _prefix: Prefix) -> H::Out {
174 *key
175}
176
177pub struct PrefixedKey<H>(PhantomData<H>);
179
180impl<H> Clone for PrefixedKey<H> {
181 fn clone(&self) -> Self {
182 Self(Default::default())
183 }
184}
185
186impl<H> core::fmt::Debug for PrefixedKey<H> {
187 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
188 core::write!(f, "PrefixedKey")
189 }
190}
191
192impl<H: KeyHasher> KeyFunction<H> for PrefixedKey<H> {
193 type Key = Vec<u8>;
194
195 fn key(hash: &H::Out, prefix: Prefix) -> Vec<u8> {
196 prefixed_key::<H>(hash, prefix)
197 }
198}
199
200pub fn prefixed_key<H: KeyHasher>(key: &H::Out, prefix: Prefix) -> Vec<u8> {
202 let mut prefixed_key = Vec::with_capacity(key.as_ref().len() + prefix.0.len() + 1);
203 prefixed_key.extend_from_slice(prefix.0);
204 if let Some(last) = prefix.1 {
205 prefixed_key.push(last);
206 }
207 prefixed_key.extend_from_slice(key.as_ref());
208 prefixed_key
209}
210
211#[derive(Clone, Debug)]
216#[deprecated(since = "0.22.0")]
217pub struct LegacyPrefixedKey<H: KeyHasher>(PhantomData<H>);
218
219#[allow(deprecated)]
220impl<H: KeyHasher> KeyFunction<H> for LegacyPrefixedKey<H> {
221 type Key = Vec<u8>;
222
223 fn key(hash: &H::Out, prefix: Prefix) -> Vec<u8> {
224 legacy_prefixed_key::<H>(hash, prefix)
225 }
226}
227
228#[deprecated(since = "0.22.0")]
231pub fn legacy_prefixed_key<H: KeyHasher>(key: &H::Out, prefix: Prefix) -> Vec<u8> {
232 let mut prefixed_key = Vec::with_capacity(key.as_ref().len() + prefix.0.len() + 1);
233 if let Some(last) = prefix.1 {
234 let mut prev = 0x01u8;
235 for i in prefix.0.iter() {
236 prefixed_key.push((prev << 4) + (*i >> 4));
237 prev = *i;
238 }
239 prefixed_key.push((prev << 4) + (last >> 4));
240 } else {
241 prefixed_key.push(0);
242 prefixed_key.extend_from_slice(prefix.0);
243 }
244 prefixed_key.extend_from_slice(key.as_ref());
245 prefixed_key
246}
247
248impl<H, KF, T> Default for MemoryDB<H, KF, T, RandomState>
249where
250 H: KeyHasher,
251 T: for<'a> From<&'a [u8]>,
252 KF: KeyFunction<H>,
253{
254 fn default() -> Self {
255 Self::from_null_node(&[0u8][..], [0u8][..].into())
256 }
257}
258
259impl<H, KF, T, S> MemoryDB<H, KF, T, S>
261where
262 H: KeyHasher,
263 T: Default,
264 KF: KeyFunction<H>,
265 S: BuildHasher + Default,
266{
267 pub fn remove_and_purge(&mut self, key: &<H as KeyHasher>::Out, prefix: Prefix) -> Option<T> {
270 if key == &self.hashed_null_node {
271 return None
272 }
273 let key = KF::key(key, prefix);
274 match self.data.entry(key) {
275 Entry::Occupied(mut entry) =>
276 if entry.get().1 == 1 {
277 let (value, _) = entry.remove();
278 Some(value)
279 } else {
280 entry.get_mut().1 -= 1;
281 None
282 },
283 Entry::Vacant(entry) => {
284 let value = T::default();
285 entry.insert((value, -1));
286 None
287 },
288 }
289 }
290
291 #[inline]
295 pub fn shrink_to_fit(&mut self) {
296 #[cfg(feature = "std")]
297 self.data.shrink_to_fit();
298 }
299}
300
301impl<H, KF, T, S> MemoryDB<H, KF, T, S>
302where
303 H: KeyHasher,
304 T: for<'a> From<&'a [u8]>,
305 KF: KeyFunction<H>,
306 S: BuildHasher + Default,
307{
308 pub fn from_null_node(null_key: &[u8], null_node_data: T) -> Self {
310 Self::from_null_node_with_hasher(null_key, null_node_data, S::default())
311 }
312
313 pub fn from_null_node_with_hasher(null_key: &[u8], null_node_data: T, hasher: S) -> Self {
315 MemoryDB {
316 data: Map::with_hasher(hasher),
317 hashed_null_node: H::hash(null_key),
318 null_node_data,
319 _kf: Default::default(),
320 }
321 }
322
323 pub fn new(data: &[u8]) -> Self {
325 Self::from_null_node(data, data.into())
326 }
327
328 pub fn default_with_root() -> (Self, H::Out) {
330 let db = Self::new(&[0u8][..]);
331 let root = db.hashed_null_node;
332
333 (db, root)
334 }
335
336 pub fn with_hasher(hasher: S) -> Self {
338 Self::from_null_node_with_hasher(&[0u8][..], [0u8][..].into(), hasher)
339 }
340
341 pub fn clear(&mut self) {
363 self.data.clear();
364 }
365
366 pub fn purge(&mut self) {
368 self.data.retain(|_, (_, rc)| {
369 let keep = *rc != 0;
370 keep
371 });
372 }
373
374 pub fn drain(&mut self) -> Map<KF::Key, (T, i32), S> {
376 mem::take(&mut self.data)
377 }
378
379 pub fn raw(&self, key: &<H as KeyHasher>::Out, prefix: Prefix) -> Option<(&T, i32)> {
385 if key == &self.hashed_null_node {
386 return Some((&self.null_node_data, 1))
387 }
388 self.data.get(&KF::key(key, prefix)).map(|(value, count)| (value, *count))
389 }
390
391 pub fn consolidate(&mut self, mut other: Self) {
393 for (key, (value, rc)) in other.drain() {
394 match self.data.entry(key) {
395 Entry::Occupied(mut entry) => {
396 if entry.get().1 < 0 {
397 entry.get_mut().0 = value;
398 }
399
400 entry.get_mut().1 += rc;
401 },
402 Entry::Vacant(entry) => {
403 entry.insert((value, rc));
404 },
405 }
406 }
407 }
408
409 pub fn keys(&self) -> Map<KF::Key, i32> {
411 self.data
412 .iter()
413 .filter_map(|(k, v)| if v.1 != 0 { Some((k.clone(), v.1)) } else { None })
414 .collect()
415 }
416}
417
418impl<H, KF, T, S> PlainDB<H::Out, T> for MemoryDB<H, KF, T, S>
419where
420 H: KeyHasher,
421 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
422 KF: Send + Sync + KeyFunction<H>,
423 KF::Key: Borrow<[u8]> + for<'a> From<&'a [u8]>,
424 S: BuildHasher + Default + Send + Sync,
425{
426 fn get(&self, key: &H::Out) -> Option<T> {
427 match self.data.get(key.as_ref()) {
428 Some(&(ref d, rc)) if rc > 0 => Some(d.clone()),
429 _ => None,
430 }
431 }
432
433 fn contains(&self, key: &H::Out) -> bool {
434 match self.data.get(key.as_ref()) {
435 Some(&(_, x)) if x > 0 => true,
436 _ => false,
437 }
438 }
439
440 fn emplace(&mut self, key: H::Out, value: T) {
441 match self.data.entry(key.as_ref().into()) {
442 Entry::Occupied(mut entry) => {
443 let &mut (ref mut old_value, ref mut rc) = entry.get_mut();
444 if *rc <= 0 {
445 *old_value = value;
446 }
447 *rc += 1;
448 },
449 Entry::Vacant(entry) => {
450 entry.insert((value, 1));
451 },
452 }
453 }
454
455 fn remove(&mut self, key: &H::Out) {
456 match self.data.entry(key.as_ref().into()) {
457 Entry::Occupied(mut entry) => {
458 let &mut (_, ref mut rc) = entry.get_mut();
459 *rc -= 1;
460 },
461 Entry::Vacant(entry) => {
462 let value = T::default();
463 entry.insert((value, -1));
464 },
465 }
466 }
467}
468
469impl<H, KF, T, S> PlainDBRef<H::Out, T> for MemoryDB<H, KF, T, S>
470where
471 H: KeyHasher,
472 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
473 KF: Send + Sync + KeyFunction<H>,
474 KF::Key: Borrow<[u8]> + for<'a> From<&'a [u8]>,
475 S: BuildHasher + Default + Send + Sync,
476{
477 fn get(&self, key: &H::Out) -> Option<T> {
478 PlainDB::get(self, key)
479 }
480 fn contains(&self, key: &H::Out) -> bool {
481 PlainDB::contains(self, key)
482 }
483}
484
485impl<H, KF, T, S> HashDB<H, T> for MemoryDB<H, KF, T, S>
486where
487 H: KeyHasher,
488 T: Default + PartialEq<T> + AsRef<[u8]> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
489 KF: KeyFunction<H> + Send + Sync,
490 S: BuildHasher + Default + Send + Sync,
491{
492 fn get(&self, key: &H::Out, prefix: Prefix) -> Option<T> {
493 if key == &self.hashed_null_node {
494 return Some(self.null_node_data.clone())
495 }
496
497 let key = KF::key(key, prefix);
498 match self.data.get(&key) {
499 Some(&(ref d, rc)) if rc > 0 => Some(d.clone()),
500 _ => None,
501 }
502 }
503
504 fn contains(&self, key: &H::Out, prefix: Prefix) -> bool {
505 if key == &self.hashed_null_node {
506 return true
507 }
508
509 let key = KF::key(key, prefix);
510 match self.data.get(&key) {
511 Some(&(_, x)) if x > 0 => true,
512 _ => false,
513 }
514 }
515
516 fn emplace(&mut self, key: H::Out, prefix: Prefix, value: T) {
517 if value == self.null_node_data {
518 return
519 }
520
521 let key = KF::key(&key, prefix);
522 match self.data.entry(key) {
523 Entry::Occupied(mut entry) => {
524 let &mut (ref mut old_value, ref mut rc) = entry.get_mut();
525 if *rc <= 0 {
526 *old_value = value;
527 }
528 *rc += 1;
529 },
530 Entry::Vacant(entry) => {
531 entry.insert((value, 1));
532 },
533 }
534 }
535
536 fn insert(&mut self, prefix: Prefix, value: &[u8]) -> H::Out {
537 if T::from(value) == self.null_node_data {
538 return self.hashed_null_node
539 }
540
541 let key = H::hash(value);
542 HashDB::emplace(self, key, prefix, value.into());
543 key
544 }
545
546 fn remove(&mut self, key: &H::Out, prefix: Prefix) {
547 if key == &self.hashed_null_node {
548 return
549 }
550
551 let key = KF::key(key, prefix);
552 match self.data.entry(key) {
553 Entry::Occupied(mut entry) => {
554 let &mut (_, ref mut rc) = entry.get_mut();
555 *rc -= 1;
556 },
557 Entry::Vacant(entry) => {
558 let value = T::default();
559 entry.insert((value, -1));
560 },
561 }
562 }
563}
564
565impl<H, KF, T, S> HashDBRef<H, T> for MemoryDB<H, KF, T, S>
566where
567 H: KeyHasher,
568 T: Default + PartialEq<T> + AsRef<[u8]> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
569 KF: KeyFunction<H> + Send + Sync,
570 S: BuildHasher + Default + Send + Sync,
571{
572 fn get(&self, key: &H::Out, prefix: Prefix) -> Option<T> {
573 HashDB::get(self, key, prefix)
574 }
575 fn contains(&self, key: &H::Out, prefix: Prefix) -> bool {
576 HashDB::contains(self, key, prefix)
577 }
578}
579
580impl<H, KF, T, S> AsPlainDB<H::Out, T> for MemoryDB<H, KF, T, S>
581where
582 H: KeyHasher,
583 T: Default + PartialEq<T> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
584 KF: KeyFunction<H> + Send + Sync,
585 KF::Key: Borrow<[u8]> + for<'a> From<&'a [u8]>,
586 S: BuildHasher + Default + Send + Sync,
587{
588 fn as_plain_db(&self) -> &dyn PlainDB<H::Out, T> {
589 self
590 }
591 fn as_plain_db_mut(&mut self) -> &mut dyn PlainDB<H::Out, T> {
592 self
593 }
594}
595
596impl<H, KF, T, S> AsHashDB<H, T> for MemoryDB<H, KF, T, S>
597where
598 H: KeyHasher,
599 T: Default + PartialEq<T> + AsRef<[u8]> + for<'a> From<&'a [u8]> + Clone + Send + Sync,
600 KF: KeyFunction<H> + Send + Sync,
601 S: BuildHasher + Default + Send + Sync,
602{
603 fn as_hash_db(&self) -> &dyn HashDB<H, T> {
604 self
605 }
606 fn as_hash_db_mut(&mut self) -> &mut dyn HashDB<H, T> {
607 self
608 }
609}
610
611#[cfg(test)]
612mod tests {
613 use super::{HashDB, HashKey, KeyHasher, MemoryDB};
614 use hash_db::EMPTY_PREFIX;
615 use keccak_hasher::KeccakHasher;
616
617 #[test]
618 fn memorydb_remove_and_purge() {
619 let hello_bytes = b"Hello world!";
620 let hello_key = KeccakHasher::hash(hello_bytes);
621
622 let mut m = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
623 m.remove(&hello_key, EMPTY_PREFIX);
624 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1);
625 m.purge();
626 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1);
627 m.insert(EMPTY_PREFIX, hello_bytes);
628 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, 0);
629 m.purge();
630 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX), None);
631
632 let mut m = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
633 assert!(m.remove_and_purge(&hello_key, EMPTY_PREFIX).is_none());
634 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, -1);
635 m.insert(EMPTY_PREFIX, hello_bytes);
636 m.insert(EMPTY_PREFIX, hello_bytes);
637 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX).unwrap().1, 1);
638 assert_eq!(&*m.remove_and_purge(&hello_key, EMPTY_PREFIX).unwrap(), hello_bytes);
639 assert_eq!(m.raw(&hello_key, EMPTY_PREFIX), None);
640 assert!(m.remove_and_purge(&hello_key, EMPTY_PREFIX).is_none());
641 }
642
643 #[test]
644 fn consolidate() {
645 let mut main = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
646 let mut other = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
647 let remove_key = other.insert(EMPTY_PREFIX, b"doggo");
648 main.remove(&remove_key, EMPTY_PREFIX);
649
650 let insert_key = other.insert(EMPTY_PREFIX, b"arf");
651 main.emplace(insert_key, EMPTY_PREFIX, "arf".as_bytes().to_vec());
652
653 let negative_remove_key = other.insert(EMPTY_PREFIX, b"negative");
654 other.remove(&negative_remove_key, EMPTY_PREFIX); other.remove(&negative_remove_key, EMPTY_PREFIX); main.remove(&negative_remove_key, EMPTY_PREFIX); main.consolidate(other);
659
660 assert_eq!(main.raw(&remove_key, EMPTY_PREFIX).unwrap(), (&"doggo".as_bytes().to_vec(), 0));
661 assert_eq!(main.raw(&insert_key, EMPTY_PREFIX).unwrap(), (&"arf".as_bytes().to_vec(), 2));
662 assert_eq!(
663 main.raw(&negative_remove_key, EMPTY_PREFIX).unwrap(),
664 (&"negative".as_bytes().to_vec(), -2),
665 );
666 }
667
668 #[test]
669 fn default_works() {
670 let mut db = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default();
671 let hashed_null_node = KeccakHasher::hash(&[0u8][..]);
672 assert_eq!(db.insert(EMPTY_PREFIX, &[0u8][..]), hashed_null_node);
673
674 let (db2, root) = MemoryDB::<KeccakHasher, HashKey<_>, Vec<u8>>::default_with_root();
675 assert!(db2.contains(&root, EMPTY_PREFIX));
676 assert!(db.contains(&root, EMPTY_PREFIX));
677 }
678}