feat: Update RAM/ROM memory records for new block structure (#4806)

The 4th wire of RAM/ROM read/write gates is generated at proving time as
a linear combination of the first three wires scaled by powers of a
challenge. To know on which rows to perform this calculation, we must
store the indices of read/write gates in the proving key.

Previously, these indices were set in the builder based on num_gates,
then offset later in the composer to account for the fact that public
inputs and/or ecc op gates are placed at the top of the trace. This PR
updates this model to work with the new block structure in the builders.
In particular, in the builder we set the gate index of ram/rom
read/writes based on the row within the block. (Currently ram/rom gates
are added to `main`, eventually they'll be in `aux`). Then, we get the
correct execution trace row index by later offsetting these values based
on the offset at which the black containing these gates is placed in the
trace.

The method for performing this offset has been moved to
execution_trace.hpp since it fundamentally depends on the ordering of
the execution trace.

Note: This is one step towards getting RAM/ROM to work with the new
block structure. It will also be necessary to remove the interleaving of
arithmetic gates into read/write gates. This will come in a follow on.

barretenberg/cpp/src/barretenberg/flavor/flavor.hpp

@@ -302,6 +302,9 @@ concept IsPlonkFlavor = IsAnyOf<T, plonk::flavor::Standard, plonk::flavor::Ultra
 template <typename T>
 concept IsUltraPlonkFlavor = IsAnyOf<T, plonk::flavor::Ultra>;
 
+template <typename T> 
+concept IsUltraPlonkOrHonk = IsAnyOf<T, plonk::flavor::Ultra, UltraFlavor, GoblinUltraFlavor>;
+
 template <typename T> 
 concept IsHonkFlavor = IsAnyOf<T, UltraFlavor, GoblinUltraFlavor>;
 

barretenberg/cpp/src/barretenberg/plonk/composer/ultra_composer.cpp

@@ -188,8 +188,6 @@ std::shared_ptr<proving_key> UltraComposer::compute_proving_key(CircuitBuilder&
 
     construct_sorted_polynomials(circuit, subgroup_size);
 
-    populate_memory_read_write_records<Flavor>(circuit, circuit_proving_key);
-
     return circuit_proving_key;
 }
 

barretenberg/cpp/src/barretenberg/proof_system/arithmetization/arithmetization.hpp

@@ -48,9 +48,12 @@ template <typename FF, size_t NUM_WIRES, size_t NUM_SELECTORS> class ExecutionTr
 
     Wires wires; // vectors of indices into a witness variables array
     Selectors selectors;
+    bool has_ram_rom = false; // does the block contain RAM/ROM gates
 
     bool operator==(const ExecutionTraceBlock& other) const = default;
 
+    size_t size() { return std::get<0>(this->wires).size(); }
+
     void reserve(size_t size_hint)
     {
         for (auto& w : wires) {
@@ -139,9 +142,17 @@ template <typename FF_> class UltraArith {
 
     struct TraceBlocks {
         UltraTraceBlock pub_inputs;
+        UltraTraceBlock arithmetic;
+        UltraTraceBlock sort;
+        UltraTraceBlock elliptic;
+        UltraTraceBlock aux;
+        UltraTraceBlock lookup;
         UltraTraceBlock main;
 
-        auto get() { return RefArray{ pub_inputs, main }; }
+        // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): update to aux.has_ram_rom = true
+        TraceBlocks() { main.has_ram_rom = true; }
+
+        auto get() { return RefArray{ pub_inputs, arithmetic, sort, elliptic, aux, lookup, main }; }
 
         bool operator==(const TraceBlocks& other) const = default;
     };
@@ -224,9 +235,23 @@ template <typename FF_> class UltraHonkArith {
     struct TraceBlocks {
         UltraHonkTraceBlock ecc_op;
         UltraHonkTraceBlock pub_inputs;
+        UltraHonkTraceBlock arithmetic;
+        UltraHonkTraceBlock sort;
+        UltraHonkTraceBlock elliptic;
+        UltraHonkTraceBlock aux;
+        UltraHonkTraceBlock lookup;
+        UltraHonkTraceBlock busread;
+        UltraHonkTraceBlock poseidon_external;
+        UltraHonkTraceBlock poseidon_internal;
         UltraHonkTraceBlock main;
 
-        auto get() { return RefArray{ ecc_op, pub_inputs, main }; }
+        TraceBlocks() { main.has_ram_rom = true; }
+
+        auto get()
+        {
+            return RefArray{ ecc_op,  pub_inputs,        arithmetic,        sort, elliptic, aux, lookup,
+                             busread, poseidon_external, poseidon_internal, main };
+        }
 
         bool operator==(const TraceBlocks& other) const = default;
     };

barretenberg/cpp/src/barretenberg/proof_system/circuit_builder/ultra_circuit_builder.cpp

@@ -2033,7 +2033,8 @@ template <typename Arithmetization> void UltraCircuitBuilder_<Arithmetization>::
     blocks.main.populate_wires(
         record.index_witness, record.value_column1_witness, record.value_column2_witness, record.record_witness);
 
-    record.gate_index = this->num_gates;
+    // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): set gate index based on block containing ram/rom
+    record.gate_index = this->blocks.main.size() - 1;
     ++this->num_gates;
 }
 
@@ -2052,7 +2053,8 @@ void UltraCircuitBuilder_<Arithmetization>::create_sorted_ROM_gate(RomRecord& re
     blocks.main.populate_wires(
         record.index_witness, record.value_column1_witness, record.value_column2_witness, record.record_witness);
 
-    record.gate_index = this->num_gates;
+    // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): set gate index based on block containing ram/rom
+    record.gate_index = this->blocks.main.size() - 1;
     ++this->num_gates;
 }
 
@@ -2097,7 +2099,8 @@ template <typename Arithmetization> void UltraCircuitBuilder_<Arithmetization>::
     blocks.main.populate_wires(
         record.index_witness, record.timestamp_witness, record.value_witness, record.record_witness);
 
-    record.gate_index = this->num_gates;
+    // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): set gate index based on block containing ram/rom
+    record.gate_index = this->blocks.main.size() - 1;
     ++this->num_gates;
 }
 
@@ -2117,7 +2120,8 @@ void UltraCircuitBuilder_<Arithmetization>::create_sorted_RAM_gate(RamRecord& re
     blocks.main.populate_wires(
         record.index_witness, record.timestamp_witness, record.value_witness, record.record_witness);
 
-    record.gate_index = this->num_gates;
+    // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): set gate index based on block containing ram/rom
+    record.gate_index = this->blocks.main.size() - 1;
     ++this->num_gates;
 }
 
@@ -2131,7 +2135,8 @@ template <typename Arithmetization>
 void UltraCircuitBuilder_<Arithmetization>::create_final_sorted_RAM_gate(RamRecord& record, const size_t ram_array_size)
 {
     record.record_witness = this->add_variable(0);
-    record.gate_index = this->num_gates;
+    // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): set gate index based on block containing ram/rom
+    record.gate_index = this->blocks.main.size(); // no -1 since we havent added the gate yet
 
     // TODO(https://github.com/AztecProtocol/barretenberg/issues/879): This method used to add a single arithmetic gate
     // with two purposes: (1) to provide wire values to the previous RAM gate via shifts, and (2) to perform a
@@ -3323,7 +3328,8 @@ template <typename Arithmetization> bool UltraCircuitBuilder_<Arithmetization>::
         }
     };
     // For each gate
-    for (size_t i = 0; i < this->num_gates; i++) {
+    // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): only checking the main block. check all blocks
+    for (size_t i = 0; i < this->blocks.main.size(); i++) {
         FF q_arith_value;
         FF q_aux_value;
         FF q_elliptic_value;
@@ -3375,7 +3381,8 @@ template <typename Arithmetization> bool UltraCircuitBuilder_<Arithmetization>::
         FF w_2_shifted_value;
         FF w_3_shifted_value;
         FF w_4_shifted_value;
-        if (i < (this->num_gates - 1)) {
+        // TODO(https://github.com/AztecProtocol/barretenberg/issues/867): gate index based on block containing ram/rom
+        if (i < (this->blocks.main.size() - 1)) {
             w_1_shifted_value = this->get_variable(blocks.main.w_l()[i + 1]);
             w_2_shifted_value = this->get_variable(blocks.main.w_r()[i + 1]);
             w_3_shifted_value = this->get_variable(blocks.main.w_o()[i + 1]);

barretenberg/cpp/src/barretenberg/proof_system/circuit_builder/ultra_circuit_builder.test.cpp

@@ -712,6 +712,40 @@ TEST(ultra_circuit_constructor, rom)
     EXPECT_EQ(result, true);
 }
 
+/**
+ * @brief A simple-as-possible RAM read test, for easier debugging
+ *
+ */
+TEST(ultra_circuit_constructor, ram_simple)
+{
+    UltraCircuitBuilder builder;
+
+    // Initialize a length 1 RAM array with a single value
+    fr ram_value = 5;
+    uint32_t ram_value_idx = builder.add_variable(ram_value);
+    size_t ram_id = builder.create_RAM_array(/*array_size=*/1);
+    builder.init_RAM_element(ram_id, /*index_value=*/0, ram_value_idx);
+
+    // Read from the RAM array we just created (at the 0th index)
+    uint32_t read_idx = builder.add_variable(0);
+    uint32_t a_idx = builder.read_RAM_array(ram_id, read_idx);
+
+    // Use the result in a simple arithmetic gate
+    builder.create_big_add_gate({
+        a_idx,
+        builder.zero_idx,
+        builder.zero_idx,
+        builder.zero_idx,
+        -1,
+        0,
+        0,
+        0,
+        builder.get_variable(ram_value_idx),
+    });
+
+    EXPECT_TRUE(builder.check_circuit());
+}
+
 TEST(ultra_circuit_constructor, ram)
 {
     UltraCircuitBuilder circuit_constructor = UltraCircuitBuilder();

barretenberg/cpp/src/barretenberg/proof_system/composer/composer_lib.hpp

@@ -6,43 +6,6 @@
 
 namespace bb {
 
-/**
- * @brief Copy memory read/write record data into proving key
- * @details Prover needs to know which gates contain a read/write 'record' witness on the 4th wire. This wire value can
- * only be fully computed once the first 3 wire polynomials have been committed to. The 4th wire on these gates will be
- * a random linear combination of the first 3 wires, using the plookup challenge `eta`. Because we shift the gates by
- * the number of public inputs, we need to update the records with the public_inputs offset
- *
- * @tparam Flavor
- * @param circuit
- * @param proving_key
- */
-template <typename Flavor>
-void populate_memory_read_write_records(const typename Flavor::CircuitBuilder& circuit,
-                                        const std::shared_ptr<typename Flavor::ProvingKey>& proving_key)
-{
-    // Determine offset of conventional gates in execution trace
-    auto offset = static_cast<uint32_t>(circuit.public_inputs.size());
-    if (Flavor::has_zero_row) {
-        offset += 1;
-    }
-    if constexpr (IsGoblinFlavor<Flavor>) {
-        offset += static_cast<uint32_t>(circuit.num_ecc_op_gates);
-    }
-    auto add_public_inputs_offset = [offset](uint32_t gate_index) { return gate_index + offset; };
-    proving_key->memory_read_records = std::vector<uint32_t>();
-    proving_key->memory_write_records = std::vector<uint32_t>();
-
-    std::transform(circuit.memory_read_records.begin(),
-                   circuit.memory_read_records.end(),
-                   std::back_inserter(proving_key->memory_read_records),
-                   add_public_inputs_offset);
-    std::transform(circuit.memory_write_records.begin(),
-                   circuit.memory_write_records.end(),
-                   std::back_inserter(proving_key->memory_write_records),
-                   add_public_inputs_offset);
-}
-
 template <typename Flavor>
 std::array<typename Flavor::Polynomial, 4> construct_lookup_table_polynomials(
     const typename Flavor::CircuitBuilder& circuit, size_t dyadic_circuit_size, size_t additional_offset = 0)

barretenberg/cpp/src/barretenberg/proof_system/execution_trace/execution_trace.cpp

@@ -14,6 +14,10 @@ void ExecutionTrace_<Flavor>::populate(Builder& builder,
 
     add_wires_and_selectors_to_proving_key(trace_data, builder, proving_key);
 
+    if constexpr (IsUltraPlonkOrHonk<Flavor>) {
+        add_memory_records_to_proving_key(trace_data, builder, proving_key);
+    }
+
     if constexpr (IsGoblinFlavor<Flavor>) {
         add_ecc_op_wires_to_proving_key(builder, proving_key);
     }
@@ -45,6 +49,22 @@ void ExecutionTrace_<Flavor>::add_wires_and_selectors_to_proving_key(
     }
 }
 
+template <class Flavor>
+void ExecutionTrace_<Flavor>::add_memory_records_to_proving_key(
+    TraceData& trace_data, Builder& builder, const std::shared_ptr<typename Flavor::ProvingKey>& proving_key)
+    requires IsUltraPlonkOrHonk<Flavor>
+{
+    ASSERT(proving_key->memory_read_records.empty() && proving_key->memory_write_records.empty());
+
+    // Update indices of RAM/ROM reads/writes based on where block containing these gates sits in the trace
+    for (auto& index : builder.memory_read_records) {
+        proving_key->memory_read_records.emplace_back(index + trace_data.ram_rom_offset);
+    }
+    for (auto& index : builder.memory_write_records) {
+        proving_key->memory_write_records.emplace_back(index + trace_data.ram_rom_offset);
+    }
+}
+
 template <class Flavor>
 typename ExecutionTrace_<Flavor>::TraceData ExecutionTrace_<Flavor>::construct_trace_data(Builder& builder,
                                                                                           size_t dyadic_circuit_size)
@@ -57,7 +77,7 @@ typename ExecutionTrace_<Flavor>::TraceData ExecutionTrace_<Flavor>::construct_t
     uint32_t offset = Flavor::has_zero_row ? 1 : 0; // Offset at which to place each block in the trace polynomials
     // For each block in the trace, populate wire polys, copy cycles and selector polys
     for (auto& block : builder.blocks.get()) {
-        auto block_size = static_cast<uint32_t>(block.wires[0].size());
+        auto block_size = static_cast<uint32_t>(block.size());
 
         // Update wire polynomials and copy cycles
         // NB: The order of row/column loops is arbitrary but needs to be row/column to match old copy_cycle code
@@ -82,6 +102,11 @@ typename ExecutionTrace_<Flavor>::TraceData ExecutionTrace_<Flavor>::construct_t
             }
         }
 
+        // Store the offset of the block containing RAM/ROM read/write gates for use in updating memory records
+        if (block.has_ram_rom) {
+            trace_data.ram_rom_offset = offset;
+        }
+
         offset += block_size;
     }
     return trace_data;

barretenberg/cpp/src/barretenberg/proof_system/execution_trace/execution_trace.hpp

@@ -21,6 +21,8 @@ template <class Flavor> class ExecutionTrace_ {
         std::array<Polynomial, Builder::Arithmetization::NUM_SELECTORS> selectors;
         // A vector of sets (vectors) of addresses into the wire polynomials whose values are copy constrained
         std::vector<CyclicPermutation> copy_cycles;
+        // The starting index in the trace of the block containing RAM/RAM read/write gates
+        uint32_t ram_rom_offset = 0;
 
         TraceData(size_t dyadic_circuit_size, Builder& builder)
         {
@@ -54,6 +56,23 @@ template <class Flavor> class ExecutionTrace_ {
                                                        Builder& builder,
                                                        const std::shared_ptr<typename Flavor::ProvingKey>& proving_key);
 
+    /**
+     * @brief Add the memory records indicating which rows correspond to RAM/ROM reads/writes
+     * @details The 4th wire of RAM/ROM read/write gates is generated at proving time as a linear combination of the
+     * first three wires scaled by powers of a challenge. To know on which rows to perform this calculation, we must
+     * store the indices of read/write gates in the proving key. In the builder, we store the row index of these gates
+     * within the block containing them. To obtain the row index in the trace at large, we simply increment these
+     * indices by the offset at which that block is placed into the trace.
+     *
+     * @param trace_data
+     * @param builder
+     * @param proving_key
+     */
+    static void add_memory_records_to_proving_key(TraceData& trace_data,
+                                                  Builder& builder,
+                                                  const std::shared_ptr<typename Flavor::ProvingKey>& proving_key)
+        requires IsUltraPlonkOrHonk<Flavor>;
+
     /**
      * @brief Construct wire polynomials, selector polynomials and copy cycles from raw circuit data
      *

barretenberg/cpp/src/barretenberg/sumcheck/instance/prover_instance.hpp

@@ -89,8 +89,6 @@ template <class Flavor> class ProverInstance_ {
 
         sorted_polynomials = construct_sorted_list_polynomials<Flavor>(circuit, dyadic_circuit_size);
 
-        populate_memory_read_write_records<Flavor>(circuit, proving_key);
-
         verification_key = std::make_shared<VerificationKey>(proving_key);
     }