public_data_tree_trace.cpp

Go to the documentation of this file.

#include "barretenberg/vm2/tracegen/public_data_tree_trace.hpp"
 
#include <algorithm>
#include <cstdint>
#include <limits>
#include <unordered_map>
#include <vector>
 
#include "barretenberg/aztec/aztec_constants.hpp"
#include "barretenberg/common/assert.hpp"
#include "barretenberg/vm2/common/field.hpp"
#include "barretenberg/vm2/generated/columns.hpp"
#include "barretenberg/vm2/generated/relations/lookups_public_data_check.hpp"
#include "barretenberg/vm2/generated/relations/lookups_public_data_squash.hpp"
#include "barretenberg/vm2/generated/relations/perms_public_data_check.hpp"
#include "barretenberg/vm2/generated/relations/perms_sstore.hpp"
#include "barretenberg/vm2/generated/relations/perms_tx.hpp"
#include "barretenberg/vm2/simulation/events/public_data_tree_check_event.hpp"
#include "barretenberg/vm2/simulation/lib/db_types.hpp"
#include "barretenberg/vm2/tracegen/lib/discard_reconstruction.hpp"
#include "barretenberg/vm2/tracegen/lib/interaction_def.hpp"
 
// This tracegen file will generate both the public data check trace and the public data squash trace from the same
// events. These two traces are sorted in a different order, so we need to process them separately.
namespace bb::avm2::tracegen {
 
using simulation::PublicDataTreeLeafPreimage;
using simulation::PublicDataTreeReadWriteEvent;
 
namespace {
 
struct EventWithDiscard {
    simulation::PublicDataTreeReadWriteEvent event;
    bool discard;
};
 
void process_public_data_tree_check_trace(const std::vector<EventWithDiscard>& events_with_discard,
                                          const std::unordered_map<FF, uint32_t>& first_write_per_slot,
                                          const std::unordered_map<FF, FF>& last_value_per_slot,
                                          TraceContainer& trace)
{
    using C = Column;
 
    // This is a shifted trace, so we start at 1
    uint32_t row = 1;
 
    uint32_t write_idx = AVM_PUBLIC_INPUTS_AVM_ACCUMULATED_DATA_PUBLIC_DATA_WRITES_ROW_IDX;
 
    for (size_t i = 0; i < events_with_discard.size(); i++) {
        bool end = i == events_with_discard.size() - 1;
        const auto& event = events_with_discard[i].event;
        const bool discard = events_with_discard[i].discard;
 
        bool exists = event.low_leaf_preimage.leaf.slot == event.leaf_slot;
        FF slot_low_leaf_slot_diff = event.leaf_slot - event.low_leaf_preimage.leaf.slot;
 
        bool next_slot_is_nonzero = false;
        FF next_slot = 0;
        if (!exists) {
            next_slot_is_nonzero = event.low_leaf_preimage.nextKey != 0;
            next_slot = event.low_leaf_preimage.nextKey;
        }
 
        bool write = event.write_data.has_value();
        // Note: the protocol_write and non_protocol_write are set by the MultiPermutationBuilder, but I also set them
        // here for clarity/testing. Can remove, but would need to change the NegativeSetProtocolWrite test.
        bool protocol_write = event.execution_id == std::numeric_limits<uint32_t>::max();
        bool should_insert = !exists && write;
        bool nondiscarded_write = write && !discard;
        bool sel_write_to_public_inputs =
            nondiscarded_write && first_write_per_slot.at(event.leaf_slot) == event.execution_id;
        FF final_value = nondiscarded_write ? last_value_per_slot.at(event.leaf_slot) : 0;
 
        FF intermediate_root = 0;
        PublicDataTreeLeafPreimage updated_low_leaf = PublicDataTreeLeafPreimage::empty();
        FF updated_low_leaf_hash = 0;
        FF new_leaf_hash = 0;
        AppendOnlyTreeSnapshot next_snapshot = event.prev_snapshot;
 
        if (write) {
            next_snapshot = event.write_data->next_snapshot;
            intermediate_root = event.write_data->intermediate_root;
            updated_low_leaf = event.write_data->updated_low_leaf_preimage;
            updated_low_leaf_hash = event.write_data->updated_low_leaf_hash;
            new_leaf_hash = event.write_data->new_leaf_hash;
        }
        uint32_t clk = event.execution_id;
        uint32_t clk_diff = end ? 0 : events_with_discard[i + 1].event.execution_id - clk;
 
        uint32_t public_data_writes_length = write_idx -
                                             AVM_PUBLIC_INPUTS_AVM_ACCUMULATED_DATA_PUBLIC_DATA_WRITES_ROW_IDX +
                                             static_cast<uint32_t>(sel_write_to_public_inputs);
 
        trace.set(row,
                  { {
                      { C::public_data_check_sel, 1 },
                      { C::public_data_check_not_end, !end },
                      { C::public_data_check_end, end },
                      { C::public_data_check_value, event.value },
                      { C::public_data_check_slot, event.slot },
                      { C::public_data_check_root, event.prev_snapshot.root },
                      { C::public_data_check_address, event.contract_address },
                      { C::public_data_check_write_root, next_snapshot.root },
                      { C::public_data_check_tree_size_before_write, event.prev_snapshot.next_available_leaf_index },
                      { C::public_data_check_tree_size_after_write, next_snapshot.next_available_leaf_index },
                      { C::public_data_check_write, write },
                      { C::public_data_check_protocol_write, protocol_write },
                      { C::public_data_check_non_protocol_write, write && !protocol_write },
                      { C::public_data_check_clk, clk },
                      { C::public_data_check_discard, discard },
                      { C::public_data_check_low_leaf_slot, event.low_leaf_preimage.leaf.slot },
                      { C::public_data_check_low_leaf_value, event.low_leaf_preimage.leaf.value },
                      { C::public_data_check_low_leaf_next_index, event.low_leaf_preimage.nextIndex },
                      { C::public_data_check_low_leaf_next_slot, event.low_leaf_preimage.nextKey },
                      { C::public_data_check_updated_low_leaf_value, updated_low_leaf.leaf.value },
                      { C::public_data_check_updated_low_leaf_next_index, updated_low_leaf.nextIndex },
                      { C::public_data_check_updated_low_leaf_next_slot, updated_low_leaf.nextKey },
                      { C::public_data_check_low_leaf_index, event.low_leaf_index },
                      { C::public_data_check_clk_diff_lo, static_cast<uint16_t>(clk_diff) },
                      { C::public_data_check_clk_diff_hi, clk_diff >> 16 },
                      { C::public_data_check_leaf_slot, event.leaf_slot },
                      { C::public_data_check_siloing_separator, DOM_SEP__PUBLIC_LEAF_SLOT },
                      { C::public_data_check_leaf_not_exists, !exists },
                      { C::public_data_check_leaf_slot_low_leaf_slot_diff_inv,
                        slot_low_leaf_slot_diff }, // Will be inverted in batch later
                      { C::public_data_check_next_slot_is_nonzero, next_slot_is_nonzero },
                      { C::public_data_check_next_slot_inv, next_slot }, // Will be inverted in batch later
                      { C::public_data_check_low_leaf_hash, event.low_leaf_hash },
                      { C::public_data_check_intermediate_root, intermediate_root },
                      { C::public_data_check_tree_height, PUBLIC_DATA_TREE_HEIGHT },
                      { C::public_data_check_const_three, 3 },
                      { C::public_data_check_const_four, 4 },
                      { C::public_data_check_merkle_hash_separator, DOM_SEP__PUBLIC_DATA_MERKLE },
                      { C::public_data_check_updated_low_leaf_hash, updated_low_leaf_hash },
                      { C::public_data_check_should_insert, should_insert },
                      { C::public_data_check_new_leaf_hash, new_leaf_hash },
                      { C::public_data_check_write_idx, write_idx },
                      { C::public_data_check_non_discarded_write, nondiscarded_write },
                      { C::public_data_check_sel_write_to_public_inputs, sel_write_to_public_inputs },
                      { C::public_data_check_final_value, final_value },
                      { C::public_data_check_public_data_writes_length, public_data_writes_length },
                      { C::public_data_check_length_pi_idx,
                        AVM_PUBLIC_INPUTS_AVM_ACCUMULATED_DATA_ARRAY_LENGTHS_PUBLIC_DATA_WRITES_ROW_IDX },
                  } });
        row++;
        if (sel_write_to_public_inputs) {
            write_idx++;
        }
    }
}
 
void process_squashing_trace(const std::vector<PublicDataTreeReadWriteEvent>& nondiscarded_writes,
                             const std::unordered_map<FF, uint32_t>& first_write_per_slot,
                             const std::unordered_map<FF, FF>& last_value_per_slot,
                             TraceContainer& trace)
{
    using C = Column;
 
    // This is a shifted trace, so we start at 1
    uint32_t row = 1;
 
    for (size_t i = 0; i < nondiscarded_writes.size(); i++) {
        bool end = i == nondiscarded_writes.size() - 1;
        const auto& event = nondiscarded_writes[i];
 
        uint32_t clk = event.execution_id;
 
        bool leaf_slot_increase = false;
        bool check_clock = false;
        uint32_t clk_diff = 0;
 
        if (!end) {
            const auto& next_event = nondiscarded_writes[i + 1];
 
            if (event.leaf_slot == next_event.leaf_slot) {
                BB_ASSERT_LT(
                    event.execution_id, next_event.execution_id, "Execution id is not less than next execution id");
                // We check strictly increasing clk diff in the squash trace.
                clk_diff = next_event.execution_id - event.execution_id - 1;
                check_clock = true;
            } else {
                BB_ASSERT_LT(static_cast<uint256_t>(event.leaf_slot),
                             static_cast<uint256_t>(next_event.leaf_slot),
                             "Leaf slot is not less than next leaf slot");
                leaf_slot_increase = true;
            }
        }
 
        bool sel_write_to_public_inputs = first_write_per_slot.at(event.leaf_slot) == event.execution_id;
        FF final_value = last_value_per_slot.at(event.leaf_slot);
 
        trace.set(row,
                  { {
                      { C::public_data_squash_sel, 1 },
                      { C::public_data_squash_leaf_slot, event.leaf_slot },
                      { C::public_data_squash_value, event.value },
                      { C::public_data_squash_clk, clk },
                      { C::public_data_squash_write_to_public_inputs, sel_write_to_public_inputs },
                      { C::public_data_squash_leaf_slot_increase, leaf_slot_increase },
                      { C::public_data_squash_check_clock, check_clock },
                      { C::public_data_squash_clk_diff_lo, static_cast<uint16_t>(clk_diff) },
                      { C::public_data_squash_clk_diff_hi, clk_diff >> 16 },
                      { C::public_data_squash_final_value, final_value },
                  } });
        row++;
    }
}
 
} // namespace
 
void PublicDataTreeTraceBuilder::process(
    const simulation::EventEmitterInterface<simulation::PublicDataTreeCheckEvent>::Container& events,
    TraceContainer& trace)
{
    std::vector<EventWithDiscard> events_with_discard;
    std::unordered_map<FF, uint32_t> first_write_clk_per_slot;
    std::unordered_map<FF, FF> last_write_value_per_slot;
 
    events_with_discard.reserve(events.size());
    // Initial order of events will be chronological order.
    process_with_discard(events, [&](const simulation::PublicDataTreeReadWriteEvent& event, bool discard) {
        events_with_discard.push_back({ event, discard });
        if (!discard && event.write_data.has_value()) {
            if (!first_write_clk_per_slot.contains(event.leaf_slot)) {
                first_write_clk_per_slot[event.leaf_slot] = event.execution_id;
            }
            last_write_value_per_slot[event.leaf_slot] = event.value;
        }
    });
 
    // Sort by clk in ascending order (reads will come with clk=0)
    std::ranges::sort(events_with_discard.begin(),
                      events_with_discard.end(),
                      [](const EventWithDiscard& a, const EventWithDiscard& b) {
                          return a.event.execution_id < b.event.execution_id;
                      });
 
    process_public_data_tree_check_trace(
        events_with_discard, first_write_clk_per_slot, last_write_value_per_slot, trace);
 
    std::vector<PublicDataTreeReadWriteEvent> nondiscarded_writes;
    nondiscarded_writes.reserve(events_with_discard.size());
    // Use only nondiscarded writes
    for (const auto& event_with_discard : events_with_discard) {
        if (event_with_discard.event.write_data.has_value() && !event_with_discard.discard) {
            nondiscarded_writes.push_back(event_with_discard.event);
        }
    }
 
    // Sort by slot, and then by clk
    std::ranges::sort(nondiscarded_writes.begin(),
                      nondiscarded_writes.end(),
                      [](const PublicDataTreeReadWriteEvent& a, const PublicDataTreeReadWriteEvent& b) {
                          if (a.leaf_slot == b.leaf_slot) {
                              return a.execution_id < b.execution_id;
                          }
                          return static_cast<uint256_t>(a.leaf_slot) < static_cast<uint256_t>(b.leaf_slot);
                      });
 
    process_squashing_trace(nondiscarded_writes, first_write_clk_per_slot, last_write_value_per_slot, trace);
 
    // Batch invert the columns.
    trace.invert_columns(
        { { Column::public_data_check_leaf_slot_low_leaf_slot_diff_inv, Column::public_data_check_next_slot_inv } });
}
 
// This is a sorted trace, so we can't use sequential lookups. Instead, we'll use generic lookups and clk indexed ones
// for the ones targeting precomputed or public inputs.
const InteractionDefinition PublicDataTreeTraceBuilder::interactions =
    InteractionDefinition()
        // Public data read/write
        .add<InteractionType::LookupGeneric, lookup_public_data_check_silo_poseidon2_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_low_leaf_slot_validation_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_low_leaf_next_slot_validation_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_low_leaf_poseidon2_0_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_low_leaf_poseidon2_1_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_updated_low_leaf_poseidon2_0_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_updated_low_leaf_poseidon2_1_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_low_leaf_merkle_check_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_new_leaf_poseidon2_0_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_new_leaf_poseidon2_1_settings>()
        .add<InteractionType::LookupGeneric, lookup_public_data_check_new_leaf_merkle_check_settings>()
        .add<InteractionType::MultiPermutation, perm_sstore_storage_write_settings, perm_tx_balance_update_settings>(
            Column::public_data_check_write)
        .add<InteractionType::Permutation, perm_public_data_check_squashing_settings>()
        .add<InteractionType::LookupIntoIndexedByRow,
             lookup_public_data_check_write_writes_length_to_public_inputs_settings>()
        .add<InteractionType::LookupIntoIndexedByRow, lookup_public_data_check_clk_diff_range_lo_settings>()
        .add<InteractionType::LookupIntoIndexedByRow, lookup_public_data_check_clk_diff_range_hi_settings>()
        .add<InteractionType::LookupIntoIndexedByRow,
             lookup_public_data_check_write_public_data_to_public_inputs_settings>()
        // Public data squash
        .add<InteractionType::LookupGeneric, lookup_public_data_squash_leaf_slot_increase_ff_gt_settings>()
        .add<InteractionType::LookupIntoIndexedByRow, lookup_public_data_squash_clk_diff_range_lo_settings>()
        .add<InteractionType::LookupIntoIndexedByRow, lookup_public_data_squash_clk_diff_range_hi_settings>();
 
} // namespace bb::avm2::tracegen