diff --git a/Cargo.lock b/Cargo.lock index 294a1b0e64..76b36e08a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3023,6 +3023,7 @@ checksum = "707907fe3c25f5424cce2cb7e1cbcafee6bdbe735ca90ef77c29e84591e5b9da" dependencies = [ "equivalent", "hashbrown 0.15.0", + "serde", ] [[package]] @@ -6811,6 +6812,19 @@ dependencies = [ "tonic 0.12.3", ] +[[package]] +name = "tari_jellyfish" +version = "1.7.0-pre.3" +dependencies = [ + "borsh", + "digest 0.10.7", + "indexmap 2.6.0", + "serde", + "tari_crypto", + "tari_hashing", + "thiserror 2.0.3", +] + [[package]] name = "tari_key_manager" version = "1.7.0-pre.3" @@ -6996,6 +7010,7 @@ dependencies = [ "tari_common_types", "tari_crypto", "tari_hashing", + "tari_jellyfish", "tari_utilities", "thiserror 2.0.3", ] diff --git a/Cargo.toml b/Cargo.toml index 5da5560363..fccab33a70 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ members = [ "common_sqlite", "infrastructure/libtor", "infrastructure/metrics", + "infrastructure/jellyfish", "infrastructure/shutdown", "infrastructure/storage", "infrastructure/tari_script", @@ -44,6 +45,7 @@ members = [ ] [workspace.dependencies] +tari_jellyfish = { path = "infrastructure/jellyfish" } tari_comms = { path = "comms/core" } tari_comms_dht = { path = "comms/dht", default-features = false } tari_common = { path = "common" } diff --git a/applications/minotari_app_grpc/proto/sidechain_types.proto b/applications/minotari_app_grpc/proto/sidechain_types.proto index db9a0f28bb..857d94a396 100644 --- a/applications/minotari_app_grpc/proto/sidechain_types.proto +++ b/applications/minotari_app_grpc/proto/sidechain_types.proto @@ -97,6 +97,7 @@ message CommitProof { message CommitProofV1 { bytes command = 1; SidechainBlockCommitProof commit_proof = 2; + bytes encoded_inclusion_proof = 3; } message SidechainBlockCommitProof { diff --git a/applications/minotari_app_grpc/src/conversions/sidechain_feature.rs b/applications/minotari_app_grpc/src/conversions/sidechain_feature.rs index 464cf5bb97..82181d9a9d 100644 --- a/applications/minotari_app_grpc/src/conversions/sidechain_feature.rs +++ b/applications/minotari_app_grpc/src/conversions/sidechain_feature.rs @@ -361,6 +361,9 @@ impl TryFrom for CommandCommitProofV1 { Ok(CommandCommitProofV1 { command: command.try_into()?, commit_proof: value.commit_proof.ok_or("commit_proof not provided")?.try_into()?, + + inclusion_proof: borsh::from_slice(&value.encoded_inclusion_proof) + .map_err(|e| format!("Failed to decode SparseMerkleProofExt: {e}"))?, }) } } @@ -370,6 +373,10 @@ impl From<&CommandCommitProofV1> for grpc::CommitProofV1 { Self { command: grpc::EvictAtom::from(value.command()).encode_to_vec(), commit_proof: Some(value.commit_proof().into()), + // Encode since the type is complex + // TODO: making this fallible is a pain - we may need to implement the proto for this + encoded_inclusion_proof: borsh::to_vec(value.inclusion_proof()) + .expect("Failed to encode SparseMerkleProofExt"), } } } diff --git a/base_layer/core/src/proto/sidechain_feature.proto b/base_layer/core/src/proto/sidechain_feature.proto index 1971578d1c..b079d8d4b0 100644 --- a/base_layer/core/src/proto/sidechain_feature.proto +++ b/base_layer/core/src/proto/sidechain_feature.proto @@ -80,6 +80,7 @@ message CommitProof { message CommitProofV1 { bytes command = 1; SidechainBlockCommitProof commit_proof = 2; + bytes encoded_inclusion_proof = 3; } message SidechainBlockCommitProof { diff --git a/base_layer/core/src/proto/sidechain_feature.rs b/base_layer/core/src/proto/sidechain_feature.rs index 7a790f1652..af19b054fa 100644 --- a/base_layer/core/src/proto/sidechain_feature.rs +++ b/base_layer/core/src/proto/sidechain_feature.rs @@ -355,6 +355,8 @@ impl TryFrom for CommandCommitProofV1 for CommandCommitProofV1> for proto::types::CommitProofV1 { fn from(value: &CommandCommitProofV1) -> Self { Self { + // Encode since command is generic command: proto::types::EvictAtom::from(value.command()).encode_to_vec(), commit_proof: Some(value.commit_proof().into()), + // Encode since the type is complex + // TODO: making this fallible is a pain - we may need to implement the proto for this + encoded_inclusion_proof: borsh::to_vec(value.inclusion_proof()) + .expect("Failed to encode SparseMerkleProofExt"), } } } diff --git a/base_layer/sidechain/Cargo.toml b/base_layer/sidechain/Cargo.toml index 00a7516e60..7f9b10b2af 100644 --- a/base_layer/sidechain/Cargo.toml +++ b/base_layer/sidechain/Cargo.toml @@ -9,6 +9,7 @@ tari_hashing = { workspace = true } tari_crypto = { version = "0.21.0", features = ["borsh"] } tari_utilities = "0.8" tari_common_types = { workspace = true } +tari_jellyfish = { workspace = true } log = "0.4.22" thiserror = "2.0" diff --git a/base_layer/sidechain/src/commit_proof.rs b/base_layer/sidechain/src/commit_proof.rs index 9ea4c90fae..a8d9409e91 100644 --- a/base_layer/sidechain/src/commit_proof.rs +++ b/base_layer/sidechain/src/commit_proof.rs @@ -12,6 +12,7 @@ use tari_hashing::{ layer2::{block_hasher, vote_signature_hasher}, ValidatorNodeHashDomain, }; +use tari_jellyfish::{LeafKey, SparseMerkleProofExt, TreeHash}; use super::error::SidechainProofValidationError; use crate::{ @@ -29,8 +30,12 @@ pub enum CommandCommitProof { } impl CommandCommitProof { - pub fn new(command: C, commit_proof: SidechainBlockCommitProof) -> Self { - Self::V1(CommandCommitProofV1 { command, commit_proof }) + pub fn new(command: C, commit_proof: SidechainBlockCommitProof, inclusion_proof: SparseMerkleProofExt) -> Self { + Self::V1(CommandCommitProofV1 { + command, + commit_proof, + inclusion_proof, + }) } pub fn command(&self) -> &C { @@ -71,10 +76,9 @@ impl CommandCommitProof { #[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize, BorshSerialize, BorshDeserialize)] pub struct CommandCommitProofV1 { - // TODO: Implement MerkleProof - // command_merkle_proof: MerkleProof, pub command: C, pub commit_proof: SidechainBlockCommitProof, + pub inclusion_proof: SparseMerkleProofExt, } impl CommandCommitProofV1 { @@ -86,12 +90,26 @@ impl CommandCommitProofV1 { &self.commit_proof } + pub fn inclusion_proof(&self) -> &SparseMerkleProofExt { + &self.inclusion_proof + } + + fn validate_inclusion_proof(&self, command: &Command) -> Result<(), SidechainProofValidationError> { + let command_hash = TreeHash::new(command.hash().into_array()); + // Command JMT uses an identity mapping between hashes and keys. + let key = LeafKey::new(command_hash); + let root_hash = TreeHash::new(self.commit_proof.header.command_merkle_root.into_array()); + self.inclusion_proof.verify_inclusion(&root_hash, &key, &command_hash)?; + Ok(()) + } + pub fn validate_committed( &self, quorum_threshold: usize, check_vn: &CheckVnFunc<'_>, ) -> Result<(), SidechainProofValidationError> { let command = self.command.to_command(); + self.validate_inclusion_proof(&command)?; self.commit_proof .validate_committed(&command, quorum_threshold, check_vn) } diff --git a/base_layer/sidechain/src/error.rs b/base_layer/sidechain/src/error.rs index 65777d32f7..7268c0c6af 100644 --- a/base_layer/sidechain/src/error.rs +++ b/base_layer/sidechain/src/error.rs @@ -9,6 +9,8 @@ pub enum SidechainProofValidationError { InvalidProof { details: String }, #[error("Internal error: {details}")] InternalError { details: String }, + #[error("Jellyfish proof verification error: {0}")] + JmtProofVerifyError(#[from] tari_jellyfish::JmtProofVerifyError), } impl SidechainProofValidationError { diff --git a/hashing/src/layer2.rs b/hashing/src/layer2.rs index 72d1b537d2..762900e0f1 100644 --- a/hashing/src/layer2.rs +++ b/hashing/src/layer2.rs @@ -19,6 +19,10 @@ pub fn tari_hasher64(label: &'static str) -> TariDomainHash TariDomainHasher::::new_with_label(label) } +pub fn tari_hasher32(label: &'static str) -> TariDomainHasher { + TariDomainHasher::::new_with_label(label) +} + fn tari_consensus_hasher(label: &'static str) -> TariConsensusHasher { TariConsensusHasher::new_with_label(label) } diff --git a/infrastructure/jellyfish/Cargo.toml b/infrastructure/jellyfish/Cargo.toml new file mode 100644 index 0000000000..1d1cbc9968 --- /dev/null +++ b/infrastructure/jellyfish/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "tari_jellyfish" +version.workspace = true +edition.workspace = true +authors.workspace = true + +[dependencies] +tari_hashing = { workspace = true } +tari_crypto = "0.21.0" + +serde = { version = "1.0", features = ["derive"] } +borsh = "1.5" +digest = "0.10" +thiserror = "2.0" +indexmap = { version = "2.6", features = ["serde"] } \ No newline at end of file diff --git a/infrastructure/jellyfish/src/bit_iter.rs b/infrastructure/jellyfish/src/bit_iter.rs new file mode 100644 index 0000000000..bd77b2b890 --- /dev/null +++ b/infrastructure/jellyfish/src/bit_iter.rs @@ -0,0 +1,54 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +use std::ops::Range; + +/// An iterator over a hash value that generates one bit for each iteration. +pub struct BitIterator<'a> { + /// The reference to the bytes that represent the `HashValue`. + bytes: &'a [u8], + pos: Range, + // invariant hash_bytes.len() == HashValue::LENGTH; + // invariant pos.end == hash_bytes.len() * 8; +} + +impl<'a> BitIterator<'a> { + /// Constructs a new `BitIterator` using given `HashValue`. + pub fn new(bytes: &'a [u8]) -> Self { + BitIterator { + bytes, + pos: 0..bytes.len() * 8, + } + } + + /// Returns the `index`-th bit in the bytes. + fn get_bit(&self, index: usize) -> bool { + // MIRAI annotations - important? + // assume!(index < self.pos.end); // assumed precondition + // assume!(self.hash_bytes.len() == 32); // invariant + // assume!(self.pos.end == self.hash_bytes.len() * 8); // invariant + let pos = index / 8; + let bit = 7 - index % 8; + (self.bytes[pos] >> bit) & 1 != 0 + } +} + +impl<'a> Iterator for BitIterator<'a> { + type Item = bool; + + fn next(&mut self) -> Option { + self.pos.next().map(|x| self.get_bit(x)) + } + + fn size_hint(&self) -> (usize, Option) { + self.pos.size_hint() + } +} + +impl<'a> DoubleEndedIterator for BitIterator<'a> { + fn next_back(&mut self) -> Option { + self.pos.next_back().map(|x| self.get_bit(x)) + } +} + +impl<'a> ExactSizeIterator for BitIterator<'a> {} diff --git a/infrastructure/jellyfish/src/error.rs b/infrastructure/jellyfish/src/error.rs new file mode 100644 index 0000000000..ac13aff7cd --- /dev/null +++ b/infrastructure/jellyfish/src/error.rs @@ -0,0 +1,30 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +use crate::{LeafKey, TreeHash}; + +#[derive(Debug, thiserror::Error)] +pub enum JmtProofVerifyError { + #[error("Sparse Merkle Tree proof has more than 256 ({num_siblings}) siblings.")] + TooManySiblings { num_siblings: usize }, + #[error("Keys do not match. Key in proof: {actual_key}. Expected key: {expected_key}.")] + KeyMismatch { actual_key: LeafKey, expected_key: LeafKey }, + #[error("Value hashes do not match. Value hash in proof: {actual}. Expected value hash: {expected}.")] + ValueMismatch { actual: TreeHash, expected: TreeHash }, + #[error("Expected inclusion proof. Found non-inclusion proof.")] + ExpectedInclusionProof, + #[error("Expected non-inclusion proof, but key exists in proof.")] + ExpectedNonInclusionProof, + #[error( + "Key would not have ended up in the subtree where the provided key in proof is the only existing key, if it \ + existed. So this is not a valid non-inclusion proof." + )] + InvalidNonInclusionProof, + #[error( + "Root hashes do not match. Actual root hash: {actual_root_hash}. Expected root hash: {expected_root_hash}." + )] + RootHashMismatch { + actual_root_hash: TreeHash, + expected_root_hash: TreeHash, + }, +} diff --git a/infrastructure/jellyfish/src/hash.rs b/infrastructure/jellyfish/src/hash.rs new file mode 100644 index 0000000000..662769322e --- /dev/null +++ b/infrastructure/jellyfish/src/hash.rs @@ -0,0 +1,77 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +use std::{ + fmt::{Display, Formatter}, + ops::{Deref, DerefMut}, +}; + +use borsh::{BorshDeserialize, BorshSerialize}; +use serde::{Deserialize, Serialize}; + +#[derive( + Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord, BorshSerialize, BorshDeserialize, Serialize, Deserialize, +)] +pub struct TreeHash([u8; 32]); + +impl TreeHash { + pub const fn new(bytes: [u8; 32]) -> Self { + Self(bytes) + } + + pub const fn zero() -> Self { + Self([0; 32]) + } + + pub const fn into_array(self) -> [u8; 32] { + self.0 + } + + pub fn try_from_bytes(bytes: &[u8]) -> Result { + if bytes.len() != 32 { + return Err(TreeHashSizeError); + } + let mut arr = [0u8; 32]; + arr.copy_from_slice(bytes); + Ok(Self(arr)) + } +} + +impl From<[u8; 32]> for TreeHash { + fn from(bytes: [u8; 32]) -> Self { + Self(bytes) + } +} + +impl Deref for TreeHash { + type Target = [u8; 32]; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for TreeHash { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} + +impl> PartialEq for TreeHash { + fn eq(&self, other: &T) -> bool { + self.0 == other.as_ref() + } +} + +impl Display for TreeHash { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + for b in self.0 { + write!(f, "{:02x}", b)?; + } + Ok(()) + } +} + +#[derive(Debug, thiserror::Error)] +#[error("Invalid TreeHash byte size. Must be 32 bytes.")] +pub struct TreeHashSizeError; diff --git a/infrastructure/jellyfish/src/lib.rs b/infrastructure/jellyfish/src/lib.rs new file mode 100644 index 0000000000..09e4e491de --- /dev/null +++ b/infrastructure/jellyfish/src/lib.rs @@ -0,0 +1,21 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +mod hash; +pub use hash::*; + +mod tree; +pub use tree::*; + +mod types; +pub use types::*; + +mod error; +pub use error::*; + +mod store; + +pub use store::*; + +mod bit_iter; +pub mod memory_store; diff --git a/infrastructure/jellyfish/src/memory_store.rs b/infrastructure/jellyfish/src/memory_store.rs new file mode 100644 index 0000000000..dd0054c39b --- /dev/null +++ b/infrastructure/jellyfish/src/memory_store.rs @@ -0,0 +1,66 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +use std::{collections::HashMap, fmt, fmt::Debug}; + +use crate::{JmtStorageError, Node, NodeKey, StaleTreeNode, TreeNode, TreeStoreReader, TreeStoreWriter}; + +#[derive(Debug, Default)] +pub struct MemoryTreeStore

{ + pub nodes: HashMap>, + pub stale_nodes: Vec, +} + +impl

MemoryTreeStore

{ + pub fn new() -> Self { + Self { + nodes: HashMap::new(), + stale_nodes: Vec::new(), + } + } + + pub fn clear_stale_nodes(&mut self) { + for stale in self.stale_nodes.drain(..) { + self.nodes.remove(stale.as_node_key()); + } + } +} + +impl TreeStoreReader

for MemoryTreeStore

{ + fn get_node(&self, key: &NodeKey) -> Result, JmtStorageError> { + self.nodes + .get(key) + .map(|node| node.clone().into_node()) + .ok_or_else(|| JmtStorageError::NotFound(key.clone())) + } +} + +impl

TreeStoreWriter

for MemoryTreeStore

{ + fn insert_node(&mut self, key: NodeKey, node: Node

) -> Result<(), JmtStorageError> { + let node = TreeNode::new_latest(node); + self.nodes.insert(key, node); + Ok(()) + } + + fn record_stale_tree_node(&mut self, stale: StaleTreeNode) -> Result<(), JmtStorageError> { + self.stale_nodes.push(stale); + Ok(()) + } +} + +impl fmt::Display for MemoryTreeStore

{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + writeln!(f, "MemoryTreeStore")?; + writeln!(f, " Nodes:")?; + let mut store = self.nodes.iter().collect::>(); + store.sort_by_key(|(key, _)| *key); + for (key, node) in store { + writeln!(f, " {}: {:?}", key, node)?; + } + writeln!(f, " Stale Nodes:")?; + for stale in &self.stale_nodes { + writeln!(f, " {}", stale.as_node_key())?; + } + Ok(()) + } +} diff --git a/infrastructure/jellyfish/src/store.rs b/infrastructure/jellyfish/src/store.rs new file mode 100644 index 0000000000..73c77ccdb0 --- /dev/null +++ b/infrastructure/jellyfish/src/store.rs @@ -0,0 +1,92 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +use serde::{Deserialize, Serialize}; + +use crate::{JmtStorageError, Node, NodeKey}; + +/// Implementers are able to read nodes from a tree store. +pub trait TreeStoreReader

{ + /// Gets node by key, if it exists. + fn get_node(&self, key: &NodeKey) -> Result, JmtStorageError>; +} + +/// Implementers are able to insert nodes to a tree store. +pub trait TreeStoreWriter

{ + /// Inserts the node under a new, unique key (i.e. never an update). + fn insert_node(&mut self, key: NodeKey, node: Node

) -> Result<(), JmtStorageError>; + + /// Marks the given tree part for a (potential) future removal by an arbitrary external pruning + /// process. + fn record_stale_tree_node(&mut self, part: StaleTreeNode) -> Result<(), JmtStorageError>; +} + +/// Implementers are able to read and write nodes to a tree store. +pub trait TreeStore

: TreeStoreReader

+ TreeStoreWriter

{} +impl + TreeStoreWriter

> TreeStore

for S {} + +impl> TreeStoreReader

for &T { + fn get_node(&self, key: &NodeKey) -> Result, JmtStorageError> { + (*self).get_node(key) + } +} + +impl> TreeStoreWriter

for &mut T { + fn insert_node(&mut self, key: NodeKey, node: Node

) -> Result<(), JmtStorageError> { + (*self).insert_node(key, node) + } + + fn record_stale_tree_node(&mut self, part: StaleTreeNode) -> Result<(), JmtStorageError> { + (*self).record_stale_tree_node(part) + } +} + +/// A part of a tree that may become stale (i.e. need eventual pruning). +#[derive(Clone, PartialEq, Eq, Hash, Debug, Serialize, Deserialize)] +pub enum StaleTreeNode { + /// A single node to be removed. + Node(NodeKey), + /// An entire subtree of descendants of a specific node (including itself). + Subtree(NodeKey), +} + +impl StaleTreeNode { + pub fn into_node_key(self) -> NodeKey { + match self { + Self::Node(key) | Self::Subtree(key) => key, + } + } + + pub fn as_node_key(&self) -> &NodeKey { + match self { + Self::Node(key) | Self::Subtree(key) => key, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum TreeNode

{ + V1(Node

), +} + +impl

TreeNode

{ + pub fn new_latest(node: Node

) -> Self { + Self::new_v1(node) + } + + pub fn new_v1(node: Node

) -> Self { + Self::V1(node) + } + + pub fn as_node(&self) -> &Node

{ + match self { + Self::V1(node) => node, + } + } + + pub fn into_node(self) -> Node

{ + match self { + Self::V1(node) => node, + } + } +} diff --git a/infrastructure/jellyfish/src/tree.rs b/infrastructure/jellyfish/src/tree.rs new file mode 100644 index 0000000000..18216ff073 --- /dev/null +++ b/infrastructure/jellyfish/src/tree.rs @@ -0,0 +1,790 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +// Copyright 2021 Radix Publishing Ltd incorporated in Jersey (Channel Islands). +// +// Licensed under the Radix License, Version 1.0 (the "License"); you may not use this +// file except in compliance with the License. You may obtain a copy of the License at: +// +// radixfoundation.org/licenses/LICENSE-v1 +// +// The Licensor hereby grants permission for the Canonical version of the Work to be +// published, distributed and used under or by reference to the Licensor's trademark +// Radix ® and use of any unregistered trade names, logos or get-up. +// +// The Licensor provides the Work (and each Contributor provides its Contributions) on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, +// including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, +// MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. +// +// Whilst the Work is capable of being deployed, used and adopted (instantiated) to create +// a distributed ledger it is your responsibility to test and validate the code, together +// with all logic and performance of that code under all foreseeable scenarios. +// +// The Licensor does not make or purport to make and hereby excludes liability for all +// and any representation, warranty or undertaking in any form whatsoever, whether express +// or implied, to any entity or person, including any representation, warranty or +// undertaking, as to the functionality security use, value or other characteristics of +// any distributed ledger nor in respect the functioning or value of any tokens which may +// be created stored or transferred using the Work. The Licensor does not warrant that the +// Work or any use of the Work complies with any law or regulation in any territory where +// it may be implemented or used or that it will be appropriate for any specific purpose. +// +// Neither the licensor nor any current or former employees, officers, directors, partners, +// trustees, representatives, agents, advisors, contractors, or volunteers of the Licensor +// shall be liable for any direct or indirect, special, incidental, consequential or other +// losses of any kind, in tort, contract or otherwise (including but not limited to loss +// of revenue, income or profits, or loss of use or data, or loss of reputation, or loss +// of any economic or other opportunity of whatsoever nature or howsoever arising), arising +// out of or in connection with (without limitation of any use, misuse, of any ledger system +// or use made or its functionality or any performance or operation of any code or protocol +// caused by bugs or programming or logic errors or otherwise); +// +// A. any offer, purchase, holding, use, sale, exchange or transmission of any +// cryptographic keys, tokens or assets created, exchanged, stored or arising from any +// interaction with the Work; +// +// B. any failure in a transmission or loss of any token or assets keys or other digital +// artefacts due to errors in transmission; +// +// C. bugs, hacks, logic errors or faults in the Work or any communication; +// +// D. system software or apparatus including but not limited to losses caused by errors +// in holding or transmitting tokens by any third-party; +// +// E. breaches or failure of security including hacker attacks, loss or disclosure of +// password, loss of private key, unauthorised use or misuse of such passwords or keys; +// +// F. any losses including loss of anticipated savings or other benefits resulting from +// use of the Work or any changes to the Work (however implemented). +// +// You are solely responsible for; testing, validating and evaluation of all operation +// logic, functionality, security and appropriateness of using the Work for any commercial +// or non-commercial purpose and for any reproduction or redistribution by You of the +// Work. You assume all risks associated with Your use of the Work and the exercise of +// permissions under this License. + +// This file contains code sourced from https://github.com/aptos-labs/aptos-core/tree/1.0.4 +// This original source is licensed under https://github.com/aptos-labs/aptos-core/blob/1.0.4/LICENSE +// +// The code in this file has been implemented by Radix® pursuant to an Apache 2 licence and has +// been modified by Radix® and is now licensed pursuant to the Radix® Open-Source Licence. +// +// Each sourced code fragment includes an inline attribution to the original source file in a +// comment starting "SOURCE: ..." +// +// Modifications from the original source are captured in two places: +// * Initial changes to get the code functional/integrated are marked by inline "INITIAL-MODIFICATION: ..." comments +// * Subsequent changes to the code are captured in the git commit history +// +// The following notice is retained from the original source +// Copyright (c) Aptos +// SPDX-License-Identifier: Apache-2.0 + +use std::{ + collections::{BTreeMap, HashMap}, + marker::PhantomData, +}; + +use super::{ + store::TreeStoreReader, + types::{ + Child, + InternalNode, + IteratedLeafKey, + JmtStorageError, + LeafKey, + LeafNode, + Nibble, + NibblePath, + Node, + NodeKey, + SparseMerkleProof, + SparseMerkleProofExt, + SparseMerkleRangeProof, + Version, + SPARSE_MERKLE_PLACEHOLDER_HASH, + }, + LeafKeyRef, + TreeHash, +}; + +// INITIAL-MODIFICATION: the original used a known key size (32) as a limit +const SANITY_NIBBLE_LIMIT: usize = 1000; + +pub type ProofValue

= (TreeHash, P, Version); + +// SOURCE: https://github.com/radixdlt/radixdlt-scrypto/blob/ca8e553c31a956c0851c1855291efe4a47fb5c97/radix-engine-stores/src/hash_tree/jellyfish.rs +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/storage/jellyfish-merkle/src/lib.rs#L329 +/// The Jellyfish Merkle tree data structure. See [`crate`] for description. +pub struct JellyfishMerkleTree<'a, R, P> { + reader: &'a R, + _payload: PhantomData

, +} + +impl<'a, R: 'a + TreeStoreReader

, P: Clone> JellyfishMerkleTree<'a, R, P> { + /// Creates a `JellyfishMerkleTree` backed by the given [`TreeReader`](trait.TreeReader.html). + pub fn new(reader: &'a R) -> Self { + Self { + reader, + _payload: PhantomData, + } + } + + /// Get the node hash from the cache if cache is provided, otherwise (for test only) compute it. + fn get_hash(node_key: &NodeKey, node: &Node

, hash_cache: Option<&HashMap>) -> TreeHash { + if let Some(cache) = hash_cache { + match cache.get(node_key.nibble_path()) { + Some(hash) => *hash, + None => unreachable!("{:?} can not be found in hash cache", node_key), + } + } else { + node.hash() + } + } + + /// For each value set: + /// Returns the new nodes and values in a batch after applying `value_set`. For + /// example, if after transaction `T_i` the committed state of tree in the persistent storage + /// looks like the following structure: + /// + /// ```text + /// S_i + /// / \ + /// . . + /// . . + /// / \ + /// o x + /// / \ + /// A B + /// storage (disk) + /// ``` + /// + /// where `A` and `B` denote the states of two adjacent accounts, and `x` is a sibling subtree + /// of the path from root to A and B in the tree. Then a `value_set` produced by the next + /// transaction `T_{i+1}` modifies other accounts `C` and `D` exist in the subtree under `x`, a + /// new partial tree will be constructed in memory and the structure will be: + /// + /// ```text + /// S_i | S_{i+1} + /// / \ | / \ + /// . . | . . + /// . . | . . + /// / \ | / \ + /// / x | / x' + /// o<-------------+- / \ + /// / \ | C D + /// A B | + /// storage (disk) | cache (memory) + /// ``` + /// + /// With this design, we are able to query the global state in persistent storage and + /// generate the proposed tree delta based on a specific root hash and `value_set`. For + /// example, if we want to execute another transaction `T_{i+1}'`, we can use the tree `S_i` in + /// storage and apply the `value_set` of transaction `T_{i+1}`. Then if the storage commits + /// the returned batch, the state `S_{i+1}` is ready to be read from the tree by calling + /// [`get_with_proof`](struct.JellyfishMerkleTree.html#method.get_with_proof). Anything inside + /// the batch is not reachable from public interfaces before being committed. + pub fn batch_put_value_set)>>( + &self, + value_set: I, + node_hashes: Option<&HashMap>, + persisted_version: Option, + version: Version, + ) -> Result<(TreeHash, TreeUpdateBatch

), JmtStorageError> { + let value_set = value_set.into_iter().collect::>(); + let deduped_and_sorted_kvs = value_set.iter().map(|(k, v)| (k, v.as_ref())).collect::>(); + + let mut batch = TreeUpdateBatch::new(); + let root_node_opt = if let Some(persisted_version) = persisted_version { + self.batch_insert_at( + &NodeKey::new_empty_path(persisted_version), + version, + deduped_and_sorted_kvs.as_slice(), + 0, + node_hashes, + &mut batch, + )? + } else { + Self::batch_update_subtree( + &NodeKey::new_empty_path(version), + version, + deduped_and_sorted_kvs.as_slice(), + 0, + node_hashes, + &mut batch, + )? + }; + + let node_key = NodeKey::new_empty_path(version); + let root_hash = if let Some(root_node) = root_node_opt { + let hash = root_node.hash(); + batch.put_node(node_key, root_node); + hash + } else { + batch.put_node(node_key, Node::Null); + SPARSE_MERKLE_PLACEHOLDER_HASH + }; + + Ok((root_hash, batch)) + } + + fn batch_insert_at( + &self, + node_key: &NodeKey, + version: Version, + kvs: &[(&LeafKey, Option<&(TreeHash, P)>)], + depth: usize, + hash_cache: Option<&HashMap>, + batch: &mut TreeUpdateBatch

, + ) -> Result>, JmtStorageError> { + let node = self.reader.get_node(node_key)?; + batch.put_stale_node(node_key.clone(), version, &node); + + match node { + Node::Internal(internal_node) => { + // There is a small possibility that the old internal node is intact. + // Traverse all the path touched by `kvs` from this internal node. + let range_iter = NibbleRangeIterator::new(kvs, depth); + // INITIAL-MODIFICATION: there was a par_iter (conditionally) used here + let new_children = range_iter + .map(|(left, right)| { + self.insert_at_child( + node_key, + &internal_node, + version, + kvs, + left, + right, + depth, + hash_cache, + batch, + ) + }) + .collect::, JmtStorageError>>()?; + + // Reuse the current `InternalNode` in memory to create a new internal node. + let mut old_children = internal_node.into_children(); + let mut new_created_children: Vec<(Nibble, Node

)> = Vec::new(); + for (child_nibble, child_option) in new_children { + if let Some(child) = child_option { + new_created_children.push((child_nibble, child)); + } else { + old_children.swap_remove(&child_nibble); + } + } + + if old_children.is_empty() && new_created_children.is_empty() { + return Ok(None); + } + if old_children.len() <= 1 && new_created_children.len() <= 1 { + if let Some((new_nibble, new_child)) = new_created_children.first() { + if let Some((old_nibble, _old_child)) = old_children.iter().next() { + if old_nibble == new_nibble && new_child.is_leaf() { + return Ok(Some(new_child.clone())); + } + } else if new_child.is_leaf() { + return Ok(Some(new_child.clone())); + } else { + // Nothing to do + } + } else { + let (old_child_nibble, old_child) = old_children.iter().next().expect("must exist"); + if old_child.is_leaf() { + let old_child_node_key = node_key.gen_child_node_key(old_child.version, *old_child_nibble); + let old_child_node = self.reader.get_node(&old_child_node_key)?; + batch.put_stale_node(old_child_node_key, version, &old_child_node); + return Ok(Some(old_child_node)); + } + } + } + + let mut new_children = old_children; + for (child_index, new_child_node) in new_created_children { + let new_child_node_key = node_key.gen_child_node_key(version, child_index); + new_children.insert( + child_index, + Child::new( + Self::get_hash(&new_child_node_key, &new_child_node, hash_cache), + version, + new_child_node.node_type(), + ), + ); + batch.put_node(new_child_node_key, new_child_node); + } + let new_internal_node = InternalNode::new(new_children); + Ok(Some(new_internal_node.into())) + }, + Node::Leaf(leaf_node) => Self::batch_update_subtree_with_existing_leaf( + node_key, version, leaf_node, kvs, depth, hash_cache, batch, + ), + Node::Null => { + assert_eq!(depth, 0, "Null node can only exist at depth 0"); + Self::batch_update_subtree(node_key, version, kvs, 0, hash_cache, batch) + }, + } + } + + fn insert_at_child( + &self, + node_key: &NodeKey, + internal_node: &InternalNode, + version: Version, + kvs: &[(&LeafKey, Option<&(TreeHash, P)>)], + left: usize, + right: usize, + depth: usize, + hash_cache: Option<&HashMap>, + batch: &mut TreeUpdateBatch

, + ) -> Result<(Nibble, Option>), JmtStorageError> { + let child_index = kvs[left].0.get_nibble(depth); + let child = internal_node.child(child_index); + + let new_child_node_option = match child { + Some(child) => self.batch_insert_at( + &node_key.gen_child_node_key(child.version, child_index), + version, + &kvs[left..=right], + depth + 1, + hash_cache, + batch, + )?, + None => Self::batch_update_subtree( + &node_key.gen_child_node_key(version, child_index), + version, + &kvs[left..=right], + depth + 1, + hash_cache, + batch, + )?, + }; + + Ok((child_index, new_child_node_option)) + } + + fn batch_update_subtree_with_existing_leaf( + node_key: &NodeKey, + version: Version, + existing_leaf_node: LeafNode

, + kvs: &[(&LeafKey, Option<&(TreeHash, P)>)], + depth: usize, + hash_cache: Option<&HashMap>, + batch: &mut TreeUpdateBatch

, + ) -> Result>, JmtStorageError> { + let existing_leaf_key = existing_leaf_node.leaf_key(); + + if kvs.len() == 1 && kvs[0].0 == existing_leaf_key { + if let (key, Some((value_hash, payload))) = kvs[0] { + let new_leaf_node = Node::new_leaf(*key, *value_hash, payload.clone(), version); + Ok(Some(new_leaf_node)) + } else { + Ok(None) + } + } else { + let existing_leaf_bucket = existing_leaf_key.get_nibble(depth); + let mut isolated_existing_leaf = true; + let mut children = vec![]; + for (left, right) in NibbleRangeIterator::new(kvs, depth) { + let child_index = kvs[left].0.get_nibble(depth); + let child_node_key = node_key.gen_child_node_key(version, child_index); + if let Some(new_child_node) = if existing_leaf_bucket == child_index { + isolated_existing_leaf = false; + Self::batch_update_subtree_with_existing_leaf( + &child_node_key, + version, + existing_leaf_node.clone(), + &kvs[left..=right], + depth + 1, + hash_cache, + batch, + )? + } else { + Self::batch_update_subtree( + &child_node_key, + version, + &kvs[left..=right], + depth + 1, + hash_cache, + batch, + )? + } { + children.push((child_index, new_child_node)); + } + } + if isolated_existing_leaf { + children.push((existing_leaf_bucket, existing_leaf_node.into())); + } + + if children.is_empty() { + Ok(None) + } else if children.len() == 1 && children[0].1.is_leaf() { + let (_, child) = children.pop().expect("Must exist"); + Ok(Some(child)) + } else { + let new_internal_node = InternalNode::new( + children + .into_iter() + .map(|(child_index, new_child_node)| { + let new_child_node_key = node_key.gen_child_node_key(version, child_index); + let result = ( + child_index, + Child::new( + Self::get_hash(&new_child_node_key, &new_child_node, hash_cache), + version, + new_child_node.node_type(), + ), + ); + batch.put_node(new_child_node_key, new_child_node); + result + }) + .collect(), + ); + Ok(Some(new_internal_node.into())) + } + } + } + + fn batch_update_subtree( + node_key: &NodeKey, + version: Version, + kvs: &[(&LeafKey, Option<&(TreeHash, P)>)], + depth: usize, + hash_cache: Option<&HashMap>, + batch: &mut TreeUpdateBatch

, + ) -> Result>, JmtStorageError> { + if kvs.len() == 1 { + if let (key, Some((value_hash, payload))) = kvs[0] { + let new_leaf_node = Node::new_leaf(*key, *value_hash, payload.clone(), version); + Ok(Some(new_leaf_node)) + } else { + Ok(None) + } + } else { + let mut children = vec![]; + for (left, right) in NibbleRangeIterator::new(kvs, depth) { + let child_index = kvs[left].0.get_nibble(depth); + let child_node_key = node_key.gen_child_node_key(version, child_index); + if let Some(new_child_node) = Self::batch_update_subtree( + &child_node_key, + version, + &kvs[left..=right], + depth + 1, + hash_cache, + batch, + )? { + children.push((child_index, new_child_node)) + } + } + if children.is_empty() { + Ok(None) + } else if children.len() == 1 && children[0].1.is_leaf() { + let (_, child) = children.pop().expect("Must exist"); + Ok(Some(child)) + } else { + let new_internal_node = InternalNode::new( + children + .into_iter() + .map(|(child_index, new_child_node)| { + let new_child_node_key = node_key.gen_child_node_key(version, child_index); + let result = ( + child_index, + Child::new( + Self::get_hash(&new_child_node_key, &new_child_node, hash_cache), + version, + new_child_node.node_type(), + ), + ); + batch.put_node(new_child_node_key, new_child_node); + result + }) + .collect(), + ); + Ok(Some(new_internal_node.into())) + } + } + } + + /// Returns the value (if applicable) and the corresponding merkle proof. + pub fn get_with_proof( + &self, + key: LeafKeyRef<'_>, + version: Version, + ) -> Result<(Option>, SparseMerkleProof), JmtStorageError> { + self.get_with_proof_ext(key, version) + .map(|(value, proof_ext)| (value, proof_ext.into())) + } + + pub fn get_with_proof_ext( + &self, + key: LeafKeyRef<'_>, + version: Version, + ) -> Result<(Option>, SparseMerkleProofExt), JmtStorageError> { + // Empty tree just returns proof with no sibling hash. + let mut next_node_key = NodeKey::new_empty_path(version); + let mut siblings = vec![]; + let nibble_path = NibblePath::new_even(key.bytes.to_vec()); + let mut nibble_iter = nibble_path.nibbles(); + + for _nibble_depth in 0..SANITY_NIBBLE_LIMIT { + let next_node = self.reader.get_node(&next_node_key)?; + match next_node { + Node::Internal(internal_node) => { + let queried_child_index = nibble_iter.next().ok_or(JmtStorageError::InconsistentState)?; + let (child_node_key, mut siblings_in_internal) = internal_node.get_child_with_siblings( + &next_node_key, + queried_child_index, + Some(self.reader), + )?; + siblings.append(&mut siblings_in_internal); + next_node_key = match child_node_key { + Some(node_key) => node_key, + None => { + return Ok(( + None, + SparseMerkleProofExt::new(None, { + siblings.reverse(); + siblings + }), + )) + }, + }; + }, + Node::Leaf(leaf_node) => { + return Ok(( + if leaf_node.leaf_key().as_ref() == key { + Some((leaf_node.value_hash(), leaf_node.payload().clone(), leaf_node.version())) + } else { + None + }, + SparseMerkleProofExt::new(Some(leaf_node.into()), { + siblings.reverse(); + siblings + }), + )); + }, + Node::Null => { + return Ok((None, SparseMerkleProofExt::new(None, vec![]))); + }, + } + } + Err(JmtStorageError::InconsistentState) + } + + /// Gets the proof that shows a list of keys up to `rightmost_key_to_prove` exist at `version`. + pub fn get_range_proof( + &self, + rightmost_key_to_prove: LeafKeyRef<'_>, + version: Version, + ) -> Result { + let (leaf, proof) = self.get_with_proof(rightmost_key_to_prove, version)?; + assert!(leaf.is_some(), "rightmost_key_to_prove must exist."); + + let siblings = proof + .siblings() + .iter() + .rev() + .zip(rightmost_key_to_prove.iter_bits()) + .filter_map(|(sibling, bit)| { + // We only need to keep the siblings on the right. + if bit { + None + } else { + Some(*sibling) + } + }) + .rev() + .collect(); + Ok(SparseMerkleRangeProof::new(siblings)) + } + + fn get_root_node(&self, version: Version) -> Result, JmtStorageError> { + let root_node_key = NodeKey::new_empty_path(version); + self.reader.get_node(&root_node_key) + } + + pub fn get_root_hash(&self, version: Version) -> Result { + self.get_root_node(version).map(|n| n.hash()) + } + + pub fn get_leaf_count(&self, version: Version) -> Result { + self.get_root_node(version).map(|n| n.leaf_count()) + } + + pub fn get_all_nodes_referenced(&self, key: NodeKey) -> Result, JmtStorageError> { + let mut out_keys = vec![]; + self.get_all_nodes_referenced_impl(key, &mut out_keys)?; + Ok(out_keys) + } + + fn get_all_nodes_referenced_impl(&self, key: NodeKey, out_keys: &mut Vec) -> Result<(), JmtStorageError> { + match self.reader.get_node(&key)? { + Node::Internal(internal_node) => { + for (child_nibble, child) in internal_node.children_sorted() { + self.get_all_nodes_referenced_impl(key.gen_child_node_key(child.version, *child_nibble), out_keys)?; + } + }, + Node::Leaf(_) | Node::Null => {}, + }; + + out_keys.push(key); + Ok(()) + } +} + +/// An iterator that iterates the index range (inclusive) of each different nibble at given +/// `nibble_idx` of all the keys in a sorted key-value pairs which have the identical Hash +/// prefix (up to nibble_idx). +struct NibbleRangeIterator<'a, P> { + sorted_kvs: &'a [(&'a LeafKey, P)], + nibble_idx: usize, + pos: usize, +} + +impl<'a, P> NibbleRangeIterator<'a, P> { + fn new(sorted_kvs: &'a [(&'a LeafKey, P)], nibble_idx: usize) -> Self { + NibbleRangeIterator { + sorted_kvs, + nibble_idx, + pos: 0, + } + } +} + +impl<'a, P> Iterator for NibbleRangeIterator<'a, P> { + type Item = (usize, usize); + + fn next(&mut self) -> Option { + let left = self.pos; + if self.pos < self.sorted_kvs.len() { + let cur_nibble = self.sorted_kvs[left].0.get_nibble(self.nibble_idx); + let (mut i, mut j) = (left, self.sorted_kvs.len() - 1); + // Find the last index of the cur_nibble. + while i < j { + let mid = j - (j - i) / 2; + if self.sorted_kvs[mid].0.get_nibble(self.nibble_idx) > cur_nibble { + j = mid - 1; + } else { + i = mid; + } + } + self.pos = i + 1; + Some((left, i)) + } else { + None + } + } +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct TreeUpdateBatch

{ + pub node_batch: Vec<(NodeKey, Node

)>, + pub stale_node_index_batch: Vec, + pub num_new_leaves: usize, + pub num_stale_leaves: usize, +} + +impl TreeUpdateBatch

{ + pub fn new() -> Self { + Self { + node_batch: vec![], + stale_node_index_batch: vec![], + num_new_leaves: 0, + num_stale_leaves: 0, + } + } + + fn inc_num_new_leaves(&mut self) { + self.num_new_leaves += 1; + } + + fn inc_num_stale_leaves(&mut self) { + self.num_stale_leaves += 1; + } + + pub fn put_node(&mut self, node_key: NodeKey, node: Node

) { + if node.is_leaf() { + self.inc_num_new_leaves(); + } + self.node_batch.push((node_key, node)) + } + + pub fn put_stale_node(&mut self, node_key: NodeKey, stale_since_version: Version, node: &Node

) { + if node.is_leaf() { + self.inc_num_stale_leaves(); + } + self.stale_node_index_batch.push(StaleNodeIndex { + node_key, + stale_since_version, + }); + } +} + +/// Indicates a node becomes stale since `stale_since_version`. +#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct StaleNodeIndex { + /// The version since when the node is overwritten and becomes stale. + pub stale_since_version: Version, + /// The [`NodeKey`](node_type/struct.NodeKey.html) identifying the node associated with this + /// record. + pub node_key: NodeKey, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{jmt_node_hash, memory_store::MemoryTreeStore, StaleTreeNode, TreeStoreWriter}; + + fn leaf_key(seed: u64) -> LeafKey { + LeafKey::new(jmt_node_hash(&seed)) + } + + #[test] + fn check_merkle_proof() { + // Evaluating the functionality of the JMT Merkle proof. + let mut mem = MemoryTreeStore::new(); + let jmt = JellyfishMerkleTree::new(&mem); + + let values = [ + (leaf_key(1), Some((jmt_node_hash(&10), Some(1u64)))), + (leaf_key(2), Some((jmt_node_hash(&11), Some(2)))), + (leaf_key(3), Some((jmt_node_hash(&12), Some(3)))), + ]; + let (_, diff) = jmt.batch_put_value_set(values, None, None, 1).unwrap(); + for (k, v) in diff.node_batch { + mem.insert_node(k, v).unwrap(); + } + + for a in diff.stale_node_index_batch { + mem.record_stale_tree_node(StaleTreeNode::Node(a.node_key)).unwrap(); + } + mem.clear_stale_nodes(); + + let jmt = JellyfishMerkleTree::new(&mem); + + // This causes get_with_proof to fail with node NotFound. + let values = [ + (leaf_key(4), Some((jmt_node_hash(&13), Some(4u64)))), + (leaf_key(5), Some((jmt_node_hash(&14), Some(5)))), + (leaf_key(6), Some((jmt_node_hash(&15), Some(6)))), + ]; + let (_mr, diff) = jmt.batch_put_value_set(values, None, Some(1), 2).unwrap(); + + for (k, v) in diff.node_batch { + mem.insert_node(k, v).unwrap(); + } + for a in diff.stale_node_index_batch { + mem.record_stale_tree_node(StaleTreeNode::Node(a.node_key)).unwrap(); + } + mem.clear_stale_nodes(); + let jmt = JellyfishMerkleTree::new(&mem); + + let k = leaf_key(3); + let (_value, sparse) = jmt.get_with_proof(k.as_ref(), 2).unwrap(); + + let leaf = sparse.leaf().unwrap(); + assert_eq!(*leaf.key(), k); + assert_eq!(*leaf.value_hash(), jmt_node_hash(&12)); + // Unanswered: How do we verify the proof root matches a Merkle root? + // assert!(sparse.siblings().iter().any(|h| *h == mr)); + } +} diff --git a/infrastructure/jellyfish/src/types.rs b/infrastructure/jellyfish/src/types.rs new file mode 100644 index 0000000000..57cc213e4f --- /dev/null +++ b/infrastructure/jellyfish/src/types.rs @@ -0,0 +1,1374 @@ +// Copyright 2024 The Tari Project +// SPDX-License-Identifier: BSD-3-Clause + +// Copyright 2021 Radix Publishing Ltd incorporated in Jersey (Channel Islands). +// +// Licensed under the Radix License, Version 1.0 (the "License"); you may not use this +// file except in compliance with the License. You may obtain a copy of the License at: +// +// radixfoundation.org/licenses/LICENSE-v1 +// +// The Licensor hereby grants permission for the Canonical version of the Work to be +// published, distributed and used under or by reference to the Licensor's trademark +// Radix ® and use of any unregistered trade names, logos or get-up. +// +// The Licensor provides the Work (and each Contributor provides its Contributions) on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, +// including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, +// MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. +// +// Whilst the Work is capable of being deployed, used and adopted (instantiated) to create +// a distributed ledger it is your responsibility to test and validate the code, together +// with all logic and performance of that code under all foreseeable scenarios. +// +// The Licensor does not make or purport to make and hereby excludes liability for all +// and any representation, warranty or undertaking in any form whatsoever, whether express +// or implied, to any entity or person, including any representation, warranty or +// undertaking, as to the functionality security use, value or other characteristics of +// any distributed ledger nor in respect the functioning or value of any tokens which may +// be created stored or transferred using the Work. The Licensor does not warrant that the +// Work or any use of the Work complies with any law or regulation in any territory where +// it may be implemented or used or that it will be appropriate for any specific purpose. +// +// Neither the licensor nor any current or former employees, officers, directors, partners, +// trustees, representatives, agents, advisors, contractors, or volunteers of the Licensor +// shall be liable for any direct or indirect, special, incidental, consequential or other +// losses of any kind, in tort, contract or otherwise (including but not limited to loss +// of revenue, income or profits, or loss of use or data, or loss of reputation, or loss +// of any economic or other opportunity of whatsoever nature or howsoever arising), arising +// out of or in connection with (without limitation of any use, misuse, of any ledger system +// or use made or its functionality or any performance or operation of any code or protocol +// caused by bugs or programming or logic errors or otherwise); +// +// A. any offer, purchase, holding, use, sale, exchange or transmission of any +// cryptographic keys, tokens or assets created, exchanged, stored or arising from any +// interaction with the Work; +// +// B. any failure in a transmission or loss of any token or assets keys or other digital +// artefacts due to errors in transmission; +// +// C. bugs, hacks, logic errors or faults in the Work or any communication; +// +// D. system software or apparatus including but not limited to losses caused by errors +// in holding or transmitting tokens by any third-party; +// +// E. breaches or failure of security including hacker attacks, loss or disclosure of +// password, loss of private key, unauthorised use or misuse of such passwords or keys; +// +// F. any losses including loss of anticipated savings or other benefits resulting from +// use of the Work or any changes to the Work (however implemented). +// +// You are solely responsible for; testing, validating and evaluation of all operation +// logic, functionality, security and appropriateness of using the Work for any commercial +// or non-commercial purpose and for any reproduction or redistribution by You of the +// Work. You assume all risks associated with Your use of the Work and the exercise of +// permissions under this License. + +// This file contains code sourced from https://github.com/aptos-labs/aptos-core/tree/1.0.4 +// This original source is licensed under https://github.com/aptos-labs/aptos-core/blob/1.0.4/LICENSE +// +// The code in this file has been implemented by Radix® pursuant to an Apache 2 licence and has +// been modified by Radix® and is now licensed pursuant to the Radix® Open-Source Licence. +// +// Each sourced code fragment includes an inline attribution to the original source file in a +// comment starting "SOURCE: ..." +// +// Modifications from the original source are captured in two places: +// * Initial changes to get the code functional/integrated are marked by inline "INITIAL-MODIFICATION: ..." comments +// * Subsequent changes to the code are captured in the git commit history +// +// The following notice is retained from the original source +// Copyright (c) Aptos +// SPDX-License-Identifier: Apache-2.0 + +use std::{fmt, fmt::Display, ops::Range}; + +use borsh::{BorshDeserialize, BorshSerialize}; +use digest::consts; +use indexmap::IndexMap; +use serde::{Deserialize, Serialize}; +use tari_crypto::hash_domain; +use tari_hashing::layer2::{tari_hasher32, TariDomainHasher}; + +use crate::{bit_iter::BitIterator, error::JmtProofVerifyError, store::TreeStoreReader, TreeHash}; + +hash_domain!(JmtHashDomain, "com.tari.jmt", 0); + +pub type JmtHasher = TariDomainHasher; + +fn jmt_node_hasher() -> JmtHasher { + tari_hasher32::("Node") +} + +pub fn jmt_node_hash(data: &T) -> TreeHash { + jmt_node_hasher().chain(data).finalize_into_array().into() +} + +pub fn jmt_node_hash2(d1: &TreeHash, d2: &TreeHash) -> TreeHash { + jmt_node_hasher().chain(d1).chain(d2).finalize_into_array().into() +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/proof/definition.rs#L182 +/// A more detailed version of `SparseMerkleProof` with the only difference that all the leaf +/// siblings are explicitly set as `SparseMerkleLeafNode` instead of its hash value. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)] +pub struct SparseMerkleProofExt { + leaf: Option, + /// All siblings in this proof, including the default ones. Siblings are ordered from the bottom + /// level to the root level. + siblings: Vec, +} + +impl SparseMerkleProofExt { + /// Constructs a new `SparseMerkleProofExt` using leaf and a list of sibling nodes. + pub(crate) fn new(leaf: Option, siblings: Vec) -> Self { + Self { leaf, siblings } + } + + /// Returns the leaf node in this proof. + pub fn leaf(&self) -> Option { + self.leaf.clone() + } + + /// Returns the list of siblings in this proof. + pub fn siblings(&self) -> &[NodeInProof] { + &self.siblings + } + + /// Verifies an element whose key is `element_key` and value is `element_value` exists in the Sparse Merkle Tree + /// using the provided proof + pub fn verify_inclusion( + &self, + expected_root_hash: &TreeHash, + element_key: &LeafKey, + element_value_hash: &TreeHash, + ) -> Result<(), JmtProofVerifyError> { + self.verify(expected_root_hash, element_key, Some(element_value_hash)) + } + + /// Verifies the proof is a valid non-inclusion proof that shows this key doesn't exist in the tree. + pub fn verify_exclusion( + &self, + expected_root_hash: &TreeHash, + element_key: &LeafKey, + ) -> Result<(), JmtProofVerifyError> { + self.verify(expected_root_hash, element_key, None) + } + + /// If `element_value` is present, verifies an element whose key is `element_key` and value is + /// `element_value` exists in the Sparse Merkle Tree using the provided proof. Otherwise, + /// verifies the proof is a valid non-inclusion proof that shows this key doesn't exist in the + /// tree. + fn verify( + &self, + expected_root_hash: &TreeHash, + element_key: &LeafKey, + element_value: Option<&TreeHash>, + ) -> Result<(), JmtProofVerifyError> { + if self.siblings.len() > 256 { + return Err(JmtProofVerifyError::TooManySiblings { + num_siblings: self.siblings.len(), + }); + } + + match (element_value, &self.leaf) { + (Some(value_hash), Some(leaf)) => { + // This is an inclusion proof, so the key and value hash provided in the proof + // should match element_key and element_value_hash. `siblings` should prove the + // route from the leaf node to the root. + if element_key != leaf.key() { + return Err(JmtProofVerifyError::KeyMismatch { + actual_key: *leaf.key(), + expected_key: *element_key, + }); + } + if *value_hash != leaf.value_hash { + return Err(JmtProofVerifyError::ValueMismatch { + actual: leaf.value_hash, + expected: *value_hash, + }); + } + }, + (Some(_), None) => return Err(JmtProofVerifyError::ExpectedInclusionProof), + (None, Some(leaf)) => { + // This is a non-inclusion proof. The proof intends to show that if a leaf node + // representing `element_key` is inserted, it will break a currently existing leaf + // node represented by `proof_key` into a branch. `siblings` should prove the + // route from that leaf node to the root. + if element_key == leaf.key() { + return Err(JmtProofVerifyError::ExpectedNonInclusionProof); + } + if element_key.common_prefix_bits_len(leaf.key()) < self.siblings.len() { + return Err(JmtProofVerifyError::InvalidNonInclusionProof); + } + }, + (None, None) => { + // This is a non-inclusion proof. The proof intends to show that if a leaf node + // representing `element_key` is inserted, it will show up at a currently empty + // position. `sibling` should prove the route from this empty position to the root. + }, + } + + let current_hash = self + .leaf + .clone() + .map_or(SPARSE_MERKLE_PLACEHOLDER_HASH, |leaf| leaf.hash()); + let actual_root_hash = self + .siblings + .iter() + .zip(element_key.iter_bits().rev().skip(256 - self.siblings.len())) + .fold(current_hash, |hash, (sibling_node, bit)| { + if bit { + SparseMerkleInternalNode::new(sibling_node.hash(), hash).hash() + } else { + SparseMerkleInternalNode::new(hash, sibling_node.hash()).hash() + } + }); + + if actual_root_hash != *expected_root_hash { + return Err(JmtProofVerifyError::RootHashMismatch { + actual_root_hash, + expected_root_hash: *expected_root_hash, + }); + } + + Ok(()) + } +} + +impl From for SparseMerkleProof { + fn from(proof_ext: SparseMerkleProofExt) -> Self { + Self::new( + proof_ext.leaf, + proof_ext.siblings.into_iter().map(|node| node.hash()).collect(), + ) + } +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/proof/definition.rs#L135 +impl SparseMerkleProof { + /// Constructs a new `SparseMerkleProof` using leaf and a list of siblings. + pub fn new(leaf: Option, siblings: Vec) -> Self { + SparseMerkleProof { leaf, siblings } + } + + /// Returns the leaf node in this proof. + pub fn leaf(&self) -> Option { + self.leaf.clone() + } + + /// Returns the list of siblings in this proof. + pub fn siblings(&self) -> &[TreeHash] { + &self.siblings + } +} + +/// A proof that can be used to authenticate an element in a Sparse Merkle Tree given trusted root +/// hash. For example, `TransactionInfoToAccountProof` can be constructed on top of this structure. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SparseMerkleProof { + /// This proof can be used to authenticate whether a given leaf exists in the tree or not. + /// - If this is `Some(leaf_node)` + /// - If `leaf_node.key` equals requested key, this is an inclusion proof and `leaf_node.value_hash` equals + /// the hash of the corresponding account blob. + /// - Otherwise this is a non-inclusion proof. `leaf_node.key` is the only key that exists in the subtree + /// and `leaf_node.value_hash` equals the hash of the corresponding account blob. + /// - If this is `None`, this is also a non-inclusion proof which indicates the subtree is empty. + leaf: Option, + + /// All siblings in this proof, including the default ones. Siblings are ordered from the bottom + /// level to the root level. + siblings: Vec, +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)] +pub enum NodeInProof { + Leaf(SparseMerkleLeafNode), + Other(TreeHash), +} + +impl From for NodeInProof { + fn from(hash: TreeHash) -> Self { + Self::Other(hash) + } +} + +impl From for NodeInProof { + fn from(leaf: SparseMerkleLeafNode) -> Self { + Self::Leaf(leaf) + } +} + +impl NodeInProof { + pub fn hash(&self) -> TreeHash { + match self { + Self::Leaf(leaf) => leaf.hash(), + Self::Other(hash) => *hash, + } + } +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/proof/definition.rs#L681 +/// Note: this is not a range proof in the sense that a range of nodes is verified! +/// Instead, it verifies the entire left part of the tree up to a known rightmost node. +/// See the description below. +/// +/// A proof that can be used to authenticate a range of consecutive leaves, from the leftmost leaf to +/// the rightmost known one, in a sparse Merkle tree. For example, given the following sparse Merkle tree: +/// +/// ```text +/// root +/// / \ +/// / \ +/// / \ +/// o o +/// / \ / \ +/// a o o h +/// / \ / \ +/// o d e X +/// / \ / \ +/// b c f g +/// ``` +/// +/// if the proof wants show that `[a, b, c, d, e]` exists in the tree, it would need the siblings +/// `X` and `h` on the right. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SparseMerkleRangeProof { + /// The vector of siblings on the right of the path from root to last leaf. The ones near the + /// bottom are at the beginning of the vector. In the above example, it's `[X, h]`. + right_siblings: Vec, +} + +impl SparseMerkleRangeProof { + /// Constructs a new `SparseMerkleRangeProof`. + pub fn new(right_siblings: Vec) -> Self { + Self { right_siblings } + } + + /// Returns the right siblings. + pub fn right_siblings(&self) -> &[TreeHash] { + &self.right_siblings + } +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/proof/mod.rs#L97 +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize, BorshSerialize, BorshDeserialize)] +pub struct SparseMerkleLeafNode { + key: LeafKey, + value_hash: TreeHash, +} + +impl SparseMerkleLeafNode { + pub fn new(key: LeafKey, value_hash: TreeHash) -> Self { + SparseMerkleLeafNode { key, value_hash } + } + + pub fn key(&self) -> &LeafKey { + &self.key + } + + pub fn value_hash(&self) -> &TreeHash { + &self.value_hash + } + + pub fn hash(&self) -> TreeHash { + jmt_node_hash2(&self.key.bytes, &self.value_hash) + } +} + +pub struct SparseMerkleInternalNode { + left_child: TreeHash, + right_child: TreeHash, +} + +impl SparseMerkleInternalNode { + pub fn new(left_child: TreeHash, right_child: TreeHash) -> Self { + Self { + left_child, + right_child, + } + } + + fn hash(&self) -> TreeHash { + jmt_node_hash2(&self.left_child, &self.right_child) + } +} + +// INITIAL-MODIFICATION: we propagate usage of our own `Hash` (instead of Aptos' `HashValue`) to avoid +// sourcing the entire https://github.com/aptos-labs/aptos-core/blob/1.0.4/crates/aptos-crypto/src/hash.rs +pub const SPARSE_MERKLE_PLACEHOLDER_HASH: TreeHash = TreeHash::zero(); + +// CSOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/crates/aptos-crypto/src/hash.rs#L422 +/// An iterator over `LeafKey` that generates one bit for each iteration. +pub struct LeafKeyBitIterator<'a> { + /// The reference to the bytes that represent the `LeafKey`. + leaf_key_bytes: &'a [u8], + pos: Range, + // invariant pos.end == leaf_key_bytes.len() * 8; +} + +impl<'a> DoubleEndedIterator for LeafKeyBitIterator<'a> { + fn next_back(&mut self) -> Option { + self.pos.next_back().map(|x| self.get_bit(x)) + } +} + +impl<'a> ExactSizeIterator for LeafKeyBitIterator<'a> {} + +impl<'a> LeafKeyBitIterator<'a> { + /// Constructs a new `LeafKeyBitIterator` using given `leaf_key_bytes`. + fn new(leaf_key: LeafKeyRef<'a>) -> Self { + LeafKeyBitIterator { + leaf_key_bytes: leaf_key.bytes, + pos: (0..leaf_key.bytes.len() * 8), + } + } + + /// Returns the `index`-th bit in the bytes. + fn get_bit(&self, index: usize) -> bool { + let pos = index / 8; + let bit = 7 - index % 8; + (self.leaf_key_bytes[pos] >> bit) & 1 != 0 + } +} + +impl<'a> Iterator for LeafKeyBitIterator<'a> { + type Item = bool; + + fn next(&mut self) -> Option { + self.pos.next().map(|x| self.get_bit(x)) + } + + fn size_hint(&self) -> (usize, Option) { + self.pos.size_hint() + } +} + +// INITIAL-MODIFICATION: since we use our own `LeafKey` here, we need it to implement these for it +pub trait IteratedLeafKey { + fn iter_bits(&self) -> LeafKeyBitIterator<'_>; + + fn get_nibble(&self, index: usize) -> Nibble; +} + +impl IteratedLeafKey for LeafKey { + fn iter_bits(&self) -> LeafKeyBitIterator<'_> { + LeafKeyBitIterator::new(self.as_ref()) + } + + fn get_nibble(&self, index: usize) -> Nibble { + Nibble::from(if index % 2 == 0 { + self.bytes[index / 2] >> 4 + } else { + self.bytes[index / 2] & 0x0F + }) + } +} + +impl IteratedLeafKey for LeafKeyRef<'_> { + fn iter_bits(&self) -> LeafKeyBitIterator<'_> { + LeafKeyBitIterator::new(*self) + } + + fn get_nibble(&self, index: usize) -> Nibble { + Nibble::from(if index % 2 == 0 { + self.bytes[index / 2] >> 4 + } else { + self.bytes[index / 2] & 0x0F + }) + } +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/transaction/mod.rs#L57 +pub type Version = u64; + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/nibble/mod.rs#L20 +#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)] +#[serde(transparent)] +pub struct Nibble(u8); + +impl From for Nibble { + fn from(nibble: u8) -> Self { + assert!(nibble < 16, "Nibble out of range: {}", nibble); + Self(nibble) + } +} + +impl From for u8 { + fn from(nibble: Nibble) -> Self { + nibble.0 + } +} + +impl fmt::LowerHex for Nibble { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:x}", self.0) + } +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/types/src/nibble/nibble_path/mod.rs#L22 +/// NibblePath defines a path in Merkle tree in the unit of nibble (4 bits). +#[derive(Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)] +pub struct NibblePath { + /// Indicates the total number of nibbles in bytes. Either `bytes.len() * 2 - 1` or + /// `bytes.len() * 2`. + // Guarantees intended ordering based on the top-to-bottom declaration order of the struct's + // members. + num_nibbles: usize, + /// The underlying bytes that stores the path, 2 nibbles per byte. If the number of nibbles is + /// odd, the second half of the last byte must be 0. + bytes: Vec, +} + +/// Supports debug format by concatenating nibbles literally. For example, [0x12, 0xa0] with 3 +/// nibbles will be printed as "12a". +impl fmt::Debug for NibblePath { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.nibbles().try_for_each(|x| write!(f, "{:x}", x)) + } +} + +// INITIAL-MODIFICATION: just to show it in errors +impl fmt::Display for NibblePath { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let hex_chars = self + .bytes + .iter() + .flat_map(|b| [b >> 4, b & 15]) + .map(|b| char::from_digit(u32::from(b), 16).unwrap()) + .take(self.num_nibbles); + + for ch in hex_chars { + write!(f, "{}", ch)?; + } + Ok(()) + } +} + +/// Convert a vector of bytes into `NibblePath` using the lower 4 bits of each byte as nibble. +impl FromIterator for NibblePath { + fn from_iter>(iter: I) -> Self { + let mut nibble_path = NibblePath::new_even(vec![]); + for nibble in iter { + nibble_path.push(nibble); + } + nibble_path + } +} + +impl NibblePath { + /// Creates a new `NibblePath` from a vector of bytes assuming each byte has 2 nibbles. + pub fn new_even(bytes: Vec) -> Self { + let num_nibbles = bytes.len() * 2; + NibblePath { num_nibbles, bytes } + } + + /// Similar to `new()` but asserts that the bytes have one less nibble. + pub fn new_odd(bytes: Vec) -> Self { + assert_eq!( + bytes.last().expect("Should have odd number of nibbles.") & 0x0F, + 0, + "Last nibble must be 0." + ); + let num_nibbles = bytes.len() * 2 - 1; + NibblePath { num_nibbles, bytes } + } + + /// Adds a nibble to the end of the nibble path. + pub fn push(&mut self, nibble: Nibble) { + if self.num_nibbles % 2 == 0 { + self.bytes.push(u8::from(nibble) << 4); + } else { + self.bytes[self.num_nibbles / 2] |= u8::from(nibble); + } + self.num_nibbles += 1; + } + + /// Pops a nibble from the end of the nibble path. + pub fn pop(&mut self) -> Option { + let poped_nibble = if self.num_nibbles % 2 == 0 { + self.bytes.last_mut().map(|last_byte| { + let nibble = *last_byte & 0x0F; + *last_byte &= 0xF0; + Nibble::from(nibble) + }) + } else { + self.bytes.pop().map(|byte| Nibble::from(byte >> 4)) + }; + if poped_nibble.is_some() { + self.num_nibbles -= 1; + } + poped_nibble + } + + /// Returns the last nibble. + pub fn last(&self) -> Option { + let last_byte_option = self.bytes.last(); + if self.num_nibbles % 2 == 0 { + last_byte_option.map(|last_byte| Nibble::from(*last_byte & 0x0F)) + } else { + let last_byte = last_byte_option.expect("Last byte must exist if num_nibbles is odd."); + Some(Nibble::from(*last_byte >> 4)) + } + } + + /// Get the i-th bit. + fn get_bit(&self, i: usize) -> bool { + assert!(i < self.num_nibbles * 4); + let pos = i / 8; + let bit = 7 - i % 8; + ((self.bytes[pos] >> bit) & 1) != 0 + } + + /// Get the i-th nibble. + pub fn get_nibble(&self, i: usize) -> Nibble { + assert!(i < self.num_nibbles); + Nibble::from((self.bytes[i / 2] >> (if i % 2 == 1 { 0 } else { 4 })) & 0xF) + } + + /// Get a bit iterator iterates over the whole nibble path. + pub fn bits(&self) -> NibbleBitIterator { + NibbleBitIterator { + nibble_path: self, + pos: (0..self.num_nibbles * 4), + } + } + + /// Get a nibble iterator iterates over the whole nibble path. + pub fn nibbles(&self) -> NibbleIterator { + NibbleIterator::new(self, 0, self.num_nibbles) + } + + /// Get the total number of nibbles stored. + pub fn num_nibbles(&self) -> usize { + self.num_nibbles + } + + /// Returns `true` if the nibbles contains no elements. + pub fn is_empty(&self) -> bool { + self.num_nibbles() == 0 + } + + /// Get the underlying bytes storing nibbles. + pub fn bytes(&self) -> &[u8] { + &self.bytes + } + + pub fn into_bytes(self) -> Vec { + self.bytes + } + + pub fn truncate(&mut self, len: usize) { + assert!(len <= self.num_nibbles); + self.num_nibbles = len; + self.bytes.truncate((len + 1) / 2); + if len % 2 != 0 { + *self.bytes.last_mut().expect("must exist.") &= 0xF0; + } + } +} + +pub trait Peekable: Iterator { + /// Returns the `next()` value without advancing the iterator. + fn peek(&self) -> Option; +} + +/// BitIterator iterates a nibble path by bit. +pub struct NibbleBitIterator<'a> { + nibble_path: &'a NibblePath, + pos: Range, +} + +impl<'a> Peekable for NibbleBitIterator<'a> { + /// Returns the `next()` value without advancing the iterator. + fn peek(&self) -> Option { + if self.pos.start < self.pos.end { + Some(self.nibble_path.get_bit(self.pos.start)) + } else { + None + } + } +} + +/// BitIterator spits out a boolean each time. True/false denotes 1/0. +impl<'a> Iterator for NibbleBitIterator<'a> { + type Item = bool; + + fn next(&mut self) -> Option { + self.pos.next().map(|i| self.nibble_path.get_bit(i)) + } +} + +/// Support iterating bits in reversed order. +impl<'a> DoubleEndedIterator for NibbleBitIterator<'a> { + fn next_back(&mut self) -> Option { + self.pos.next_back().map(|i| self.nibble_path.get_bit(i)) + } +} + +/// NibbleIterator iterates a nibble path by nibble. +#[derive(Debug)] +pub struct NibbleIterator<'a> { + /// The underlying nibble path that stores the nibbles + nibble_path: &'a NibblePath, + + /// The current index, `pos.start`, will bump by 1 after calling `next()` until `pos.start == + /// pos.end`. + pos: Range, + + /// The start index of the iterator. At the beginning, `pos.start == start`. [start, pos.end) + /// defines the range of `nibble_path` this iterator iterates over. `nibble_path` refers to + /// the entire underlying buffer but the range may only be partial. + start: usize, + // invariant self.start <= self.pos.start; + // invariant self.pos.start <= self.pos.end; +} + +/// NibbleIterator spits out a byte each time. Each byte must be in range [0, 16). +impl<'a> Iterator for NibbleIterator<'a> { + type Item = Nibble; + + fn next(&mut self) -> Option { + self.pos.next().map(|i| self.nibble_path.get_nibble(i)) + } +} + +impl<'a> Peekable for NibbleIterator<'a> { + /// Returns the `next()` value without advancing the iterator. + fn peek(&self) -> Option { + if self.pos.start < self.pos.end { + Some(self.nibble_path.get_nibble(self.pos.start)) + } else { + None + } + } +} + +impl<'a> NibbleIterator<'a> { + fn new(nibble_path: &'a NibblePath, start: usize, end: usize) -> Self { + assert!(start <= end); + Self { + nibble_path, + pos: (start..end), + start, + } + } + + /// Returns a nibble iterator that iterates all visited nibbles. + pub fn visited_nibbles(&self) -> NibbleIterator<'a> { + Self::new(self.nibble_path, self.start, self.pos.start) + } + + /// Returns a nibble iterator that iterates all remaining nibbles. + pub fn remaining_nibbles(&self) -> NibbleIterator<'a> { + Self::new(self.nibble_path, self.pos.start, self.pos.end) + } + + /// Turn it into a `BitIterator`. + pub fn bits(&self) -> NibbleBitIterator<'a> { + NibbleBitIterator { + nibble_path: self.nibble_path, + pos: (self.pos.start * 4..self.pos.end * 4), + } + } + + /// Cut and return the range of the underlying `nibble_path` that this iterator is iterating + /// over as a new `NibblePath` + pub fn get_nibble_path(&self) -> NibblePath { + self.visited_nibbles().chain(self.remaining_nibbles()).collect() + } + + /// Get the number of nibbles that this iterator covers. + pub fn num_nibbles(&self) -> usize { + assert!(self.start <= self.pos.end); // invariant + self.pos.end - self.start + } + + /// Return `true` if the iteration is over. + pub fn is_finished(&self) -> bool { + self.peek().is_none() + } +} + +// INITIAL-MODIFICATION: We will use this type (instead of `Hash`) to allow for arbitrary key length +/// A leaf key (i.e. a complete nibble path). +#[derive( + Clone, Debug, Copy, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, BorshSerialize, BorshDeserialize, +)] +pub struct LeafKey { + /// The underlying bytes. + /// All leaf keys of the same tree must be of the same length - otherwise the tree's behavior + /// becomes unspecified. + /// All leaf keys must be evenly distributed across their space - otherwise the tree's + /// performance degrades. + /// TARI: always a hash, so replaced heap-allocated Vec with a Hash + pub bytes: TreeHash, +} + +impl LeafKey { + pub fn new(bytes: TreeHash) -> Self { + Self { bytes } + } + + pub fn as_ref(&self) -> LeafKeyRef<'_> { + LeafKeyRef::new(self.bytes.as_slice()) + } + + pub fn iter_bits(&self) -> BitIterator<'_> { + BitIterator::new(self.bytes.as_slice()) + } + + pub fn common_prefix_bits_len(&self, other: &LeafKey) -> usize { + self.iter_bits() + .zip(other.iter_bits()) + .take_while(|(x, y)| x == y) + .count() + } +} + +impl Display for LeafKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.bytes.fmt(f) + } +} + +// INITIAL-MODIFICATION: We will use this type (instead of `Hash`) to allow for arbitrary key length +/// A leaf key (i.e. a complete nibble path). +#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)] +pub struct LeafKeyRef<'a> { + /// The underlying bytes. + /// All leaf keys of the same tree must be of the same length - otherwise the tree's behavior + /// becomes unspecified. + /// All leaf keys must be evenly distributed across their space - otherwise the tree's + /// performance degrades. + pub bytes: &'a [u8], +} + +impl<'a> LeafKeyRef<'a> { + pub fn new(bytes: &'a [u8]) -> Self { + Self { bytes } + } +} + +impl PartialEq for LeafKeyRef<'_> { + fn eq(&self, other: &LeafKey) -> bool { + self.bytes == other.bytes.as_slice() + } +} + +// SOURCE: https://github.com/aptos-labs/aptos-core/blob/1.0.4/storage/jellyfish-merkle/src/node_type/mod.rs#L48 +/// The unique key of each node. +#[derive(Clone, Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)] +pub struct NodeKey { + /// The version at which the node is created. + version: Version, + /// The nibble path this node represents in the tree. + nibble_path: NibblePath, +} + +impl NodeKey { + /// Creates a new `NodeKey`. + pub fn new(version: Version, nibble_path: NibblePath) -> Self { + Self { version, nibble_path } + } + + /// A shortcut to generate a node key consisting of a version and an empty nibble path. + pub fn new_empty_path(version: Version) -> Self { + Self::new(version, NibblePath::new_even(vec![])) + } + + /// Gets the version. + pub fn version(&self) -> Version { + self.version + } + + /// Gets the nibble path. + pub fn nibble_path(&self) -> &NibblePath { + &self.nibble_path + } + + /// Generates a child node key based on this node key. + pub fn gen_child_node_key(&self, version: Version, n: Nibble) -> Self { + let mut node_nibble_path = self.nibble_path().clone(); + node_nibble_path.push(n); + Self::new(version, node_nibble_path) + } + + /// Generates parent node key at the same version based on this node key. + pub fn gen_parent_node_key(&self) -> Self { + let mut node_nibble_path = self.nibble_path().clone(); + assert!(node_nibble_path.pop().is_some(), "Current node key is root.",); + Self::new(self.version, node_nibble_path) + } +} + +// INITIAL-MODIFICATION: just to show it in errors +impl fmt::Display for NodeKey { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "v{}:{}", self.version, self.nibble_path) + } +} + +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub enum NodeType { + Leaf, + Null, + /// A internal node that haven't been finished the leaf count migration, i.e. None or not all + /// of the children leaf counts are known. + Internal { + leaf_count: usize, + }, +} + +/// Each child of [`InternalNode`] encapsulates a nibble forking at this node. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct Child { + /// The hash value of this child node. + pub hash: TreeHash, + /// `version`, the `nibble_path` of the [`NodeKey`] of this [`InternalNode`] the child belongs + /// to and the child's index constitute the [`NodeKey`] to uniquely identify this child node + /// from the storage. Used by `[`NodeKey::gen_child_node_key`]. + pub version: Version, + /// Indicates if the child is a leaf, or if it's an internal node, the total number of leaves + /// under it (though it can be unknown during migration). + pub node_type: NodeType, +} + +impl Child { + pub fn new(hash: TreeHash, version: Version, node_type: NodeType) -> Self { + Self { + hash, + version, + node_type, + } + } + + pub fn is_leaf(&self) -> bool { + matches!(self.node_type, NodeType::Leaf) + } + + pub fn leaf_count(&self) -> usize { + match self.node_type { + NodeType::Leaf => 1, + NodeType::Internal { leaf_count } => leaf_count, + NodeType::Null => unreachable!("Child cannot be Null"), + } + } +} + +/// [`Children`] is just a collection of children belonging to a [`InternalNode`], indexed from 0 to +/// 15, inclusive. +pub(crate) type Children = IndexMap; + +/// Represents a 4-level subtree with 16 children at the bottom level. Theoretically, this reduces +/// IOPS to query a tree by 4x since we compress 4 levels in a standard Merkle tree into 1 node. +/// Though we choose the same internal node structure as that of Patricia Merkle tree, the root hash +/// computation logic is similar to a 4-level sparse Merkle tree except for some customizations. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct InternalNode { + /// Up to 16 children. + children: Children, + /// Total number of leaves under this internal node + leaf_count: usize, +} + +impl InternalNode { + /// Creates a new Internal node. + pub fn new(mut children: Children) -> Self { + children.sort_keys(); + let leaf_count = children.values().map(Child::leaf_count).sum(); + Self { children, leaf_count } + } + + pub fn leaf_count(&self) -> usize { + self.leaf_count + } + + pub fn node_type(&self) -> NodeType { + NodeType::Internal { + leaf_count: self.leaf_count, + } + } + + pub fn hash(&self) -> TreeHash { + self.merkle_hash( + 0, // start index + 16, // the number of leaves in the subtree of which we want the hash of root + self.generate_bitmaps(), + ) + } + + pub fn children_sorted(&self) -> impl Iterator { + // let mut tmp = self.children.iter().collect::>(); + // tmp.sort_by_key(|(nibble, _)| **nibble); + // tmp.into_iter() + self.children.iter() + } + + pub fn into_children(self) -> Children { + self.children + } + + /// Gets the `n`-th child. + pub fn child(&self, n: Nibble) -> Option<&Child> { + self.children.get(&n) + } + + /// Generates `existence_bitmap` and `leaf_bitmap` as a pair of `u16`s: child at index `i` + /// exists if `existence_bitmap[i]` is set; child at index `i` is leaf node if + /// `leaf_bitmap[i]` is set. + pub fn generate_bitmaps(&self) -> (u16, u16) { + let mut existence_bitmap = 0; + let mut leaf_bitmap = 0; + for (nibble, child) in &self.children { + let i = u8::from(*nibble); + existence_bitmap |= 1u16 << i; + if child.is_leaf() { + leaf_bitmap |= 1u16 << i; + } + } + // `leaf_bitmap` must be a subset of `existence_bitmap`. + assert_eq!(existence_bitmap | leaf_bitmap, existence_bitmap); + (existence_bitmap, leaf_bitmap) + } + + /// Given a range [start, start + width), returns the sub-bitmap of that range. + fn range_bitmaps(start: u8, width: u8, bitmaps: (u16, u16)) -> (u16, u16) { + assert!(start < 16 && width.count_ones() == 1 && start % width == 0); + assert!(width <= 16 && (start + width) <= 16); + // A range with `start == 8` and `width == 4` will generate a mask 0b0000111100000000. + // use as converting to smaller integer types when 'width == 16' + #[allow(clippy::cast_possible_truncation)] + let mask = (((1u32 << width) - 1) << start) as u16; + (bitmaps.0 & mask, bitmaps.1 & mask) + } + + fn merkle_hash(&self, start: u8, width: u8, (existence_bitmap, leaf_bitmap): (u16, u16)) -> TreeHash { + // Given a bit [start, 1 << nibble_height], return the value of that range. + let (range_existence_bitmap, range_leaf_bitmap) = + Self::range_bitmaps(start, width, (existence_bitmap, leaf_bitmap)); + if range_existence_bitmap == 0 { + // No child under this subtree + SPARSE_MERKLE_PLACEHOLDER_HASH + } else if width == 1 || (range_existence_bitmap.count_ones() == 1 && range_leaf_bitmap != 0) { + // Only 1 leaf child under this subtree or reach the lowest level + #[allow(clippy::cast_possible_truncation)] + let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8); + self.child(only_child_index) + .expect("Corrupted internal node: existence_bitmap inconsistent") + .hash + } else { + let left_child = self.merkle_hash(start, width / 2, (range_existence_bitmap, range_leaf_bitmap)); + let right_child = self.merkle_hash( + start + width / 2, + width / 2, + (range_existence_bitmap, range_leaf_bitmap), + ); + SparseMerkleInternalNode::new(left_child, right_child).hash() + } + } + + fn gen_node_in_proof>( + &self, + start: u8, + width: u8, + (existence_bitmap, leaf_bitmap): (u16, u16), + (tree_reader, node_key): (&R, &NodeKey), + ) -> Result { + // Given a bit [start, 1 << nibble_height], return the value of that range. + let (range_existence_bitmap, range_leaf_bitmap) = + Self::range_bitmaps(start, width, (existence_bitmap, leaf_bitmap)); + Ok(if range_existence_bitmap == 0 { + // No child under this subtree + NodeInProof::Other(SPARSE_MERKLE_PLACEHOLDER_HASH) + } else if width == 1 || (range_existence_bitmap.count_ones() == 1 && range_leaf_bitmap != 0) { + // Only 1 leaf child under this subtree or reach the lowest level + #[allow(clippy::cast_possible_truncation)] + let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8); + let only_child = self + .child(only_child_index) + .expect("Corrupted internal node: existence_bitmap inconsistent"); + if matches!(only_child.node_type, NodeType::Leaf) { + let only_child_node_key = node_key.gen_child_node_key(only_child.version, only_child_index); + match tree_reader.get_node(&only_child_node_key)? { + Node::Internal(_) => { + unreachable!("Corrupted internal node: in-memory leaf child is internal node on disk") + }, + Node::Leaf(leaf_node) => NodeInProof::Leaf(SparseMerkleLeafNode::from(leaf_node)), + Node::Null => unreachable!("Child cannot be Null"), + } + } else { + NodeInProof::Other(only_child.hash) + } + } else { + let left_child = self.merkle_hash(start, width / 2, (range_existence_bitmap, range_leaf_bitmap)); + let right_child = self.merkle_hash( + start + width / 2, + width / 2, + (range_existence_bitmap, range_leaf_bitmap), + ); + NodeInProof::Other(SparseMerkleInternalNode::new(left_child, right_child).hash()) + }) + } + + /// Gets the child and its corresponding siblings that are necessary to generate the proof for + /// the `n`-th child. If it is an existence proof, the returned child must be the `n`-th + /// child; otherwise, the returned child may be another child. See inline explanation for + /// details. When calling this function with n = 11 (node `b` in the following graph), the + /// range at each level is illustrated as a pair of square brackets: + /// + /// ```text + /// 4 [f e d c b a 9 8 7 6 5 4 3 2 1 0] -> root level + /// --------------------------------------------------------------- + /// 3 [f e d c b a 9 8] [7 6 5 4 3 2 1 0] width = 8 + /// chs <--┘ shs <--┘ + /// 2 [f e d c] [b a 9 8] [7 6 5 4] [3 2 1 0] width = 4 + /// shs <--┘ └--> chs + /// 1 [f e] [d c] [b a] [9 8] [7 6] [5 4] [3 2] [1 0] width = 2 + /// chs <--┘ └--> shs + /// 0 [f] [e] [d] [c] [b] [a] [9] [8] [7] [6] [5] [4] [3] [2] [1] [0] width = 1 + /// ^ chs <--┘ └--> shs + /// | MSB|<---------------------- uint 16 ---------------------------->|LSB + /// height chs: `child_half_start` shs: `sibling_half_start` + /// ``` + pub fn get_child_with_siblings>( + &self, + node_key: &NodeKey, + n: Nibble, + reader: Option<&R>, + ) -> Result<(Option, Vec), JmtStorageError> { + let mut siblings = vec![]; + let (existence_bitmap, leaf_bitmap) = self.generate_bitmaps(); + + // Nibble height from 3 to 0. + for h in (0..4).rev() { + // Get the number of children of the internal node that each subtree at this height + // covers. + let width = 1 << h; + let (child_half_start, sibling_half_start) = get_child_and_sibling_half_start(n, h); + // Compute the root hash of the subtree rooted at the sibling of `r`. + if let Some(reader) = reader { + siblings.push(self.gen_node_in_proof( + sibling_half_start, + width, + (existence_bitmap, leaf_bitmap), + (reader, node_key), + )?); + } else { + siblings.push( + self.merkle_hash(sibling_half_start, width, (existence_bitmap, leaf_bitmap)) + .into(), + ); + } + + let (range_existence_bitmap, range_leaf_bitmap) = + Self::range_bitmaps(child_half_start, width, (existence_bitmap, leaf_bitmap)); + + if range_existence_bitmap == 0 { + // No child in this range. + return Ok((None, siblings)); + } + + if width == 1 || (range_existence_bitmap.count_ones() == 1 && range_leaf_bitmap != 0) { + // Return the only 1 leaf child under this subtree or reach the lowest level + // Even this leaf child is not the n-th child, it should be returned instead of + // `None` because it's existence indirectly proves the n-th child doesn't exist. + // Please read proof format for details. + #[allow(clippy::cast_possible_truncation)] + let only_child_index = Nibble::from(range_existence_bitmap.trailing_zeros() as u8); + return Ok(( + { + let only_child_version = self + .child(only_child_index) + // Should be guaranteed by the self invariants, but these are not easy to express at the moment + .expect("Corrupted internal node: child_bitmap inconsistent") + .version; + Some(node_key.gen_child_node_key(only_child_version, only_child_index)) + }, + siblings, + )); + } + } + unreachable!("Impossible to get here without returning even at the lowest level.") + } +} + +/// Given a nibble, computes the start position of its `child_half_start` and `sibling_half_start` +/// at `height` level. +pub(crate) fn get_child_and_sibling_half_start(n: Nibble, height: u8) -> (u8, u8) { + // Get the index of the first child belonging to the same subtree whose root, let's say `r` is + // at `height` that the n-th child belongs to. + // Note: `child_half_start` will be always equal to `n` at height 0. + let child_half_start = (0xFF << height) & u8::from(n); + + // Get the index of the first child belonging to the subtree whose root is the sibling of `r` + // at `height`. + let sibling_half_start = child_half_start ^ (1 << height); + + (child_half_start, sibling_half_start) +} + +/// Leaf node, capturing the value hash and carrying an arbitrary payload. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub struct LeafNode

{ + // The key of this leaf node (i.e. its full nibble path). + leaf_key: LeafKey, + // The hash of an externally-stored value. + // Note: do not confuse that value with the `payload`. + value_hash: TreeHash, + // The client payload. + // This is not the "value" whose changes are tracked by the tree (in fact, these values are + // supposed to be stored externally, and the tree only cares about their hashes - see + // `value_hash`). + // Rather, the payload is an arbitrary piece of data that the client wishes to store within the + // tree, in order to facilitate some related processing: + // - Many clients do not need it and will simply use a no-cost `()`. + // - A use-case designed by the original authors was to store a non-hashed element key as a payload (while the + // `leaf_key` contains that key's hash, to ensure the nibble paths are distributed over their space, for + // performance). + // - Our current use-case (specific to a "two layers" tree) is to store the nested tree's root metadata. + payload: P, + // The version at which this leaf was created. + version: Version, +} + +impl

LeafNode

{ + /// Creates a new leaf node. + pub fn new(leaf_key: LeafKey, value_hash: TreeHash, payload: P, version: Version) -> Self { + Self { + leaf_key, + value_hash, + payload, + version, + } + } + + /// Gets the key. + pub fn leaf_key(&self) -> &LeafKey { + &self.leaf_key + } + + /// Gets the associated value hash. + pub fn value_hash(&self) -> TreeHash { + self.value_hash + } + + /// Gets the payload. + pub fn payload(&self) -> &P { + &self.payload + } + + /// Gets the version. + pub fn version(&self) -> Version { + self.version + } + + /// Gets the leaf's hash (not to be confused with a `value_hash()`). + /// This hash incorporates the node's key and the value's hash, in order to capture certain + /// changes within a sparse merkle tree (consider 2 trees, both containing a single element with + /// the same value, but stored under different keys - we want their root hashes to differ). + pub fn leaf_hash(&self) -> TreeHash { + jmt_node_hash2(&self.leaf_key.bytes, &self.value_hash) + } +} + +impl From> for SparseMerkleLeafNode { + fn from(leaf_node: LeafNode) -> Self { + Self::new(leaf_node.leaf_key, leaf_node.value_hash) + } +} + +/// The concrete node type of [`JellyfishMerkleTree`]. +#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] +pub enum Node

{ + /// A wrapper of [`InternalNode`]. + Internal(InternalNode), + /// A wrapper of [`LeafNode`]. + Leaf(LeafNode

), + /// Represents empty tree only + Null, +} + +impl

From for Node

{ + fn from(node: InternalNode) -> Self { + Node::Internal(node) + } +} + +impl From> for Node

{ + fn from(node: LeafNode

) -> Self { + Node::Leaf(node) + } +} + +impl

Node

{ + // /// Creates the [`Internal`](Node::Internal) variant. + // #[cfg(any(test, feature = "fuzzing"))] + // pub fn new_internal(children: Children) -> Self { + // Node::Internal(InternalNode::new(children)) + // } + + /// Creates the [`Leaf`](Node::Leaf) variant. + pub fn new_leaf(leaf_key: LeafKey, value_hash: TreeHash, payload: P, version: Version) -> Self { + Node::Leaf(LeafNode::new(leaf_key, value_hash, payload, version)) + } + + /// Returns `true` if the node is a leaf node. + pub fn is_leaf(&self) -> bool { + matches!(self, Node::Leaf(_)) + } + + pub fn leaf(&self) -> Option<&LeafNode

> { + match self { + Node::Leaf(leaf) => Some(leaf), + _ => None, + } + } + + /// Returns `NodeType` + pub fn node_type(&self) -> NodeType { + match self { + // The returning value will be used to construct a `Child` of a internal node, while an + // internal node will never have a child of Node::Null. + Self::Leaf(_) => NodeType::Leaf, + Self::Internal(n) => n.node_type(), + Self::Null => NodeType::Null, + } + } + + /// Returns leaf count if known + pub fn leaf_count(&self) -> usize { + match self { + Node::Leaf(_) => 1, + Node::Internal(internal_node) => internal_node.leaf_count, + Node::Null => 0, + } + } + + /// Computes the hash of nodes. + pub fn hash(&self) -> TreeHash { + match self { + Node::Internal(internal_node) => internal_node.hash(), + Node::Leaf(leaf_node) => leaf_node.leaf_hash(), + Node::Null => SPARSE_MERKLE_PLACEHOLDER_HASH, + } + } +} + +// INITIAL-MODIFICATION: we propagate usage of our own error enum (instead of `std::io::ErrorKind` +// used by Aptos) to allow for no-std build. +/// Error originating from underlying storage failure / inconsistency. +#[derive(Debug, thiserror::Error)] +pub enum JmtStorageError { + #[error("A node {0} expected to exist (according to JMT logic) was not found in the storage")] + NotFound(NodeKey), + + #[error("Nodes read from the storage are violating some JMT property (e.g. form a cycle).")] + InconsistentState, + + #[error("Unexpected error: {0}")] + UnexpectedError(String), + + #[error("Attempted to insert node {0} that already exists")] + Conflict(NodeKey), +}