jmt_blake3/
overlay.rs

1use std::collections::HashMap;
2
3use anyhow::Result;
4use tracing::instrument;
5
6use crate::{
7    storage::{TreeReader, TreeWriter},
8    JellyfishMerkleTree, KeyHash, MissingRootError, OwnedValue, RootHash, Version,
9};
10
11// There's some similarity between the OverlayTree and the TreeCache, but it's
12// not exactly clear how to share code between them (the TreeCache operates at
13// the level of the tree internals, not at the level of its external interface),
14// so it's easier and better for now to just make a new wrapper.
15
16/// A wrapper around a [`JellyfishMerkleTree`] that buffers pending writes to the
17/// tree, and overlays the effect of those writes on the tree state for reading.
18pub struct WriteOverlay<R> {
19    reader: R,
20    overlay: HashMap<KeyHash, OwnedValue>,
21    version: Version,
22}
23
24impl<R> WriteOverlay<R> {
25    /// Use this [`Version`] with [`Self::new`] to specify that the writes
26    /// should be committed with version `0`.
27    pub const PRE_GENESIS_VERSION: Version = u64::MAX;
28}
29
30impl<R> WriteOverlay<R>
31where
32    R: TreeReader + Sync,
33{
34    /// Constructs a new [`WriteOverlay`] with the given `reader` and `version`.
35    ///
36    /// All reads performed with `get` will use `version` when querying the
37    /// underlying backing store.  The buffered writes created with `put` will
38    /// be written as `version + 1`, so `version` should probably be the latest
39    /// version if `commit` will be called.
40    ///
41    /// To initialize an empty tree, use [`Self::PRE_GENESIS_VERSION`] here.
42    pub fn new(reader: R, version: Version) -> Self {
43        Self {
44            reader,
45            version,
46            overlay: Default::default(),
47        }
48    }
49
50    fn tree(&self) -> JellyfishMerkleTree<'_, R> {
51        JellyfishMerkleTree::new(&self.reader)
52    }
53
54    /// Gets a value by key.
55    ///
56    /// This method reflects the results of any pending writes made by `put`.
57    #[instrument(name = "WriteOverlay::get", skip(self, key))]
58    pub async fn get(&self, key: KeyHash) -> Result<Option<OwnedValue>> {
59        if let Some(value) = self.overlay.get(&key) {
60            tracing::trace!(?key, value = ?hex::encode(&value), "read from cache");
61            Ok(Some(value.clone()))
62        } else {
63            match self.tree().get(key, self.version).await {
64                Ok(Some(value)) => {
65                    tracing::trace!(version = ?self.version, ?key, value = ?hex::encode(&value), "read from tree");
66                    Ok(Some(value))
67                }
68                Ok(None) => {
69                    tracing::trace!(version = ?self.version, ?key, "key not found in tree");
70                    Ok(None)
71                }
72                // This allows for using the Overlay on an empty database without errors
73                Err(e) if e.downcast_ref::<MissingRootError>().is_some() => {
74                    tracing::trace!(version = ?self.version, "no data available at this version");
75                    Ok(None)
76                }
77                Err(e) => Err(e),
78            }
79        }
80    }
81
82    /// Gets a value by key alongside an ICS23 existence proof of that value.
83    ///
84    /// This method does *not* reflect results of any pending writes to the WriteOverlay. An error
85    /// will be returned if the key exists in the WriteOverlay, or if the key does not exist in the
86    /// tree.
87    #[instrument(name = "WriteOverlay::get_with_proof", skip(self, key))]
88    pub async fn get_with_proof(
89        &self,
90        key: Vec<u8>,
91    ) -> Result<(OwnedValue, ics23_blake3::ExistenceProof)> {
92        if self.overlay.contains_key(&key.clone().into()) {
93            return Err(anyhow::anyhow!("key is not yet committed to tree"));
94        }
95        let proof = self.tree().get_with_ics23_proof(key, self.version).await?;
96
97        Ok((proof.value.clone(), proof))
98    }
99
100    /// Puts a key/value pair in the overlay.
101    ///
102    /// Assuming it is not overwritten by a subsequent `put`, the value will be
103    /// written to the tree when `commit` is called.
104    #[instrument(name = "WriteOverlay::put", skip(self, key, value))]
105    pub fn put(&mut self, key: KeyHash, value: OwnedValue) {
106        tracing::trace!(?key, value = ?hex::encode(&value));
107        *self.overlay.entry(key).or_default() = value;
108    }
109
110    /// Clears the overlay, committing all pending writes to the provided
111    /// `writer` and returning the new [`RootHash`] and [`Version`].
112    ///
113    /// The overlay will then point at the newly written state and tree version.
114    #[instrument(name = "WriteOverlay::commit", skip(self, writer))]
115    pub async fn commit<W>(&mut self, mut writer: W) -> Result<(RootHash, Version)>
116    where
117        W: TreeWriter + Sync,
118    {
119        // Committing an empty write overlay will panic in the underlying JMT, so we provide a more
120        // meaningful panic message here instead.
121        if self.overlay.is_empty() {
122            panic!("cannot commit empty `WriteOverlay`");
123        }
124
125        let overlay = std::mem::take(&mut self.overlay);
126        // We use wrapping_add here so that we can write `new_version = 0` by
127        // overflowing `PRE_GENESIS_VERSION`.
128        let new_version = self.version.wrapping_add(1);
129        tracing::trace!(old_version = ?self.version, new_version, ?overlay);
130        let (root_hash, batch) = self
131            .tree()
132            .put_value_set(overlay.into_iter().collect(), new_version)
133            .await?;
134
135        writer.write_node_batch(&batch.node_batch).await?;
136        tracing::trace!(?root_hash, "wrote node batch to backing store");
137
138        // Now that we've successfully written the new nodes, update the version.
139        self.version = new_version;
140
141        Ok((root_hash, new_version))
142    }
143}
jmt_blake3/overlay.rs

jmt_blake3/
overlay.rs
