refactor: put node hasher in charge of labeling nodes and determining…

… kinds

benchtop/src/nomt.rs

@@ -1,7 +1,7 @@
 use crate::{backend::Transaction, timer::Timer, workload::Workload};
 use fxhash::FxHashMap;
 use nomt::{
-    Blake3Hasher, KeyPath, KeyReadWrite, Nomt, Options, Overlay, Session, SessionParams,
+    trie::KeyPath, Blake3Hasher, KeyReadWrite, Nomt, Options, Overlay, Session, SessionParams,
     WitnessMode,
 };
 use sha2::Digest;

core/src/proof/multi_proof.rs

@@ -7,7 +7,7 @@ use crate::{
         path_proof::{hash_path, shared_bits},
         KeyOutOfScope, PathProof, PathProofTerminal,
     },
-    trie::{InternalData, KeyPath, LeafData, Node, NodeHasher, NodeHasherExt, TERMINATOR},
+    trie::{InternalData, KeyPath, LeafData, Node, NodeHasher, TERMINATOR},
 };
 
 #[cfg(not(feature = "std"))]
@@ -498,7 +498,7 @@ mod tests {
 
     use crate::{
         proof::{PathProof, PathProofTerminal},
-        trie::{self, InternalData, LeafData, NodeHasher, NodeHasherExt, TERMINATOR},
+        trie::{self, InternalData, LeafData, Node, NodeHasher, NodeKind, TERMINATOR},
         trie_pos::TriePosition,
     };
 
@@ -888,8 +888,36 @@ mod tests {
     pub struct Blake3Hasher;
 
     impl NodeHasher for Blake3Hasher {
-        fn hash_node(data: &trie::NodePreimage) -> [u8; 32] {
-            blake3::hash(data).into()
+        fn hash_leaf(data: &trie::LeafData) -> [u8; 32] {
+            let mut hasher = blake3::Hasher::new();
+            hasher.update(&data.key_path);
+            hasher.update(&data.value_hash);
+            let mut hash: [u8; 32] = hasher.finalize().into();
+
+            // Label with MSB
+            hash[0] |= 0b10000000;
+            hash
+        }
+
+        fn hash_internal(data: &trie::InternalData) -> [u8; 32] {
+            let mut hasher = blake3::Hasher::new();
+            hasher.update(&data.left);
+            hasher.update(&data.right);
+            let mut hash: [u8; 32] = hasher.finalize().into();
+
+            // Label with MSB
+            hash[0] &= 0b01111111;
+            hash
+        }
+
+        fn node_kind(node: &Node) -> NodeKind {
+            if node[0] >> 7 == 1 {
+                NodeKind::Leaf
+            } else if node == &TERMINATOR {
+                NodeKind::Terminator
+            } else {
+                NodeKind::Internal
+            }
         }
     }
 

core/src/proof/path_proof.rs

@@ -1,8 +1,6 @@
 //! Proving and verifying inclusion, non-inclusion, and updates to the trie.
 
-use crate::trie::{
-    self, InternalData, KeyPath, LeafData, Node, NodeHasher, NodeHasherExt, NodeKind, TERMINATOR,
-};
+use crate::trie::{self, InternalData, KeyPath, LeafData, Node, NodeHasher, NodeKind, TERMINATOR};
 use crate::trie_pos::TriePosition;
 
 use bitvec::prelude::*;
@@ -295,7 +293,7 @@ pub fn verify_update<H: NodeHasher>(
                 *sibling
             };
 
-            match (NodeKind::of(&cur_node), NodeKind::of(&sibling)) {
+            match (NodeKind::of::<H>(&cur_node), NodeKind::of::<H>(&sibling)) {
                 (NodeKind::Terminator, NodeKind::Terminator) => {}
                 (NodeKind::Leaf, NodeKind::Terminator) => {}
                 (NodeKind::Terminator, NodeKind::Leaf) => {

core/src/trie.rs

@@ -13,10 +13,11 @@
 //!
 //! All node preimages are 512 bits.
 
-/// A node in the binary trie. In this schema, it is always 256 bits and is the hash of an
-/// underlying structure, represented as [`NodePreimage`].
+/// A node in the binary trie. In this schema, it is always 256 bits and is the hash of either
+/// an [`LeafData`] or [`InternalData`], or zeroed if it's a [`TERMINATOR`].
 ///
-/// The first bit of the [`Node`] is used to indicate the kind of node
+/// [`Node`]s are labeled by the [`NodeHasher`] used to indicate whether they are leaves or internal
+/// nodes. Typically, this is done by setting the MSB.
 pub type Node = [u8; 32];
 
 /// The path to a key. All paths have a 256 bit fixed length.
@@ -25,12 +26,6 @@ pub type KeyPath = [u8; 32];
 /// The hash of a value. In this schema, it is always 256 bits.
 pub type ValueHash = [u8; 32];
 
-/// The preimage of a node. In this schema, it is always 512 bits.
-///
-/// Note that this encoding itself does not contain information about
-/// whether the node is a leaf or internal node.
-pub type NodePreimage = [u8; 64];
-
 /// The terminator hash is a special node hash value denoting an empty sub-tree.
 /// Concretely, when this appears at a given location in the trie,
 /// it implies that no key with a path beginning with the location has a value.
@@ -39,18 +34,18 @@ pub type NodePreimage = [u8; 64];
 pub const TERMINATOR: Node = [0u8; 32];
 
 /// Whether the node hash indicates the node is a leaf.
-pub fn is_leaf(hash: &Node) -> bool {
-    hash[0] >> 7 == 1
+pub fn is_leaf<H: NodeHasher>(hash: &Node) -> bool {
+    H::node_kind(hash) == NodeKind::Leaf
 }
 
 /// Whether the node hash indicates the node is an internal node.
-pub fn is_internal(hash: &Node) -> bool {
-    hash[0] >> 7 == 0 && !is_terminator(&hash)
+pub fn is_internal<H: NodeHasher>(hash: &Node) -> bool {
+    H::node_kind(hash) == NodeKind::Internal
 }
 
 /// Whether the node holds the special `EMPTY_SUBTREE` value.
-pub fn is_terminator(hash: &Node) -> bool {
-    hash == &TERMINATOR
+pub fn is_terminator<H: NodeHasher>(hash: &Node) -> bool {
+    H::node_kind(hash) == NodeKind::Terminator
 }
 
 /// The kind of a node.
@@ -66,14 +61,8 @@ pub enum NodeKind {
 
 impl NodeKind {
     /// Get the kind of the provided node.
-    pub fn of(node: &Node) -> Self {
-        if is_leaf(node) {
-            NodeKind::Leaf
-        } else if is_terminator(node) {
-            NodeKind::Terminator
-        } else {
-            NodeKind::Internal
-        }
+    pub fn of<H: NodeHasher>(node: &Node) -> Self {
+        H::node_kind(node)
     }
 }
 
@@ -86,16 +75,6 @@ pub struct InternalData {
     pub right: Node,
 }
 
-impl InternalData {
-    /// Encode the internal node.
-    pub fn encode(&self) -> NodePreimage {
-        let mut node = [0u8; 64];
-        node[0..32].copy_from_slice(&self.left[..]);
-        node[32..64].copy_from_slice(&self.right[..]);
-        node
-    }
-}
-
 /// The data of a leaf node.
 #[derive(Debug, Default, Clone, PartialEq, Eq)]
 #[cfg_attr(
@@ -112,63 +91,24 @@ pub struct LeafData {
     pub value_hash: ValueHash,
 }
 
-impl LeafData {
-    /// Encode the leaf node.
-    pub fn encode(&self) -> NodePreimage {
-        let mut node = [0u8; 64];
-        self.encode_into(&mut node[..]);
-        node
-    }
-
-    /// Encode the leaf node into the given slice. It must have length at least 64 or this panics.
-    pub fn encode_into(&self, buf: &mut [u8]) {
-        buf[0..32].copy_from_slice(&self.key_path[..]);
-        buf[32..64].copy_from_slice(&self.value_hash[..]);
-    }
-
-    /// Decode the leaf node. Fails if the provided slice is not 64 bytes.
-    pub fn decode(buf: &[u8]) -> Option<Self> {
-        if buf.len() != 64 {
-            None
-        } else {
-            let mut leaf = LeafData {
-                key_path: Default::default(),
-                value_hash: Default::default(),
-            };
-            leaf.key_path.copy_from_slice(&buf[..32]);
-            leaf.value_hash.copy_from_slice(&buf[32..]);
-            Some(leaf)
-        }
-    }
-}
-
 /// A trie node hash function specialized for 64 bytes of data.
+///
+/// Note that it is illegal for the produced hash to equal [0; 32], as this value is reserved
+/// for the terminator node.
+///
+/// A node hasher should domain-separate internal and leaf nodes in some specific way. The
+/// recommended approach for binary hashes is to set the MSB to 0 or 1 depending on the node kind.
+/// However, for other kinds of hashes (e.g. Poseidon2 or other algebraic hashes), other labeling
+/// schemes may be required.
 pub trait NodeHasher {
-    /// Hash a node, encoded as exactly 64 bytes of data. This should not
-    /// domain-separate the hash.
-    fn hash_node(data: &NodePreimage) -> [u8; 32];
-}
+    /// Hash a leaf. This should domain-separate the hash
+    /// according to the node kind.
+    fn hash_leaf(data: &LeafData) -> [u8; 32];
 
-pub trait NodeHasherExt: NodeHasher {
-    /// Hash an internal node. This returns a domain-separated hash.
-    fn hash_internal(internal: &InternalData) -> Node {
-        let data = internal.encode();
-        let mut hash = Self::hash_node(&data);
+    /// Hash an internal node. This should domain-separate
+    /// the hash according to the node kind.
+    fn hash_internal(data: &InternalData) -> [u8; 32];
 
-        // set msb to 0.
-        hash[0] &= 0b01111111;
-        hash
-    }
-
-    /// Hash a leaf node. This returns a domain-separated hash.
-    fn hash_leaf(leaf: &LeafData) -> Node {
-        let data = leaf.encode();
-        let mut hash = Self::hash_node(&data);
-
-        // set msb to 1
-        hash[0] |= 0b10000000;
-        hash
-    }
+    /// Get the kind of the given node.
+    fn node_kind(node: &Node) -> NodeKind;
 }
-
-impl<T: NodeHasher> NodeHasherExt for T {}

core/src/update.rs

@@ -1,6 +1,6 @@
 //! Trie update logic helpers.
 
-use crate::trie::{self, KeyPath, LeafData, Node, NodeHasher, NodeHasherExt, ValueHash};
+use crate::trie::{self, KeyPath, LeafData, Node, NodeHasher, ValueHash};
 
 use bitvec::prelude::*;
 
@@ -252,16 +252,43 @@ pub fn build_trie<H: NodeHasher>(
 
 #[cfg(test)]
 mod tests {
-    use super::{
-        bitvec, build_trie, trie, BitVec, LeafData, Msb0, Node, NodeHasher, NodeHasherExt,
-        WriteNode,
-    };
+    use crate::trie::{NodeKind, TERMINATOR};
+
+    use super::{bitvec, build_trie, trie, BitVec, LeafData, Msb0, Node, NodeHasher, WriteNode};
 
     struct DummyNodeHasher;
 
     impl NodeHasher for DummyNodeHasher {
-        fn hash_node(data: &trie::NodePreimage) -> [u8; 32] {
-            blake3::hash(data).into()
+        fn hash_leaf(data: &trie::LeafData) -> [u8; 32] {
+            let mut hasher = blake3::Hasher::new();
+            hasher.update(&data.key_path);
+            hasher.update(&data.value_hash);
+            let mut hash: [u8; 32] = hasher.finalize().into();
+
+            // Label with MSB
+            hash[0] |= 0b10000000;
+            hash
+        }
+
+        fn hash_internal(data: &trie::InternalData) -> [u8; 32] {
+            let mut hasher = blake3::Hasher::new();
+            hasher.update(&data.left);
+            hasher.update(&data.right);
+            let mut hash: [u8; 32] = hasher.finalize().into();
+
+            // Label with MSB
+            hash[0] &= 0b01111111;
+            hash
+        }
+
+        fn node_kind(node: &Node) -> NodeKind {
+            if node[0] >> 7 == 1 {
+                NodeKind::Leaf
+            } else if node == &TERMINATOR {
+                NodeKind::Terminator
+            } else {
+                NodeKind::Internal
+            }
         }
     }
 

examples/witness_verification/src/main.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use nomt_core::{
     proof,
-    trie::{LeafData, NodeHasher},
+    trie::{self, LeafData, Node, NodeHasher, NodeKind, TERMINATOR},
 };
 
 // Hash nodes using blake3. The Hasher below will be utilized in the witness
@@ -10,8 +10,36 @@ use nomt_core::{
 pub struct Blake3Hasher;
 
 impl NodeHasher for Blake3Hasher {
-    fn hash_node(data: &nomt_core::trie::NodePreimage) -> [u8; 32] {
-        blake3::hash(data).into()
+    fn hash_leaf(data: &trie::LeafData) -> [u8; 32] {
+        let mut hasher = blake3::Hasher::new();
+        hasher.update(&data.key_path);
+        hasher.update(&data.value_hash);
+        let mut hash: [u8; 32] = hasher.finalize().into();
+
+        // Label with MSB
+        hash[0] |= 0b10000000;
+        hash
+    }
+
+    fn hash_internal(data: &trie::InternalData) -> [u8; 32] {
+        let mut hasher = blake3::Hasher::new();
+        hasher.update(&data.left);
+        hasher.update(&data.right);
+        let mut hash: [u8; 32] = hasher.finalize().into();
+
+        // Label with MSB
+        hash[0] &= 0b01111111;
+        hash
+    }
+
+    fn node_kind(node: &Node) -> NodeKind {
+        if node[0] >> 7 == 1 {
+            NodeKind::Leaf
+        } else if node == &TERMINATOR {
+            NodeKind::Terminator
+        } else {
+            NodeKind::Internal
+        }
     }
 }
 

fuzz/fuzz_targets/api_surface.rs

@@ -5,7 +5,9 @@ use std::collections::{HashMap, HashSet};
 use arbitrary::Arbitrary;
 use libfuzzer_sys::fuzz_target;
 
-use nomt::{Blake3Hasher, KeyPath, KeyReadWrite, Nomt, Options, SessionParams, Value, WitnessMode};
+use nomt::{
+    trie::KeyPath, Blake3Hasher, KeyReadWrite, Nomt, Options, SessionParams, Value, WitnessMode,
+};
 
 fuzz_target!(|run: Run| {
     let db = open_db(run.commit_concurrency);