diff --git a/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/node_store/cached_content_addressed_tree_store.hpp b/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/node_store/cached_content_addressed_tree_store.hpp index 9b113365e1a..f91fd9bf775 100644 --- a/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/node_store/cached_content_addressed_tree_store.hpp +++ b/barretenberg/cpp/src/barretenberg/crypto/merkle_tree/node_store/cached_content_addressed_tree_store.hpp @@ -19,6 +19,7 @@ #include #include #include +#include template <> struct std::hash { std::size_t operator()(const uint256_t& k) const { return k.data[0]; } @@ -729,36 +730,48 @@ void ContentAddressedCachedTreeStore::persist_node(const std::opt uint32_t level, WriteTransaction& tx) { - // If the optional hash does not have a value then it means it's the zero tree value at this level - // If it has a value but that value is not in our stores then it means it is referencing a node - // created in a previous block, so that will need to have it's reference count increased - if (!optional_hash.has_value()) { - return; - } - fr hash = optional_hash.value(); + struct StackObject { + std::optional opHash; + uint32_t lvl; + }; + std::vector stack; + stack.push_back({ .opHash = optional_hash, .lvl = level }); + + while (!stack.empty()) { + StackObject so = stack.back(); + stack.pop_back(); + + // If the optional hash does not have a value then it means it's the zero tree value at this level + // If it has a value but that value is not in our stores then it means it is referencing a node + // created in a previous block, so that will need to have it's reference count increased + if (!so.opHash.has_value()) { + continue; + } + fr hash = so.opHash.value(); - if (level == depth_) { - // this is a leaf - persist_leaf_pre_image(hash, tx); - } + if (so.lvl == depth_) { + // this is a leaf + persist_leaf_pre_image(hash, tx); + } - // std::cout << "Persisting node hash " << hash << " at level " << level << std::endl; + // std::cout << "Persisting node hash " << hash << " at level " << so.lvl << std::endl; + auto nodePayloadIter = nodes_.find(hash); + if (nodePayloadIter == nodes_.end()) { + // need to increase the stored node's reference count here + dataStore_->increment_node_reference_count(hash, tx); + continue; + } - auto nodePayloadIter = nodes_.find(hash); - if (nodePayloadIter == nodes_.end()) { - // need to increase the stored node's reference count here - dataStore_->increment_node_reference_count(hash, tx); - return; - } - NodePayload nodeData = nodePayloadIter->second; - dataStore_->set_or_increment_node_reference_count(hash, nodeData, tx); - if (nodeData.ref != 1) { - // If the node now has a ref count greater then 1, we don't continue. - // It means that the entire sub-tree underneath already exists - return; + NodePayload nodeData = nodePayloadIter->second; + dataStore_->set_or_increment_node_reference_count(hash, nodeData, tx); + if (nodeData.ref != 1) { + // If the node now has a ref count greater then 1, we don't continue. + // It means that the entire sub-tree underneath already exists + continue; + } + stack.push_back({ .opHash = nodePayloadIter->second.left, .lvl = so.lvl + 1 }); + stack.push_back({ .opHash = nodePayloadIter->second.right, .lvl = so.lvl + 1 }); } - persist_node(nodePayloadIter->second.left, level + 1, tx); - persist_node(nodePayloadIter->second.right, level + 1, tx); } template @@ -1048,29 +1061,38 @@ void ContentAddressedCachedTreeStore::remove_node(const std::opti std::optional maxIndex, WriteTransaction& tx) { - if (!optional_hash.has_value()) { - return; - } - fr hash = optional_hash.value(); - - // we need to retrieve the node and decrement it's reference count - // std::cout << "Decrementing ref count for node " << hash << ", level " << level << std::endl; - NodePayload nodeData; - dataStore_->decrement_node_reference_count(hash, nodeData, tx); - - if (nodeData.ref != 0) { - // node was not deleted, we don't continue the search - return; - } - - // the node was deleted, if it was a leaf then we need to remove the pre-image - if (level == depth_) { - remove_leaf(hash, maxIndex, tx); + struct StackObject { + std::optional opHash; + uint32_t lvl; + }; + std::vector stack; + stack.push_back({ .opHash = optional_hash, .lvl = level }); + + while (!stack.empty()) { + StackObject so = stack.back(); + stack.pop_back(); + + if (!so.opHash.has_value()) { + continue; + } + fr hash = so.opHash.value(); + // we need to retrieve the node and decrement it's reference count + // std::cout << "Decrementing ref count for node " << hash << ", level " << so.lvl << std::endl; + NodePayload nodeData; + dataStore_->decrement_node_reference_count(hash, nodeData, tx); + + if (nodeData.ref != 0) { + // node was not deleted, we don't continue the search + continue; + } + // the node was deleted, if it was a leaf then we need to remove the pre-image + if (so.lvl == depth_) { + remove_leaf(hash, maxIndex, tx); + } + // push the child nodes to the stack + stack.push_back({ .opHash = std::optional(nodeData.left), .lvl = so.lvl + 1 }); + stack.push_back({ .opHash = std::optional(nodeData.right), .lvl = so.lvl + 1 }); } - - // now recursively remove the next level - remove_node(std::optional(nodeData.left), level + 1, maxIndex, tx); - remove_node(std::optional(nodeData.right), level + 1, maxIndex, tx); } template void ContentAddressedCachedTreeStore::initialise()