fix: Store lagrange forms of selector polys w/ Ultra (AztecProtocol/b…

…arretenberg#255)

* store lagrange forms of selector polynomials when serializing pk for Ultra

* added comment to ultra_selector_properties

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

@@ -46,10 +46,19 @@ namespace plonk {
 
 std::vector<ComposerBase::SelectorProperties> ultra_selector_properties()
 {
+    // When reading and writing the proving key from a buffer we must precompute the Lagrange form of certain selector
+    // polynomials. In order to avoid a new selector type and definitions in the polynomial manifest, we can instead
+    // store the Lagrange forms of all the selector polynomials.
+    //
+    // This workaround increases the memory footprint of the prover, and is a possible place of improvement in the
+    // future. Below is the previous state showing where the Lagrange form is necessary for a selector:
+    //     { "q_m", true },         { "q_c", true },    { "q_1", true },        { "q_2", true },
+    //     { "q_3", true },         { "q_4", false },   { "q_arith", false },   { "q_sort", false },
+    //     { "q_elliptic", false }, { "q_aux", false }, { "table_type", true },
     std::vector<ComposerBase::SelectorProperties> result{
-        { "q_m", true },         { "q_c", true },    { "q_1", true },        { "q_2", true },
-        { "q_3", true },         { "q_4", false },   { "q_arith", false },   { "q_sort", false },
-        { "q_elliptic", false }, { "q_aux", false }, { "table_type", true },
+        { "q_m", true },        { "q_c", true },   { "q_1", true },        { "q_2", true },
+        { "q_3", true },        { "q_4", true },   { "q_arith", true },    { "q_sort", true },
+        { "q_elliptic", true }, { "q_aux", true }, { "table_type", true },
     };
     return result;
 }

barretenberg/cpp/src/barretenberg/proof_system/proving_key/proving_key.test.cpp

@@ -3,6 +3,7 @@
 #include "proving_key.hpp"
 #include "serialize.hpp"
 #include "barretenberg/plonk/composer/standard_composer.hpp"
+#include "barretenberg/plonk/composer/ultra_composer.hpp"
 
 #ifndef __wasm__
 #include <filesystem>
@@ -47,6 +48,42 @@ TEST(proving_key, proving_key_from_serialized_key)
     EXPECT_EQ(p_key.contains_recursive_proof, proving_key->contains_recursive_proof);
 }
 
+// Test proving key serialization/deserialization to/from buffer using UltraComposer
+TEST(proving_key, proving_key_from_serialized_key_ultra)
+{
+    plonk::UltraComposer composer = plonk::UltraComposer();
+    fr a = fr::one();
+    composer.add_public_variable(a);
+
+    bonk::proving_key& p_key = *composer.compute_proving_key();
+    auto pk_buf = to_buffer(p_key);
+    auto pk_data = from_buffer<bonk::proving_key_data>(pk_buf);
+    auto crs = std::make_unique<bonk::FileReferenceStringFactory>("../srs_db/ignition");
+    auto proving_key =
+        std::make_shared<bonk::proving_key>(std::move(pk_data), crs->get_prover_crs(pk_data.circuit_size + 1));
+
+    // Loop over all pre-computed polys for the given composer type and ensure equality
+    // between original proving key polynomial store and the polynomial store that was
+    // serialized/deserialized from buffer
+    bonk::PrecomputedPolyList precomputed_poly_list(p_key.composer_type);
+    bool all_polys_are_equal{ true };
+    for (size_t i = 0; i < precomputed_poly_list.size(); ++i) {
+        std::string poly_id = precomputed_poly_list[i];
+        barretenberg::polynomial input_poly = p_key.polynomial_store.get(poly_id);
+        barretenberg::polynomial output_poly = proving_key->polynomial_store.get(poly_id);
+        all_polys_are_equal = all_polys_are_equal && (input_poly == output_poly);
+    }
+
+    // Check that all pre-computed polynomials are equal
+    EXPECT_EQ(all_polys_are_equal, true);
+
+    // Check equality of other proving_key_data data
+    EXPECT_EQ(p_key.composer_type, proving_key->composer_type);
+    EXPECT_EQ(p_key.circuit_size, proving_key->circuit_size);
+    EXPECT_EQ(p_key.num_public_inputs, proving_key->num_public_inputs);
+    EXPECT_EQ(p_key.contains_recursive_proof, proving_key->contains_recursive_proof);
+}
+
 // Test that a proving key can be serialized/deserialized using mmap
 #ifndef __wasm__
 TEST(proving_key, proving_key_from_mmaped_key)

barretenberg/cpp/src/barretenberg/proof_system/types/polynomial_manifest.hpp

@@ -234,6 +234,10 @@ class PrecomputedPolyList {
             case PolynomialSource::SELECTOR: // monomial and fft
                 precomputed_poly_ids.emplace_back(label);
                 precomputed_poly_ids.emplace_back(label + "_fft");
+                // Store all lagrange forms of selector polynomials for ultra
+                if (composer_type == plonk::ComposerType::PLOOKUP) {
+                    precomputed_poly_ids.emplace_back(label + "_lagrange");
+                }
                 break;
             case PolynomialSource::PERMUTATION: // monomial, fft, and lagrange
                 precomputed_poly_ids.emplace_back(label);