membership.hpp

#pragma once
#include "barretenberg/stdlib/hash/pedersen/pedersen.hpp"
#include "barretenberg/stdlib/primitives/byte_array/byte_array.hpp"
#include "barretenberg/stdlib/primitives/field/field.hpp"
#include "hash_path.hpp"
 
namespace proof_system::plonk {
namespace stdlib {
namespace merkle_tree {
 
template <typename Builder> using bit_vector = std::vector<bool_t<Builder>>;
template <typename Builder>
field_t<Builder> compute_subtree_root(hash_path<Builder> const& hashes,
                                      field_t<Builder> const& value,
                                      bit_vector<Builder> const& index,
                                      size_t at_height,
                                      bool const is_updating_tree = false)
{
    auto current = value;
    for (size_t i = at_height; i < hashes.size(); ++i) {
        // get the parity bit at this level of the tree (get_bit returns bool so we know this is 0 or 1)
        bool_t<Builder> path_bit = index[i];
 
        // reconstruct the two inputs we need to hash
        // if `path_bit = false`, we know `current` is the left leaf and `hashes[i].second` is the right leaf
        // if `path_bit = true`, we know `current` is the right leaf and `hashes[i].first` is the left leaf
        // We don't need to explicitly check that hashes[i].first = current iff !path bit , or that hashes[i].second =
        // current iff path_bit If either of these does not hold, then the final computed merkle root will not match
        field_t<Builder> left = field_t<Builder>::conditional_assign(path_bit, hashes[i].first, current);
        field_t<Builder> right = field_t<Builder>::conditional_assign(path_bit, current, hashes[i].second);
        if (is_updating_tree) {
            current = pedersen_hash<Builder>::hash({ left, right }, 0);
        } else {
            current = pedersen_hash<Builder>::hash_skip_field_validation({ left, right }, 0);
        }
    }
 
    return current;
}
 
template <typename Builder>
bool_t<Builder> check_subtree_membership(field_t<Builder> const& root,
                                         hash_path<Builder> const& hashes,
                                         field_t<Builder> const& value,
                                         bit_vector<Builder> const& index,
                                         size_t at_height,
                                         bool const is_updating_tree = false)
{
    return (compute_subtree_root(hashes, value, index, at_height, is_updating_tree) == root);
}
 
template <typename Builder>
void assert_check_subtree_membership(field_t<Builder> const& root,
                                     hash_path<Builder> const& hashes,
                                     field_t<Builder> const& value,
                                     bit_vector<Builder> const& index,
                                     size_t at_height,
                                     bool const is_updating_tree = false,
                                     std::string const& msg = "assert_check_subtree_membership")
{
    auto exists = check_subtree_membership(root, hashes, value, index, at_height, is_updating_tree);
    exists.assert_equal(true, msg);
}
 
template <typename Builder>
bool_t<Builder> check_membership(field_t<Builder> const& root,
                                 hash_path<Builder> const& hashes,
                                 field_t<Builder> const& value,
                                 bit_vector<Builder> const& index,
                                 bool const is_updating_tree = false)
{
    return check_subtree_membership(root, hashes, value, index, 0, is_updating_tree);
}
 
template <typename Builder>
void assert_check_membership(field_t<Builder> const& root,
                             hash_path<Builder> const& hashes,
                             field_t<Builder> const& value,
                             bit_vector<Builder> const& index,
                             bool const is_updating_tree = false,
                             std::string const& msg = "assert_check_membership")
{
    auto exists = stdlib::merkle_tree::check_membership(root, hashes, value, index, is_updating_tree);
    exists.assert_equal(true, msg);
}
 
template <typename Builder>
void update_membership(field_t<Builder> const& new_root,
                       field_t<Builder> const& new_value,
                       field_t<Builder> const& old_root,
                       hash_path<Builder> const& old_hashes,
                       field_t<Builder> const& old_value,
                       bit_vector<Builder> const& index,
                       std::string const& msg = "update_membership")
{
    // Check that the old_value, is in the tree given by old_root, at index.
    assert_check_membership(old_root, old_hashes, old_value, index, false, msg + "_old_value");
 
    // Check that the new_value, is in the tree given by new_root, at index.
    assert_check_membership(new_root, old_hashes, new_value, index, true, msg + "_new_value");
}
 
template <typename Builder>
field_t<Builder> update_memberships(field_t<Builder> old_root,
                                    std::vector<field_t<Builder>> const& new_roots,
                                    std::vector<field_t<Builder>> const& new_values,
                                    std::vector<field_t<Builder>> const& old_values,
                                    std::vector<hash_path<Builder>> const& old_paths,
                                    std::vector<bit_vector<Builder>> const& old_indicies)
{
    for (size_t i = 0; i < old_indicies.size(); i++) {
        update_membership(
            new_roots[i], new_values[i], old_root, old_paths[i], old_values[i], old_indicies[i], "update_memberships");
 
        old_root = new_roots[i];
    }
 
    return old_root;
}
 
template <typename Builder>
void update_subtree_membership(field_t<Builder> const& new_root,
                               field_t<Builder> const& new_subtree_root,
                               field_t<Builder> const& old_root,
                               hash_path<Builder> const& old_hashes,
                               field_t<Builder> const& old_subtree_root,
                               bit_vector<Builder> const& index,
                               size_t at_height,
                               std::string const& msg = "update_subtree_membership")
{
    // Check that the old_subtree_root, is in the tree given by old_root, at index and at_height.
    assert_check_subtree_membership(
        old_root, old_hashes, old_subtree_root, index, at_height, false, msg + "_old_subtree");
 
    // Check that the new_subtree_root, is in the tree given by new_root, at index and at_height.
    // By extracting partner hashes from `old_hashes`, we also validate both membership proofs use
    // identical merkle trees (apart from the leaf that is being updated)
    assert_check_subtree_membership(
        new_root, old_hashes, new_subtree_root, index, at_height, true, msg + "_new_subtree");
}
 
template <typename Builder> field_t<Builder> compute_tree_root(std::vector<field_t<Builder>> const& input)
{
    // Check if the input vector size is a power of 2.
    ASSERT(input.size() > 0);
    ASSERT(!(input.size() & (input.size() - 1)) == true);
    auto layer = input;
    while (layer.size() > 1) {
        std::vector<field_t<Builder>> next_layer(layer.size() / 2);
        for (size_t i = 0; i < next_layer.size(); ++i) {
            next_layer[i] = pedersen_hash<Builder>::hash({ layer[i * 2], layer[i * 2 + 1] });
        }
        layer = std::move(next_layer);
    }
 
    return layer[0];
}
 
template <typename Builder>
bool_t<Builder> check_tree(field_t<Builder> const& root, std::vector<field_t<Builder>> const& values)
{
    return compute_tree_root(values) == root;
}
 
template <typename Builder>
void assert_check_tree(field_t<Builder> const& root, std::vector<field_t<Builder>> const& values)
{
    auto valid = check_tree(root, values);
    valid.assert_equal(true, "assert_check_tree");
}
 
template <typename Builder>
void batch_update_membership(field_t<Builder> const& new_root,
                             field_t<Builder> const& old_root,
                             hash_path<Builder> const& old_path,
                             std::vector<field_t<Builder>> const& new_values,
                             field_t<Builder> const& start_index,
                             std::string const& msg = "batch_update_membership")
{
    size_t height = numeric::get_msb(new_values.size());
    auto zero_subtree_root = field_t<Builder>(zero_hash_at_height(height));
 
    auto rollup_root = compute_tree_root(new_values);
 
    update_subtree_membership(
        new_root, rollup_root, old_root, old_path, zero_subtree_root, start_index.decompose_into_bits(), height, msg);
}
 
} // namespace merkle_tree
} // namespace stdlib
} // namespace proof_system::plonk