Add op_sha256tree

.gitignore

@@ -23,3 +23,6 @@ target/
 # Node.js
 /target
 /node_modules
+
+# Mac generates these with the benchmarking tools
+.DS_Store

fuzz/fuzz_targets/operators.rs

@@ -20,10 +20,11 @@ use clvmr::more_ops::{
 };
 use clvmr::reduction::Response;
 use clvmr::secp_ops::{op_secp256k1_verify, op_secp256r1_verify};
+use clvmr::sha_tree_op::op_sha256_tree;
 
 type Opf = fn(&mut Allocator, NodePtr, Cost) -> Response;
 
-const FUNS: [Opf; 46] = [
+const FUNS: [Opf; 47] = [
     op_if as Opf,
     op_cons as Opf,
     op_first as Opf,
@@ -73,6 +74,8 @@ const FUNS: [Opf; 46] = [
     op_secp256r1_verify as Opf,
     // keccak operator
     op_keccak256 as Opf,
+    // shatree operator
+    op_sha256_tree as Opf,
 ];
 
 fuzz_target!(|data: &[u8]| {

op-tests/test-sha256tree-hash.txt

@@ -0,0 +1,32 @@
+; This file was generated by tools/generate-sha256tree-tests.py
+
+sha256tree (((((0 . 0x1f7061813451ffb2830d0aab69351d593dab69e339835167) . (0 . 0xd09ba26eb2ba7592a67659ca19fc923a1631b38cd5aebb0b53baa7c68d2d89faad5c2f8e2d9ffeb9714f50a59a21ab4e)) . 0x02) . (((0x177ba5f23f0e140f805254e78fac8c2bee4fe850fc37ead18cb51314c9181e31 . 0x666f6f626172) . (0x02 . 0x666f6f626172)) . 0)) . (((0xc4c6b5a3bc74be71b5631a1dfd4cc0aae3ee363ecbb0214d . (0x01 . 0)) . (0x06d17ae983679967e26dd593cf637c98951adb3e90f349ec421212eaa30109e51b12120fc9ac3d6e527e383aaad939ae . (0 . 0x5c8b34f9fac0f365af055ef59678414793a7a5859715ea8f76bfa5ea13ff7c49))) . (0x01 . 0))) => 0x6768523ba0f32733a54551b101688a6300311e0b10e212049226d98bcd7f869d | 22579
+sha256tree (0xcdbe2cffebfda5b5dfe499e05894967d10d0e55d5bd0c9d3 . (0x16f905cd314124664cf8070035bff4faa40fdea48447770a . (0x01 . 0))) => 0x157c1e83a68a2c6cbbd92aa7ff82552b0cb2679885858695f85ed9404126ac80 | 4739
+sha256tree 0xc3041da8237d1d10610210bbad12e483fcce196c80cd189f650e6d86fda57f92 => 0x291b0356fa82ec28b5db8254c56f368273478de1fc6c8d7b85ca80706958cfa0 | 1035
+sha256tree ((0xfac7754adb5942ea853a150bbfe72c4165d7a36b35bf125240b81e5764cf52b04c6ec11682626aa49e3872e68979808a . 0) . (0x666f6f626172 . 0x666f6f626172)) => 0x9ff9d9c782362027084f4ef5ce131621d35ecde7e9cf5320b36ad7e45cf1c7e1 | 4803
+sha256tree (((0 . 0x02) . (0x18dd1f6517ae25cf64f9ddc232b7a756c3dfe0b60221d62d . 0x02)) . 0) => 0x591fbedc01c82c80596e03d4dcd75532e40520577008c6d5c6764ac2a984aa27 | 5995
+sha256tree ((0x666f6f626172 . 0x02) . 0x2a7776815b9b83c9c9049b6f5218739552e0dba6937b5e69943776aef44c9dff78f688d40641dddf52bf6818417523df) => 0xfd2359fd4190f1424be253889a1604cf8e106f0b1b791165e262dd8a652d0a05 | 3547
+sha256tree 0x68ed3d3389d15111a797fdf9434045841b3a3753d93310d593cd3c80ff27e2213e1faf566c4565ffede6bc89ec96bf89 => 0x8c9d292643d2b79627b488e277710209b0bf33e2526f96ba64019d17faf0e7d8 | 1035
+sha256tree (((0 . 0x119c26c141c87f1359bbfb2a82afd7a060043f446760e7f60229a485388eaa68) . (0x01 . 0x02)) . (0xdcc048db813648ccb794f644e9cafc6094bbe1bea5ef388f . (0x02 . 0x9c7483847dc6e351081a37c527bbda1e24b015e31f60d353b96370b77f26530d))) => 0xdec44a4d2d8752c305825b09d9e0b24031628f33f041354ccb2fc457131e750c | 8635
+sha256tree ((0 . (((0 . 0xa5c8efe9fc85db1b35d96edf0d1fef862c30fea73d1c5318de02dd1504cb88e8fded7f1ef8479b9f37204f368ac8e379) . (0x02 . 0x0df17fa06f30f7429162e282f00722e520e0a2676c69cfd2b92f678b28729524)) . 0x666f6f626172)) . (0x02 . (((0x8302df3c2a02bcca20648da3a1aa96d3f9dbd47960d90a53 . 0) . (0x01 . 0x02)) . ((0xa31acd511f4560b50be8bf18d47c1f422bab6bcdda0e69ea0152ad49c28833a78c2d4613ec8aa512cdf3a948eadcd7b9 . 0x02) . (0x6b3b39ef6fe8fb41a13a495599cadb2ded8da6ddd25518b1 . 0x01))))) => 0x7f6f3a1dd8af7087774438b0ab2770dcaacf3f42d632f3544a9f5b63c7378212 | 18747
+sha256tree 0x51cde38bba7b68ba39903e1e83c152969d896608efde132ce9f9f4bcdc6833e5fa3d21238a6cf5cceb132adf1d371a60 => 0xa2d4a03a19d0ed8bf52df8a7e26433e98c672be16ffc2137ceb7209be333c98b | 1035
+sha256tree (((0x02 . 0xa21186440ff93ecba8304bfee4997e3618b0c84c31b21533) . (0 . 0x28b70cdb54d9cc014c6aea054082229d47ea269a9b280770)) . ((0xa8dc249448c80f5004f13abd6401ad4bcb6b52a97c5a359b . 0x0f0df3479340a2d5a7d7c47cb0ea209f60953c9f85a423b5) . (0x02 . 0x666f6f626172))) => 0x1ba6e28d4eb3b955faa2d6f348261a6be3a5427d44489647ab06e366de0b2479 | 9763
+sha256tree (0x104e55ad9b1db96359ed05d6ed3ec4c207c3b6213b46b974 . 0x01) => 0xeaadf0be76b05cbbdaaae7f3d805e678039814c0a5620794ae8961052f123400 | 2227
+sha256tree (((0x6f815e4db56a306afba86916d4b29b194ceed7f492fc8a33a85d9a15cfe0536d . 0x01) . 0x95199b2f16397b23a2a70704c656f51b28e3d3d22680823d76ae4029ec4eb71a12760c82b4441a9d39b5baec010f6b7d) . 0x3c370f10cd9ec3d572cfa6a8724856b6c68a0e60912fa563f0a182323bcd6d31) => 0x4c71c2c6b8d98390697a61bd228bf1113b685f835e97552cf62feeef0d7069d1 | 4931
+sha256tree 0x7a9aa9d0946b816656299f1bd609cd0692602ce26f2f124649955b4ba54b9e8d16dcf5b392bed7189da48020f79c5bbd => 0x28a2096bc8056d3b3287a605b2dd183489e4eeb4113093d841243c2a08d83719 | 1035
+sha256tree ((((0xb941c50724b5c81213f40b621efbec33d30d2b1839ca80c81a7728bde1ba73fa . 0x8bd83d18af286870a712bb04d61d4b7fe29489cd0c7fced8dec078b631a9b0d9) . (0x02 . 0x02)) . (0x02 . 0x02)) . (((0x01 . 0x01) . (0 . 0x666f6f626172)) . ((0x01 . 0x02) . (0x01 . 0x666f6f626172)))) => 0x883633e181c8d3aee1667d530d4a2807ef242c4f961678f301b0b40ee55ab948 | 17427
+sha256tree ((0 . 0x01) . 0x01) => 0x9b28d247c983abca94e6961e6cfe6623cb9e9415d0f9e0cb04a46b3ae43369bd | 3483
+sha256tree (0x02 . 0xd6af8d56786198cde353189f35566ece010a122c7c6f181991803d7cff673b7f) => 0x938ba2209eb55f5daa22c51677f94adbbdf84566c6fb3b89ff0fa8d08d5905f2 | 2291
+sha256tree (((0 . (0xa79fbde48fc3ddeccdd64774441ab7bbd6a34f5ad1a6d78e3a3a960eea3dfbf0 . 0x666f6f626172)) . ((0xe02c71ebcfdfef147f01ed578a0844726265fb0e86cf091f93308866491f72d2 . 0x50e86350e1ddae6f2289626bd722568991a62b94f9f9ef1d4e19aa9cf5c8aee7) . (0xe98e1cfcff183a5637969e9ae20e07acb20bd662fa4262b9 . 0x02))) . (0xbb9437c1ba52ab7ca465599c57ab6d71991fca001d67274028407a245406f34f . ((0xb9a4001e4d6b805d3a3d8b6b7976ef91810056bbaca737739b8e0bda6fd5b661 . 0x02) . (0x01 . 0x666f6f626172)))) => 0xca1aa87e0b15e86b3cf1a975bb61440b607d3ded9a6d95300b19c59765d575c5 | 15107
+sha256tree ((0 . 0x02) . 0x882d2f48ed804484d9f23d38053f523e3574119071b3a7f6) => 0xb64c99c57ce8bac8cd73351de0eae02473271755f053504a7270fe02921d1a0e | 3483
+sha256tree (((0x666f6f626172 . 0x02) . (0 . 0x666f6f626172)) . 0x3d2f00fbfd0278404cd508c4afb460f55b116e5733240c4f4a30a5a790beed1c) => 0x3f6498812d3a5d119db27f679b938558e7d74c9cd921bc28bb6245823ca1f7fe | 6059
+sha256tree 0x57afc714e65003547db455176d5f85cc6594b88e5d28f87bf49ec82135fa791b09d6a102c34787fec98544486b888023 => 0xbc6c9852407944ab4bd3749b61216e3ab64c1d1e11b56be0cc30fd0344e34085 | 1035
+sha256tree ((0x3e9be38550994f6fa12c2ba30bc582ceda9f715d5a3a4af06d3a4f35d1b5d242a0b825d74a3b7e3086590b22172f5eeb . 0x02) . 0) => 0xc8d024ec5e8b001598920071454558cd6a2831742829cb507e744d79b0dbb09c | 3547
+sha256tree (0 . 0x51e7d76d53afdbbe486df4f6eb42197aa985154283ed028c) => 0xbc02de5bb1e9feccc7adc78cdc6ca567c0311df0543e46434dc795edebc67dbf | 2227
+sha256tree 0x666f6f626172 => 0xf03942eca4827c93931fee97f117479ef474c9aaa449655ddffb48886bde58ad | 971
+sha256tree (0x934128806851cb5eba9d244dfce1734a41c190bde679560a22fe90a87a4f775c . (0x02 . ((0 . 0x02) . (0x62e282852844353b813fd3a967e9af12b596a2ba912236363ed6869f2669546e . 0x02)))) => 0x9063f5dcb016de08e12d1f8b4c893a9f1c008d08cd0c543c5707a2a328feb6c9 | 7379
+sha256tree (((((0x02 . 0x02) . (0 . 0x01)) . (0xd69626fe507014737c49d9b7cf950721c7142ce1b217d2c4c654b2274615129a2374a9a5700d37188047b55ecf1d55dd . (0x4dca4553f9115e233dda1e184759e201f1de5fe11d3f2dc0 . 0xbb2aa251bf2f429ddf6a48a0d70ca6ceeb9dda50428cabde))) . (((0x02 . 0) . 0x01) . (0x16f07f1f29699553d8064284a6a324b906ff081a195075e5a9d0c6c399b2492a2e13aa69bd7a236ae2de60a2a7725306 . 0xfc7fd2354396ab20b07d6f07259c8691456ef7989feee44a04f2d87c96d7b4cecfc039e2ba04f86da20b062bc9ccda38))) . 0x1e122d35a2ac075a48dc607a8a428586a3a06591f6e20e71c95bb94dcaaf95a6165d97a7d8f90e8ae5517468cba60afc) => 0xcf43ccb9fdabc87ac031c3f372addd99e91103ef5d89e389a186e0215a620be0 | 16299
+sha256tree ((0x4295ead455b30626cd3dc0aed97dd4be40f4fe965f337484092ba6650c1e5fb0 . ((0x666f6f626172 . 0x01) . (0x231bcedef942bcc4b1c99f8b26382a7433866dc96cd617f59d39ebd5f10fe23b . 0x02))) . (((0xa27954b60506b8f2cb4c56398797eb9d640ca976b9c21126a39e14af0fdb94ac . 0x01) . 0x67a0ad6c4e30c5e702cc723d57b00ff3bf818fc553093b1970aa6a9c2698f87098fbf322ee3047d144c266b2f425da32) . 0)) => 0xa0bf50e2536ab240ad5e55d8ff761722a87fef03f1d57eb13566cb7a41185db0 | 11275
+sha256tree (0x0fb2228f690832f6f2d1f08afeeaf00352b182d49a7db18c5bc335b95609d9503ba3f97054599526ae316a63a4d72fed . 0x64acfa2fd1990e8b55989a4ff9544b7368e9eccfe3c98716a1c330db640be65b) => 0xef13a74e8149d647e668f323049257a5f0ae16ec6d1f64a2ca2404cb670f233a | 2355
+sha256tree 0x01 => 0x9dcf97a184f32623d11a73124ceb99a5709b083721e878a16d78f596718ba7b2 | 971
+sha256tree (0x02 . (0 . 0x666f6f626172)) => 0xfd1c831b92f4ca52591d92502f97d00dbe004f454873c989bd1161978cb06ab6 | 3483

op-tests/test-sha256tree.txt

@@ -0,0 +1,10 @@
+sha256tree 1 => 0x9dcf97a184f32623d11a73124ceb99a5709b083721e878a16d78f596718ba7b2 | 971
+sha256tree 0x00cafef00d => 0x60bc5062a80c4363cbe881815300d349c5524d98e3502a016466d14d1f3f25d9 | 971
+sha256tree () => 0x4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a | 971
+sha256tree (() . ()) => 0x52db9ef97986e7382ef78b8eae2dacdbb2ce823ed1396a0fb2f7f120a2b40a63 | 2227
+sha256tree (0x00cafe . 0x00f00d) => 0x005b24012b958cc8fc19b7d45438f82fb9034c24b2a77d219250316f482ba6b6 | 2227
+sha256tree (202 254 240 13) => 0xb5e2a23ac2de5105c72b9658428a3d389a9f909012eedd4b641504d981f81ac5 | 5995
+sha256tree (202 254 (240 13)) => 0x5edd993092c97418933750dc25add55a84ac9f07b2067d537d38f0044bd317ae | 7251
+sha256tree 10 20 => FAIL
+sha256tree (10 . 20) 30 => FAIL
+sha256tree (10 . 20) (20 . 30) => FAIL

src/chia_dialect.rs

@@ -16,6 +16,7 @@ use crate::more_ops::{
 };
 use crate::reduction::Response;
 use crate::secp_ops::{op_secp256k1_verify, op_secp256r1_verify};
+use crate::sha_tree_op::op_sha256_tree;
 
 // unknown operators are disallowed
 // (otherwise they are no-ops with well defined cost)
@@ -29,7 +30,11 @@ pub const LIMIT_HEAP: u32 = 0x0004;
 // This is a hard-fork and should only be enabled when it activates
 pub const ENABLE_KECCAK_OPS_OUTSIDE_GUARD: u32 = 0x0100;
 
-// The default mode when running grnerators in mempool-mode (i.e. the stricter
+// this flag enables the sha256tree op *outside* the softfork guard.
+// This is a hard-fork and should only be enabled when it activates.
+pub const ENABLE_SHA256_TREE: u32 = 0x0200;
+
+// The default mode when running generators in mempool-mode (i.e. the stricter
 // mode)
 pub const MEMPOOL_MODE: u32 = NO_UNKNOWN_OPS | LIMIT_HEAP;
 
@@ -164,6 +169,7 @@ impl Dialect for ChiaDialect {
             60 => op_modpow,
             61 => op_mod,
             62 if (flags & ENABLE_KECCAK_OPS_OUTSIDE_GUARD) != 0 => op_keccak256,
+            63 if (flags & ENABLE_SHA256_TREE) != 0 => op_sha256_tree,
             _ => {
                 return unknown_operator(allocator, o, argument_list, flags, max_cost);
             }

src/lib.rs

@@ -15,13 +15,17 @@ pub mod run_program;
 pub mod runtime_dialect;
 pub mod secp_ops;
 pub mod serde;
+pub mod sha_tree_op;
 pub mod traverse_path;
+pub mod treehash;
 
 pub use allocator::{Allocator, Atom, NodePtr, ObjectType, SExp};
 pub use chia_dialect::ChiaDialect;
 pub use run_program::run_program;
 
-pub use chia_dialect::{ENABLE_KECCAK_OPS_OUTSIDE_GUARD, LIMIT_HEAP, MEMPOOL_MODE, NO_UNKNOWN_OPS};
+pub use chia_dialect::{
+    ENABLE_KECCAK_OPS_OUTSIDE_GUARD, ENABLE_SHA256_TREE, LIMIT_HEAP, MEMPOOL_MODE, NO_UNKNOWN_OPS,
+};
 
 #[cfg(feature = "counters")]
 pub use run_program::run_program_with_counters;

src/sha_tree_op.rs

@@ -0,0 +1,11 @@
+use crate::allocator::{Allocator, NodePtr};
+use crate::cost::Cost;
+use crate::op_utils::get_args;
+use crate::reduction::Response;
+use crate::treehash::*;
+
+pub fn op_sha256_tree(a: &mut Allocator, input: NodePtr, max_cost: Cost) -> Response {
+    let [n] = get_args::<1>(a, input, "sha256tree")?;
+    // let mut cache = TreeCache::default();
+    tree_hash_costed(a, n, max_cost)
+}

src/test_ops.rs

@@ -15,6 +15,7 @@ use crate::more_ops::{
 use crate::number::Number;
 use crate::reduction::{Reduction, Response};
 use crate::secp_ops::{op_secp256k1_verify, op_secp256r1_verify};
+use crate::sha_tree_op::op_sha256_tree;
 
 use crate::error::EvalErr;
 use hex::FromHex;
@@ -106,6 +107,7 @@ fn parse_atom(a: &mut Allocator, v: &str) -> NodePtr {
             "secp256k1_verify" => a.new_atom(&[0x13, 0xd6, 0x1f, 0x00]).unwrap(),
             "secp256r1_verify" => a.new_atom(&[0x1c, 0x3a, 0x8f, 0x00]).unwrap(),
             "keccak256" => a.new_atom(&[62]).unwrap(),
+            "sha256tree" => a.new_atom(&[63]).unwrap(),
             _ => {
                 panic!("atom not supported \"{v}\"");
             }
@@ -264,6 +266,8 @@ mod tests {
     #[case("test-secp256r1")]
     #[case("test-modpow")]
     #[case("test-sha256")]
+    #[case("test-sha256tree")]
+    #[case("test-sha256tree-hash")]
     #[case("test-keccak256")]
     #[case("test-keccak256-generated")]
     fn test_ops(#[case] filename: &str) {
@@ -320,6 +324,7 @@ mod tests {
             ("secp256r1_verify", op_secp256r1_verify as Opf),
             ("modpow", op_modpow as Opf),
             ("keccak256", op_keccak256 as Opf),
+            ("sha256tree", op_sha256_tree as Opf),
         ]);
 
         println!("Test cases from: {filename}");

src/treehash.rs

@@ -0,0 +1,180 @@
+use crate::allocator::NodeVisitor;
+use crate::allocator::{Allocator, NodePtr};
+use crate::cost::check_cost;
+use crate::cost::Cost;
+use crate::more_ops::PRECOMPUTED_HASHES;
+use crate::op_utils::MALLOC_COST_PER_BYTE;
+use crate::reduction::Reduction;
+use crate::reduction::Response;
+use chia_sha2::Sha256;
+
+// the base cost is the cost of calling it to begin with
+// this is set to the same as sha256
+const SHA256TREE_BASE_COST: Cost = 87;
+// this cost is applied for every node we traverse to
+const SHA256TREE_NODE_COST: Cost = 500;
+// this is the cost for every 32 bytes in a sha256 call
+// it is set to the same as sha256
+const SHA256TREE_COST_PER_32_BYTES: Cost = 64;
+
+pub fn tree_hash_atom(bytes: &[u8]) -> [u8; 32] {
+    let mut sha256 = Sha256::new();
+    sha256.update([1]);
+    sha256.update(bytes);
+    sha256.finalize()
+}
+
+pub fn tree_hash_pair(first: &[u8; 32], rest: &[u8; 32]) -> [u8; 32] {
+    let mut sha256 = Sha256::new();
+    sha256.update([2]);
+    sha256.update(first);
+    sha256.update(rest);
+    sha256.finalize()
+}
+
+enum TreeOp {
+    SExp(NodePtr),
+    Cons,
+}
+
+// costing is done for every 32 byte chunk that is hashed
+#[inline]
+fn increment_cost_for_hash_of_bytes(size: usize, cost: &mut Cost) {
+    *cost += (size.div_ceil(32)) as u64 * SHA256TREE_COST_PER_32_BYTES;
+}
+
+// this function costs but does not cache
+// we can use it to check that the cache is properly remembering costs
+pub fn tree_hash_costed(a: &mut Allocator, node: NodePtr, cost_left: Cost) -> Response {
+    let mut hashes = Vec::new();
+    let mut ops = vec![TreeOp::SExp(node)];
+
+    let mut cost = SHA256TREE_BASE_COST;
+
+    while let Some(op) = ops.pop() {
+        match op {
+            TreeOp::SExp(node) => {
+                // we could theoretically add a COST_PER_NODE on this line in the future
+                cost += SHA256TREE_NODE_COST;
+                check_cost(cost, cost_left)?;
+                match a.node(node) {
+                    NodeVisitor::Buffer(bytes) => {
+                        // +1 byte to length because of prefix before atoms
+                        increment_cost_for_hash_of_bytes(bytes.len() + 1, &mut cost);
+                        check_cost(cost, cost_left)?;
+                        let hash = tree_hash_atom(bytes);
+                        hashes.push(hash);
+                    }
+                    NodeVisitor::U32(val) => {
+                        // +1 byte to length because of prefix before atoms
+                        increment_cost_for_hash_of_bytes(a.atom_len(node) + 1, &mut cost);
+                        check_cost(cost, cost_left)?;
+                        if (val as usize) < PRECOMPUTED_HASHES.len() {
+                            hashes.push(PRECOMPUTED_HASHES[val as usize]);
+                        } else {
+                            hashes.push(tree_hash_atom(a.atom(node).as_ref()));
+                        }
+                    }
+                    NodeVisitor::Pair(left, right) => {
+                        // 2 * 32byte hashes from a pair
+                        // + 1 byte to length because of prefix before atoms
+                        increment_cost_for_hash_of_bytes(65, &mut cost);
+                        check_cost(cost, cost_left)?;
+
+                        ops.push(TreeOp::Cons);
+                        ops.push(TreeOp::SExp(left));
+                        ops.push(TreeOp::SExp(right));
+                    }
+                }
+            }
+            TreeOp::Cons => {
+                let first = hashes.pop().unwrap();
+                let rest = hashes.pop().unwrap();
+                hashes.push(tree_hash_pair(&first, &rest));
+            }
+        }
+    }
+
+    assert!(hashes.len() == 1);
+    cost += MALLOC_COST_PER_BYTE * 32;
+    check_cost(cost, cost_left)?;
+    Ok(Reduction(cost, a.new_atom(&hashes[0])?))
+}
+
+// this function neither costs, nor caches
+// and it also returns bytes, rather than an Atom
+pub fn tree_hash(a: &Allocator, node: NodePtr) -> [u8; 32] {
+    let mut hashes = Vec::new();
+    let mut ops = vec![TreeOp::SExp(node)];
+
+    while let Some(op) = ops.pop() {
+        match op {
+            TreeOp::SExp(node) => match a.node(node) {
+                NodeVisitor::Buffer(bytes) => {
+                    hashes.push(tree_hash_atom(bytes));
+                }
+                NodeVisitor::U32(val) => {
+                    if (val as usize) < PRECOMPUTED_HASHES.len() {
+                        hashes.push(PRECOMPUTED_HASHES[val as usize]);
+                    } else {
+                        hashes.push(tree_hash_atom(a.atom(node).as_ref()));
+                    }
+                }
+                NodeVisitor::Pair(left, right) => {
+                    ops.push(TreeOp::Cons);
+                    ops.push(TreeOp::SExp(left));
+                    ops.push(TreeOp::SExp(right));
+                }
+            },
+            TreeOp::Cons => {
+                let first = hashes.pop().unwrap();
+                let rest = hashes.pop().unwrap();
+                hashes.push(tree_hash_pair(&first, &rest));
+            }
+        }
+    }
+
+    assert!(hashes.len() == 1);
+    hashes[0]
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn test_sha256_atom(buf: &[u8]) {
+        let hash = tree_hash_atom(buf);
+
+        let mut hasher = Sha256::new();
+        hasher.update([1_u8]);
+        if !buf.is_empty() {
+            hasher.update(buf);
+        }
+
+        assert_eq!(hash.as_ref(), hasher.finalize().as_slice());
+    }
+
+    #[test]
+    fn test_tree_hash_atom() {
+        test_sha256_atom(&[]);
+        for val in 0..=255 {
+            test_sha256_atom(&[val]);
+        }
+
+        for val in 0..=255 {
+            test_sha256_atom(&[0, val]);
+        }
+
+        for val in 0..=255 {
+            test_sha256_atom(&[0xff, val]);
+        }
+    }
+
+    #[test]
+    fn test_precomputed_atoms() {
+        assert_eq!(tree_hash_atom(&[]), PRECOMPUTED_HASHES[0]);
+        for val in 1..(PRECOMPUTED_HASHES.len() as u8) {
+            assert_eq!(tree_hash_atom(&[val]), PRECOMPUTED_HASHES[val as usize]);
+        }
+    }
+}

tools/generate-keccak-tests.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from eth_hash.auto import keccak
 from random import randbytes, randint, seed, sample
 from more_itertools import sliced
@@ -9,7 +10,8 @@
 
 SIZE = 100
 
-with open("../op-tests/test-keccak256-generated.txt", "w+") as f:
+p = Path(__file__).parent.parent / "op-tests/test-keccak256-generated.txt"
+with open(p, "w+") as f:
     f.write("; This file was generated by tools/generate-keccak-tests.py\n\n")
 
     for i in range(SIZE):