From aa8da37bc67310a517927a9774344ae83fcd94fd Mon Sep 17 00:00:00 2001 From: LesterEvSe Date: Wed, 14 Jan 2026 17:57:59 +0200 Subject: [PATCH 1/2] feat: add smt storage --- crates/contracts/Cargo.toml | 4 +- crates/contracts/src/lib.rs | 3 +- .../src/smt_storage/build_witness.rs | 66 ++++++ crates/contracts/src/smt_storage/mod.rs | 220 ++++++++++++++++++ .../smt_storage/source_simf/smt_storage.simf | 65 ++++++ 5 files changed, 356 insertions(+), 2 deletions(-) create mode 100644 crates/contracts/src/smt_storage/build_witness.rs create mode 100644 crates/contracts/src/smt_storage/mod.rs create mode 100644 crates/contracts/src/smt_storage/source_simf/smt_storage.simf diff --git a/crates/contracts/Cargo.toml b/crates/contracts/Cargo.toml index 0d4aeb8..91897a4 100644 --- a/crates/contracts/Cargo.toml +++ b/crates/contracts/Cargo.toml @@ -10,7 +10,7 @@ keywords = ["simplicity", "liquid", "bitcoin", "elements", "contracts"] categories = ["cryptography::cryptocurrencies"] [features] -default = ["sdk-basic", "finance-options", "finance-dcd", "finance-option-offer", "simple-storage", "bytes32-tr-storage", "array-tr-storage"] +default = ["sdk-basic", "finance-options", "finance-dcd", "finance-option-offer", "simple-storage", "bytes32-tr-storage", "array-tr-storage", "smt-storage"] sdk-basic = [] finance-options = [] finance-dcd = [] @@ -18,6 +18,8 @@ finance-option-offer = [] simple-storage = [] bytes32-tr-storage = [] array-tr-storage = [] +smt-storage = [] +swap-with-change = [] [lints] workspace = true diff --git a/crates/contracts/src/lib.rs b/crates/contracts/src/lib.rs index 2f1ebf1..75535c6 100644 --- a/crates/contracts/src/lib.rs +++ b/crates/contracts/src/lib.rs @@ -18,7 +18,8 @@ pub mod bytes32_tr_storage; pub mod finance; #[cfg(feature = "simple-storage")] pub mod simple_storage; - +#[cfg(feature = "smt-storage")] +pub mod smt_storage; #[cfg(feature = "finance-dcd")] pub use finance::dcd; #[cfg(feature = "finance-option-offer")] diff --git a/crates/contracts/src/smt_storage/build_witness.rs b/crates/contracts/src/smt_storage/build_witness.rs new file mode 100644 index 0000000..9f267c1 --- /dev/null +++ b/crates/contracts/src/smt_storage/build_witness.rs @@ -0,0 +1,66 @@ +use std::collections::HashMap; + +use simplicityhl::num::U256; +use simplicityhl::types::{ResolvedType, TypeConstructible, UIntType}; +use simplicityhl::value::{UIntValue, ValueConstructible}; +use simplicityhl::{WitnessValues, str::WitnessName}; + +#[allow(non_camel_case_types)] +pub type u256 = [u8; 32]; +pub const DEPTH: usize = 7; + +#[derive(Debug, Clone, bincode::Encode, bincode::Decode, PartialEq, Eq)] +pub struct SMTWitness { + leaf: u256, + merkle_data: [(u256, bool); DEPTH], +} + +impl SMTWitness { + #[must_use] + pub fn new(leaf: &u256, merkle_data: &[(u256, bool); DEPTH]) -> Self { + Self { + leaf: *leaf, + merkle_data: *merkle_data, + } + } +} + +impl Default for SMTWitness { + fn default() -> Self { + Self { + leaf: [0u8; 32], + merkle_data: [([0u8; 32], false); DEPTH], + } + } +} + +#[must_use] +pub fn build_smt_storage_witness(witness: &SMTWitness) -> WitnessValues { + let values: Vec = witness + .merkle_data + .iter() + .map(|(value, is_right)| { + let hash_val = + simplicityhl::Value::from(UIntValue::U256(U256::from_byte_array(*value))); + let direction_val = simplicityhl::Value::from(*is_right); + + simplicityhl::Value::product(hash_val, direction_val) + }) + .collect(); + + let element_type = simplicityhl::types::TypeConstructible::product( + UIntType::U256.into(), + ResolvedType::boolean(), + ); + + simplicityhl::WitnessValues::from(HashMap::from([ + ( + WitnessName::from_str_unchecked("LEAF"), + simplicityhl::Value::from(UIntValue::U256(U256::from_byte_array(witness.leaf))), + ), + ( + WitnessName::from_str_unchecked("MERKLE_DATA"), + simplicityhl::Value::array(values, element_type), + ), + ])) +} diff --git a/crates/contracts/src/smt_storage/mod.rs b/crates/contracts/src/smt_storage/mod.rs new file mode 100644 index 0000000..db72f14 --- /dev/null +++ b/crates/contracts/src/smt_storage/mod.rs @@ -0,0 +1,220 @@ +use std::sync::Arc; + +use simplicityhl::simplicity::bitcoin::secp256k1; +use simplicityhl::simplicity::elements::hashes::HashEngine as _; +use simplicityhl::simplicity::elements::taproot::{LeafVersion, TaprootBuilder, TaprootSpendInfo}; +use simplicityhl::simplicity::elements::{Script, Transaction}; +use simplicityhl::simplicity::hashes::{Hash, sha256}; +use simplicityhl::simplicity::jet::Elements; +use simplicityhl::simplicity::jet::elements::ElementsEnv; +use simplicityhl::simplicity::{Cmr, RedeemNode, leaf_version}; +use simplicityhl::tracker::TrackerLogLevel; +use simplicityhl::{Arguments, CompiledProgram, TemplateProgram}; +use simplicityhl_core::{ProgramError, run_program}; + +mod build_witness; + +pub use build_witness::{DEPTH, SMTWitness, build_smt_storage_witness, u256}; + +pub const SMT_STORAGE_SOURCE: &str = include_str!("source_simf/smt_storage.simf"); + +/// Get the storage template program for instantiation. +/// +/// # Panics +/// +/// Panics if the embedded source fails to compile (should never happen). +#[must_use] +pub fn get_smt_storage_template_program() -> TemplateProgram { + TemplateProgram::new(SMT_STORAGE_SOURCE).expect("INTERNAL: expected to compile successfully.") +} + +/// Get compiled storage program, panicking on failure. +/// +/// # Panics +/// +/// Panics if program instantiation fails. +#[must_use] +pub fn get_smt_storage_compiled_program() -> CompiledProgram { + let program = get_smt_storage_template_program(); + + program.instantiate(Arguments::default(), true).unwrap() +} + +/// Execute storage program with new state. +/// +/// # Errors +/// Returns error if program execution fails. +pub fn execute_smt_storage_program( + witness: &SMTWitness, + compiled_program: &CompiledProgram, + env: &ElementsEnv>, + runner_log_level: TrackerLogLevel, +) -> Result>, ProgramError> { + let witness_values = build_smt_storage_witness(witness); + Ok(run_program(compiled_program, witness_values, env, runner_log_level)?.0) +} + +fn smt_storage_script_ver(cmr: Cmr) -> (Script, LeafVersion) { + (Script::from(cmr.as_ref().to_vec()), leaf_version()) +} + +/// Computes the TapData-tagged hash of the Simplicity state. +/// +/// This involves hashing the tag "`TapData`" twice, followed by the +/// limbs of the state. +/// +/// # Panics +/// +/// This function **does not panic**. +/// All hashing operations (`sha256::Hash::engine`, `input`, `from_engine`) are +/// infallible, and iterating over the state limbs is safe. +#[must_use] +pub fn compute_tapdata_tagged_hash_of_the_state(leaf: &u256, path: &[bool; DEPTH]) -> sha256::Hash { + let tag = sha256::Hash::hash(b"TapData"); + let mut eng = sha256::Hash::engine(); + eng.input(tag.as_byte_array()); + eng.input(tag.as_byte_array()); + eng.input(leaf); + + let mut current_hash = sha256::Hash::from_engine(eng); + + // Change to valid tree hashes + let dummy_hash = [0u8; 32]; + + for is_right_direction in path { + let mut eng = sha256::Hash::engine(); + dbg!(*is_right_direction); + + if *is_right_direction { + eng.input(&dummy_hash); + eng.input(¤t_hash.to_byte_array()); + } else { + eng.input(¤t_hash.to_byte_array()); + eng.input(&dummy_hash); + } + + current_hash = sha256::Hash::from_engine(eng); + } + current_hash +} + +/// Given a Simplicity CMR and an internal key, computes the [`TaprootSpendInfo`] +/// for a Taptree with this CMR as its single leaf. +/// +/// # Panics +/// +/// This function **panics** if building the taproot tree fails (the calls to +/// `TaprootBuilder::add_leaf_with_ver` or `.add_hidden` return `Err`) or if +/// finalizing the builder fails. Those panics come from the `.expect(...)` +/// calls on the builder methods. +#[must_use] +pub fn smt_storage_taproot_spend_info( + internal_key: secp256k1::XOnlyPublicKey, + leaf: &u256, + path: &[bool; DEPTH], + cmr: Cmr, +) -> TaprootSpendInfo { + let (script, version) = smt_storage_script_ver(cmr); + let state_hash = compute_tapdata_tagged_hash_of_the_state(leaf, path); + + // Build taproot tree with hidden leaf + let builder = TaprootBuilder::new() + .add_leaf_with_ver(1, script, version) + .expect("tap tree should be valid") + .add_hidden(1, state_hash) + .expect("tap tree should be valid"); + + builder + .finalize(secp256k1::SECP256K1, internal_key) + .expect("tap tree should be valid") +} + +#[cfg(test)] +mod smt_storage_tests { + use super::*; + use anyhow::Result; + use std::sync::Arc; + + use simplicityhl::elements::confidential::{Asset, Value}; + use simplicityhl::elements::pset::{Input, Output, PartiallySignedTransaction}; + use simplicityhl::elements::{AssetId, BlockHash, OutPoint, Script, Txid}; + use simplicityhl::simplicity::elements::taproot::ControlBlock; + use simplicityhl::simplicity::jet::elements::{ElementsEnv, ElementsUtxo}; + + #[rustfmt::skip] // mangles byte vectors + fn smt_storage_unspendable_internal_key() -> secp256k1::XOnlyPublicKey { + secp256k1::XOnlyPublicKey::from_slice(&[ + 0x50, 0x92, 0x9b, 0x74, 0xc1, 0xa0, 0x49, 0x54, 0xb7, 0x8b, 0x4b, 0x60, 0x35, 0xe9, 0x7a, 0x5e, + 0x07, 0x8a, 0x5a, 0x0f, 0x28, 0xec, 0x96, 0xd5, 0x47, 0xbf, 0xee, 0x9a, 0xce, 0x80, 0x3a, 0xc0, + ]) + .expect("key should be valid") + } + + #[test] + fn test_smt_storage_mint_path() -> Result<()> { + let old_leaf = [0u8; 32]; + let mut path = [true; DEPTH]; + path[1] = false; + path[4] = false; + + let merkle_data = path.map(|is_right| ([0u8; 32], is_right)); + let witness = SMTWitness::new(&old_leaf, &merkle_data); + + let mut new_leaf = old_leaf; + new_leaf[31] = 1; + + let program = get_smt_storage_compiled_program(); + let cmr = program.commit().cmr(); + + let old_spend_info = smt_storage_taproot_spend_info( + smt_storage_unspendable_internal_key(), + &old_leaf, + &path, + cmr, + ); + let old_script_pubkey = Script::new_v1_p2tr_tweaked(old_spend_info.output_key()); + + let new_spend_info = smt_storage_taproot_spend_info( + smt_storage_unspendable_internal_key(), + &new_leaf, + &path, + cmr, + ); + let new_script_pubkey = Script::new_v1_p2tr_tweaked(new_spend_info.output_key()); + + let mut pst = PartiallySignedTransaction::new_v2(); + let outpoint0 = OutPoint::new(Txid::from_slice(&[0; 32])?, 0); + pst.add_input(Input::from_prevout(outpoint0)); + pst.add_output(Output::new_explicit( + new_script_pubkey.clone(), + 0, + AssetId::default(), + None, + )); + + let control_block = old_spend_info + .control_block(&smt_storage_script_ver(cmr)) + .expect("must get control block"); + + let env = ElementsEnv::new( + Arc::new(pst.extract_tx()?), + vec![ElementsUtxo { + script_pubkey: old_script_pubkey, + asset: Asset::default(), + value: Value::default(), + }], + 0, + cmr, + ControlBlock::from_slice(&control_block.serialize())?, + None, + BlockHash::all_zeros(), + ); + + assert!( + execute_smt_storage_program(&witness, &program, &env, TrackerLogLevel::Trace,).is_ok(), + "expected success mint path" + ); + + Ok(()) + } +} diff --git a/crates/contracts/src/smt_storage/source_simf/smt_storage.simf b/crates/contracts/src/smt_storage/source_simf/smt_storage.simf new file mode 100644 index 0000000..6eed3fc --- /dev/null +++ b/crates/contracts/src/smt_storage/source_simf/smt_storage.simf @@ -0,0 +1,65 @@ +/* + * Extends `bytes32_tr_storage` using `array_fold` for larger buffers. + * Optimized for small, fixed-size states where linear hashing is more efficient + * than Merkle Trees. By avoiding proof overhead like sibling hashes, we reduce + * witness size and simplify contract logic for small N. + */ +fn hash_array_tr_storage_with_update(elem: (u256, bool), prev_hash: u256) -> u256 { + let (hash, is_right): (u256, bool) = dbg!(elem); + let ctx: Ctx8 = jet::sha_256_ctx_8_init(); + + let new_hash: Ctx8 = match is_right { + true => { + let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, hash); + jet::sha_256_ctx_8_add_32(ctx, prev_hash) + }, + false => { + let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, prev_hash); + jet::sha_256_ctx_8_add_32(ctx, hash) + } + }; + + jet::sha_256_ctx_8_finalize(new_hash) +} + +fn script_hash_for_input_script(leaf: u256, merkle_data: [(u256, bool); 7]) -> u256 { + let tap_leaf: u256 = jet::tapleaf_hash(); + let ctx: Ctx8 = jet::tapdata_init(); + let ctx: Ctx8 = jet::sha_256_ctx_8_add_32(ctx, leaf); + let hash_leaf: u256 = jet::sha_256_ctx_8_finalize(ctx); + + let computed: u256 = array_fold::(merkle_data, hash_leaf); + let tap_node: u256 = jet::build_tapbranch(tap_leaf, computed); + + let bip0341_key: u256 = 0x50929b74c1a04954b78b4b6035e97a5e078a5a0f28ec96d547bfee9ace803ac0; + let tweaked_key: u256 = jet::build_taptweak(bip0341_key, tap_node); + + let hash_ctx1: Ctx8 = jet::sha_256_ctx_8_init(); + let hash_ctx2: Ctx8 = jet::sha_256_ctx_8_add_2(hash_ctx1, 0x5120); // Segwit v1, length 32 + let hash_ctx3: Ctx8 = jet::sha_256_ctx_8_add_32(hash_ctx2, tweaked_key); + jet::sha_256_ctx_8_finalize(hash_ctx3) +} + +fn main() { + let leaf_data: u256 = witness::LEAF; + + // Path and hash + let merkle_data: [(u256, bool); 7] = witness::MERKLE_DATA; + let (leaf1, leaf2, leaf3, leaf4): (u64, u64, u64, u64) = ::into(leaf_data); + + // Load + assert!(jet::eq_256( + script_hash_for_input_script(leaf_data, merkle_data), + unwrap(jet::input_script_hash(jet::current_index())) + )); + + // There may be arbitrary logic here + let new_leaf4: u64 = 1; + let new_leaf: u256 = <(u64, u64, u64, u64)>::into((leaf1, leaf2, leaf3, new_leaf4)); + + // Store + assert!(jet::eq_256( + script_hash_for_input_script(new_leaf, merkle_data), + unwrap(jet::output_script_hash(jet::current_index())) + )); +} \ No newline at end of file From cfc92e006dcf0ffffe59822918109c68be9668eb Mon Sep 17 00:00:00 2001 From: LesterEvSe Date: Mon, 19 Jan 2026 16:02:02 +0200 Subject: [PATCH 2/2] feat: add smt to rust and fill with some elements --- Cargo.lock | 83 +++++++++++-- crates/contracts/Cargo.toml | 3 +- crates/contracts/src/smt_storage/mod.rs | 57 ++++++--- crates/contracts/src/smt_storage/smt.rs | 159 ++++++++++++++++++++++++ 4 files changed, 274 insertions(+), 28 deletions(-) create mode 100644 crates/contracts/src/smt_storage/smt.rs diff --git a/Cargo.lock b/Cargo.lock index c77bddf..1d57b20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -337,7 +337,7 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" dependencies = [ - "getrandom", + "getrandom 0.2.16", "once_cell", "tiny-keccak", ] @@ -349,6 +349,7 @@ dependencies = [ "anyhow", "bincode", "hex", + "rand 0.9.2", "ring", "sha2", "simplicityhl", @@ -530,6 +531,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + [[package]] name = "ghost-cell" version = "0.2.6" @@ -815,6 +828,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "rand" version = "0.8.5" @@ -822,8 +841,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -833,7 +862,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -842,7 +881,16 @@ version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" dependencies = [ - "getrandom", + "getrandom 0.2.16", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", ] [[package]] @@ -891,7 +939,7 @@ checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" dependencies = [ "cc", "cfg-if", - "getrandom", + "getrandom 0.2.16", "libc", "untrusted", "windows-sys 0.52.0", @@ -981,7 +1029,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9465315bc9d4566e1724f0fffcbcc446268cb522e60f9a27bcded6b19c108113" dependencies = [ "bitcoin_hashes", - "rand", + "rand 0.8.5", "secp256k1-sys", ] @@ -1001,7 +1049,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52a44aed3002b5ae975f8624c5df3a949cfbf00479e18778b6058fcd213b76e3" dependencies = [ "bitcoin-private", - "rand", + "rand 0.8.5", "secp256k1", "secp256k1-zkp-sys", ] @@ -1107,7 +1155,7 @@ dependencies = [ "bitcoin_hashes", "byteorder", "elements", - "getrandom", + "getrandom 0.2.16", "ghost-cell", "hex-conservative", "miniscript", @@ -1133,7 +1181,7 @@ dependencies = [ "base64", "clap", "either", - "getrandom", + "getrandom 0.2.16", "itertools", "miniscript", "pest", @@ -1353,6 +1401,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasip2" +version = "1.0.1+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.106" @@ -1523,6 +1580,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "wit-bindgen" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" + [[package]] name = "yaml-rust2" version = "0.10.4" diff --git a/crates/contracts/Cargo.toml b/crates/contracts/Cargo.toml index 91897a4..dc58b4d 100644 --- a/crates/contracts/Cargo.toml +++ b/crates/contracts/Cargo.toml @@ -39,4 +39,5 @@ simplicityhl = { workspace = true } simplicityhl-core = { workspace = true } [dev-dependencies] -anyhow = "1" \ No newline at end of file +anyhow = "1" +rand = "0.9.2" \ No newline at end of file diff --git a/crates/contracts/src/smt_storage/mod.rs b/crates/contracts/src/smt_storage/mod.rs index db72f14..d7214b0 100644 --- a/crates/contracts/src/smt_storage/mod.rs +++ b/crates/contracts/src/smt_storage/mod.rs @@ -13,8 +13,10 @@ use simplicityhl::{Arguments, CompiledProgram, TemplateProgram}; use simplicityhl_core::{ProgramError, run_program}; mod build_witness; +mod smt; pub use build_witness::{DEPTH, SMTWitness, build_smt_storage_witness, u256}; +pub use smt::SparseMerkleTree; pub const SMT_STORAGE_SOURCE: &str = include_str!("source_simf/smt_storage.simf"); @@ -69,7 +71,10 @@ fn smt_storage_script_ver(cmr: Cmr) -> (Script, LeafVersion) { /// All hashing operations (`sha256::Hash::engine`, `input`, `from_engine`) are /// infallible, and iterating over the state limbs is safe. #[must_use] -pub fn compute_tapdata_tagged_hash_of_the_state(leaf: &u256, path: &[bool; DEPTH]) -> sha256::Hash { +pub fn compute_tapdata_tagged_hash_of_the_state( + leaf: &u256, + path: &[(u256, bool); DEPTH], +) -> sha256::Hash { let tag = sha256::Hash::hash(b"TapData"); let mut eng = sha256::Hash::engine(); eng.input(tag.as_byte_array()); @@ -78,19 +83,15 @@ pub fn compute_tapdata_tagged_hash_of_the_state(leaf: &u256, path: &[bool; DEPTH let mut current_hash = sha256::Hash::from_engine(eng); - // Change to valid tree hashes - let dummy_hash = [0u8; 32]; - - for is_right_direction in path { + for (hash, is_right_direction) in path { let mut eng = sha256::Hash::engine(); - dbg!(*is_right_direction); if *is_right_direction { - eng.input(&dummy_hash); + eng.input(hash); eng.input(¤t_hash.to_byte_array()); } else { eng.input(¤t_hash.to_byte_array()); - eng.input(&dummy_hash); + eng.input(hash); } current_hash = sha256::Hash::from_engine(eng); @@ -111,7 +112,7 @@ pub fn compute_tapdata_tagged_hash_of_the_state(leaf: &u256, path: &[bool; DEPTH pub fn smt_storage_taproot_spend_info( internal_key: secp256k1::XOnlyPublicKey, leaf: &u256, - path: &[bool; DEPTH], + path: &[(u256, bool); DEPTH], cmr: Cmr, ) -> TaprootSpendInfo { let (script, version) = smt_storage_script_ver(cmr); @@ -133,6 +134,7 @@ pub fn smt_storage_taproot_spend_info( mod smt_storage_tests { use super::*; use anyhow::Result; + use rand::Rng as _; use std::sync::Arc; use simplicityhl::elements::confidential::{Asset, Value}; @@ -150,26 +152,47 @@ mod smt_storage_tests { .expect("key should be valid") } + fn add_elements(smt: &mut SparseMerkleTree, num: u64) -> (u256, [u256; DEPTH], [bool; DEPTH]) { + let mut rng = rand::rng(); + + let mut leaf = [0u8; 32]; + let mut merkle_hashes = [[0u8; 32]; DEPTH]; + let mut path = [false; DEPTH]; + + for _ in 0..num { + leaf = rng.random(); + path = std::array::from_fn(|_| rng.random()); + merkle_hashes = smt.update(&leaf, path); + } + + (leaf, merkle_hashes, path) + } + #[test] fn test_smt_storage_mint_path() -> Result<()> { - let old_leaf = [0u8; 32]; - let mut path = [true; DEPTH]; - path[1] = false; - path[4] = false; + let mut smt = SparseMerkleTree::new(); + let (old_leaf, merkle_hashes, path) = add_elements(&mut smt, 30); + + let merkle_data = + std::array::from_fn(|i| (merkle_hashes[DEPTH - i - 1], path[DEPTH - i - 1])); - let merkle_data = path.map(|is_right| ([0u8; 32], is_right)); let witness = SMTWitness::new(&old_leaf, &merkle_data); + // Set last leaf qword to 1 let mut new_leaf = old_leaf; + for byte in new_leaf.iter_mut().skip(24) { + *byte = 0; + } new_leaf[31] = 1; + smt.update(&new_leaf, path); let program = get_smt_storage_compiled_program(); let cmr = program.commit().cmr(); - let old_spend_info = smt_storage_taproot_spend_info( + let old_spend_info: TaprootSpendInfo = smt_storage_taproot_spend_info( smt_storage_unspendable_internal_key(), &old_leaf, - &path, + &merkle_data, cmr, ); let old_script_pubkey = Script::new_v1_p2tr_tweaked(old_spend_info.output_key()); @@ -177,7 +200,7 @@ mod smt_storage_tests { let new_spend_info = smt_storage_taproot_spend_info( smt_storage_unspendable_internal_key(), &new_leaf, - &path, + &merkle_data, cmr, ); let new_script_pubkey = Script::new_v1_p2tr_tweaked(new_spend_info.output_key()); diff --git a/crates/contracts/src/smt_storage/smt.rs b/crates/contracts/src/smt_storage/smt.rs new file mode 100644 index 0000000..a60724e --- /dev/null +++ b/crates/contracts/src/smt_storage/smt.rs @@ -0,0 +1,159 @@ +use simplicityhl::simplicity::elements::hashes::HashEngine as _; +use simplicityhl::simplicity::hashes::{Hash, sha256}; + +use super::build_witness::{DEPTH, u256}; + +#[derive(Debug, Clone, bincode::Encode, bincode::Decode, PartialEq, Eq)] +pub enum TreeNode { + Leaf { + leaf_hash: u256, + }, + Branch { + hash: u256, + left: Box, + right: Box, + }, +} + +impl TreeNode { + pub fn get_hash(&self) -> u256 { + match self { + TreeNode::Leaf { leaf_hash } => *leaf_hash, + TreeNode::Branch { hash, .. } => *hash, + } + } +} + +pub struct SparseMerkleTree { + root: Box, + precalculate_hashes: [u256; DEPTH], +} + +impl SparseMerkleTree { + #[must_use] + pub fn new() -> Self { + let mut precalculate_hashes = [[0u8; 32]; DEPTH]; + let mut eng = sha256::Hash::engine(); + let zero = [0u8; 32]; + eng.input(&zero); + precalculate_hashes[0] = *sha256::Hash::from_engine(eng).as_byte_array(); + + for i in 1..DEPTH { + let mut eng = sha256::Hash::engine(); + eng.input(&precalculate_hashes[i - 1]); + eng.input(&precalculate_hashes[i - 1]); + precalculate_hashes[i] = *sha256::Hash::from_engine(eng).as_byte_array(); + } + + Self { + root: Box::new(TreeNode::Leaf { + leaf_hash: precalculate_hashes[0], + }), + precalculate_hashes, + } + } + + fn calculate_hash(left: &mut TreeNode, right: &mut TreeNode) -> u256 { + let mut eng = sha256::Hash::engine(); + eng.input(&left.get_hash()); + eng.input(&right.get_hash()); + *sha256::Hash::from_engine(eng).as_byte_array() + } + + // Return array of hashes + fn traverse( + defaults: &[u256], + leaf: &u256, + path: &[bool], + root: &mut Box, + hashes: &mut [u256], + ) { + let Some((is_right, remaining_path)) = path.split_first() else { + let tag = sha256::Hash::hash(b"TapData"); + let mut eng = sha256::Hash::engine(); + eng.input(tag.as_byte_array()); + eng.input(tag.as_byte_array()); + eng.input(leaf); + + **root = TreeNode::Leaf { + leaf_hash: *sha256::Hash::from_engine(eng).as_byte_array(), + }; + return; + }; + + let (child_zero, remaining_defaults) = defaults + .split_last() + .expect("Defaults length must match path length"); + + if matches!(**root, TreeNode::Leaf { .. }) { + let new_branch = Box::new(TreeNode::Branch { + hash: [0u8; 32], + left: Box::new(TreeNode::Leaf { + leaf_hash: *child_zero, + }), + right: Box::new(TreeNode::Leaf { + leaf_hash: *child_zero, + }), + }); + + *root = new_branch; + } + + let (current_hash_slot, remaining_hashes) = hashes + .split_first_mut() + .expect("Hashes length must match path length"); + + if let TreeNode::Branch { + ref mut left, + ref mut right, + ref mut hash, + } = **root + { + if *is_right { + *current_hash_slot = left.get_hash(); + Self::traverse( + remaining_defaults, + leaf, + remaining_path, + right, + remaining_hashes, + ); + } else { + *current_hash_slot = right.get_hash(); + Self::traverse( + remaining_defaults, + leaf, + remaining_path, + left, + remaining_hashes, + ); + } + + *hash = Self::calculate_hash(left, right); + } else { + unreachable!("Should be a branch at this point"); + } + } + + // For insert change 0 to leaf. + // For delete vice versa. + // And for udpate change old value to new. + // Then, recalculate hashes + pub fn update(&mut self, leaf: &u256, path: [bool; DEPTH]) -> [u256; DEPTH] { + let mut hashes = self.precalculate_hashes; + Self::traverse( + &self.precalculate_hashes, + leaf, + &path, + &mut self.root, + &mut hashes, + ); + hashes + } +} + +impl Default for SparseMerkleTree { + fn default() -> Self { + Self::new() + } +}