From ba5afe4596a22498730e4968ab1645312b4e91f7 Mon Sep 17 00:00:00 2001
From: Martun Karapetyan <martun.karapetyan@gmail.com>
Date: Mon, 16 Sep 2024 17:09:24 +0400
Subject: [PATCH] Implementing command compute-combined-Q for aggregated FRI.

---
 .../crypto3/zk/commitments/polynomial/lpc.hpp |  32 +++++-
 .../snark/systems/plonk/placeholder/proof.hpp |   2 +-
 proof-producer/README.md                      |  27 ++---
 .../nil/proof-generator/arg_parser.hpp        |   2 +
 .../include/nil/proof-generator/prover.hpp    | 103 ++++++++++++++----
 .../bin/proof-producer/src/arg_parser.cpp     |   8 +-
 6 files changed, 127 insertions(+), 47 deletions(-)

diff --git a/crypto3/libs/zk/include/nil/crypto3/zk/commitments/polynomial/lpc.hpp b/crypto3/libs/zk/include/nil/crypto3/zk/commitments/polynomial/lpc.hpp
index 95e0d99733..120adf3767 100644
--- a/crypto3/libs/zk/include/nil/crypto3/zk/commitments/polynomial/lpc.hpp
+++ b/crypto3/libs/zk/include/nil/crypto3/zk/commitments/polynomial/lpc.hpp
@@ -269,7 +269,8 @@ namespace nil {
                         return fri_proof;
                     }
 
-                    /** \brief 
+                    /** \brief Computes polynomial combined_Q. In case this function changes, 
+                               the function 'compute_theta_power_for_combined_Q' below should be changed accordingly. 
                      *  \param theta The value of challenge. When called from aggregated FRI, this values is sent from
                                 the "main prover" machine.
                      *  \param starting_power When aggregated FRI is used, the value is not zero, it's the total degree of all
@@ -350,6 +351,35 @@ namespace nil {
                         return combined_Q;
                     }
 
+                    // Computes and returns the maximal power of theta used to compute the value of Combined_Q.
+                    std::size_t compute_theta_power_for_combined_Q() const {
+                        std::size_t theta_power = 0;
+                        this->build_points_map();
+
+                        auto points = this->get_unique_points();
+
+                        for (auto const &point: points) {
+                            for (std::size_t i: this->_z.get_batches()) {
+                                for (std::size_t j = 0; j < this->_z.get_batch_size(i); j++) {
+                                    auto iter = this->_points_map[i][j].find(point);
+                                    if (iter == this->_points_map[i][j].end())
+                                        continue;
+
+                                    theta_power++;
+                                }
+                            }
+                        }
+
+                        for (std::size_t i: this->_z.get_batches()) {
+                            if (!_batch_fixed.at(i))
+                                continue;
+
+                            theta_power += this->_z.get_batch_size(i);
+                        }
+
+                        return theta_power;
+                    }
+
                     bool verify_eval(
                         const proof_type &proof,
                         const std::map<std::size_t, commitment_type> &commitments,
diff --git a/crypto3/libs/zk/include/nil/crypto3/zk/snark/systems/plonk/placeholder/proof.hpp b/crypto3/libs/zk/include/nil/crypto3/zk/snark/systems/plonk/placeholder/proof.hpp
index c5cd9b9d35..ce1948ba33 100644
--- a/crypto3/libs/zk/include/nil/crypto3/zk/snark/systems/plonk/placeholder/proof.hpp
+++ b/crypto3/libs/zk/include/nil/crypto3/zk/snark/systems/plonk/placeholder/proof.hpp
@@ -141,7 +141,7 @@ namespace nil {
 
                     // This vector contains N partial proofs, one per prover.
                     std::vector<placeholder_partial_proof<FieldType, ParamsType>> partial_proofs;
-                    typename commitment_type::aggregated_proof_type aggregated_proof;
+                    typename commitment_scheme_type::aggregated_proof_type aggregated_proof;
                 };
             }    // namespace snark
         }        // namespace zk
diff --git a/proof-producer/README.md b/proof-producer/README.md
index e9b180507d..ecf9335b0b 100644
--- a/proof-producer/README.md
+++ b/proof-producer/README.md
@@ -39,24 +39,6 @@ proof-generator --circuit <circuit-file> --assignment <assignment-file> --proof
 ```
 
 # Building from source
-1. Install dependencies:
-    ```
-    sudo apt-get install \
-        build-essential \
-        libsctp-dev \
-        libssl-dev \
-        libicu-dev \
-        lsb-release \
-        gnutls-dev \
-        pkg-config
-    ```
-
-2. Build with CMake:
-    ```mkdir build
-    cd build
-    cmake ..
-    make -j $(nrpoc)
-    ```
 
 # Sample calls to proof-producer
 
@@ -84,7 +66,12 @@ Verify generated proof:
 ./build/bin/proof-producer/proof-producer-single-threaded --stage="verify" --circuit="circuit.crct" --common-data="preprocessed_common_data.dat" --proof="proof.bin" --assignment-description-file="assignment-description.dat" -q 10
 ```
 
-Aggregate challenges
+Aggregate challenges, done once on the main prover
 ```bash
 ./build/bin/proof-producer/proof-producer-single-threaded --stage="generate-aggregated-challenge" --input-challenge-files challenge1.dat challenge2.dat --aggregated-challenge-file="aggregated_challenge.dat"
-```
\ No newline at end of file
+```
+
+Compute polynomial combined_Q, done on each prover
+```bash
+./build/bin/proof-producer/proof-producer-single-threaded --stage="generate-combined-Q" --aggregated-challenge-file="aggregated_challenge.dat" --combined-Q-starting-power=0  --commitment-state-file="commitment_state.dat" --combined-Q-polynomial-file="combined-Q.dat"
+```
diff --git a/proof-producer/bin/proof-producer/include/nil/proof-generator/arg_parser.hpp b/proof-producer/bin/proof-producer/include/nil/proof-generator/arg_parser.hpp
index 2194f46613..e678de6f64 100644
--- a/proof-producer/bin/proof-producer/include/nil/proof-generator/arg_parser.hpp
+++ b/proof-producer/bin/proof-producer/include/nil/proof-generator/arg_parser.hpp
@@ -45,6 +45,8 @@ namespace nil {
             boost::filesystem::path assignment_description_file_path;
             std::vector<boost::filesystem::path> input_challenge_files;
             boost::filesystem::path aggregated_challenge_file = "aggregated_challenge.dat";
+            boost::filesystem::path combined_Q_polynomial_file = "combined_Q.dat";
+            std::size_t combined_Q_starting_power;
             boost::log::trivial::severity_level log_level = boost::log::trivial::severity_level::info;
             CurvesVariant elliptic_curve_type = type_identity<nil::crypto3::algebra::curves::pallas>{};
             HashesVariant hash_type = type_identity<nil::crypto3::hashes::keccak_1600<256>>{};
diff --git a/proof-producer/bin/proof-producer/include/nil/proof-generator/prover.hpp b/proof-producer/bin/proof-producer/include/nil/proof-generator/prover.hpp
index ab88ecd573..3cac0de5e1 100644
--- a/proof-producer/bin/proof-producer/include/nil/proof-generator/prover.hpp
+++ b/proof-producer/bin/proof-producer/include/nil/proof-generator/prover.hpp
@@ -23,6 +23,7 @@
 #include <fstream>
 #include <random>
 #include <sstream>
+#include <optional>
 
 #include <boost/log/trivial.hpp>
 
@@ -75,7 +76,8 @@ namespace nil {
                 auto read_iter = v->begin();
                 auto status = marshalled_data.read(read_iter, v->size());
                 if (status != nil::marshalling::status_type::success) {
-                    BOOST_LOG_TRIVIAL(error) << "When reading a Marshalled structure from file " << path << ", decoding step failed";
+                    BOOST_LOG_TRIVIAL(error) << "When reading a Marshalled structure from file "
+                        << path << ", decoding step failed.";
                     return std::nullopt;
                 }
                 return marshalled_data;
@@ -104,7 +106,12 @@ namespace nil {
                 PREPROCESS = 1,
                 PROVE = 2,
                 VERIFY = 3,
-                GENERATE_AGGREGATED_CHALLENGE = 4
+                GENERATE_AGGREGATED_CHALLENGE = 4,
+                PARTIAL_PROVE = 5,
+                COMPUTE_COMBINED_Q = 6,
+                GENERATE_FRI_PROOF = 7,
+                GENERATE_LPC_INITIAL_PROOF = 8,
+                MERGE_PROOFS = 9
             };
 
             ProverStage prover_stage_from_string(const std::string& stage) {
@@ -113,7 +120,9 @@ namespace nil {
                     {"preprocess", ProverStage::PREPROCESS},
                     {"prove", ProverStage::PROVE},
                     {"verify", ProverStage::VERIFY},
-                    {"generate-aggregated-challenge", ProverStage::GENERATE_AGGREGATED_CHALLENGE}
+                    {"generate-aggregated-challenge", ProverStage::GENERATE_AGGREGATED_CHALLENGE},
+                    {"partial-prove", ProverStage::PARTIAL_PROVE},
+                    {"compute-combined-Q", ProverStage::COMPUTE_COMBINED_Q}
                 };
                 auto it = stage_map.find(stage);
                 if (it == stage_map.end()) {
@@ -132,6 +141,7 @@ namespace nil {
             using LpcParams = nil::crypto3::zk::commitments::list_polynomial_commitment_params<HashType, HashType, 2>;
             using Lpc = nil::crypto3::zk::commitments::list_polynomial_commitment<BlueprintField, LpcParams>;
             using LpcScheme = typename nil::crypto3::zk::commitments::lpc_commitment_scheme<Lpc>;
+            using polynomial_type = typename LpcScheme::polynomial_type;
             using CircuitParams = nil::crypto3::zk::snark::placeholder_circuit_params<BlueprintField>;
             using PlaceholderParams = nil::crypto3::zk::snark::placeholder_params<CircuitParams, LpcScheme>;
             using Proof = nil::crypto3::zk::snark::placeholder_proof<BlueprintField, PlaceholderParams>;
@@ -448,6 +458,40 @@ namespace nil {
                 return true;
             }
 
+            std::optional<typename BlueprintField::value_type> read_challenge(
+                    const boost::filesystem::path& input_file) {
+                using challenge_marshalling_type = nil::crypto3::marshalling::types::field_element<
+                    TTypeBase, typename BlueprintField::value_type>;
+
+                if (!nil::proof_generator::can_read_from_file(input_file.string())) {
+                    BOOST_LOG_TRIVIAL(error) << "Can't read file " << input_file;
+                    return std::nullopt;
+                }
+
+                auto marshalled_challenge = detail::decode_marshalling_from_file<challenge_marshalling_type>(
+                    input_file);
+
+                if (!marshalled_challenge) {
+                    return std::nullopt;
+                }
+
+                return marshalled_challenge->value();
+            }
+
+            bool save_challenge(const boost::filesystem::path& challenge_file,
+                                const typename BlueprintField::value_type& challenge) {
+                using challenge_marshalling_type = nil::crypto3::marshalling::types::field_element<
+                    TTypeBase, typename BlueprintField::value_type>;
+
+                BOOST_LOG_TRIVIAL(info) << "Writing challenge to " << challenge_file << std::endl;
+
+                // marshall the challenge
+                challenge_marshalling_type marshalled_challenge(challenge);
+
+                return detail::encode_marshalling_to_file<challenge_marshalling_type>
+                    (challenge_file, marshalled_challenge);
+            }
+
             void create_lpc_scheme() {
                 // Lambdas and grinding bits should be passed through preprocessor directives
                 std::size_t table_rows_log = std::ceil(std::log2(table_description_->rows_amount));
@@ -497,38 +541,51 @@ namespace nil {
                     return false;
                 }
                 BOOST_LOG_TRIVIAL(info) << "Generating aggregated challenge to " << aggregated_challenge_file;
-                // check that we can access all input files
-                for (const auto &input_file : aggregate_input_files) {
-                    BOOST_LOG_TRIVIAL(info) << "Reading challenge from " << input_file;
-                    if (!nil::proof_generator::can_read_from_file(input_file.string())) {
-                        BOOST_LOG_TRIVIAL(error) << "Can't read file " << input_file;
-                        return false;
-                    }
-                }
+
                 // create the transcript
                 using transcript_hash_type = typename PlaceholderParams::transcript_hash_type;
                 using transcript_type = crypto3::zk::transcript::fiat_shamir_heuristic_sequential<transcript_hash_type>;
-                using challenge_marshalling_type =
-                    nil::crypto3::marshalling::types::field_element<
-                        TTypeBase, typename BlueprintField::value_type>;
                 transcript_type transcript;
+
                 // read challenges from input files and add them to the transcript
                 for (const auto &input_file : aggregate_input_files) {
-                    auto challenge = detail::decode_marshalling_from_file<challenge_marshalling_type>(input_file);
+                    std::optional<typename BlueprintField::value_type> challenge = read_challenge(input_file);
                     if (!challenge) {
-                        BOOST_LOG_TRIVIAL(error) << "Failed to read challenge from " << input_file;
                         return false;
                     }
-                    transcript(challenge->value());
+                    transcript(challenge.value());
                 }
+
                 // produce the aggregated challenge
                 auto output_challenge = transcript.template challenge<BlueprintField>();
-                // marshall the challenge
-                challenge_marshalling_type marshalled_challenge(output_challenge);
-                // write the challenge to the output file
-                BOOST_LOG_TRIVIAL(info) << "Writing aggregated challenge to " << aggregated_challenge_file;
-                return detail::encode_marshalling_to_file<challenge_marshalling_type>
-                    (aggregated_challenge_file, marshalled_challenge);
+
+                return save_challenge(aggregated_challenge_file, output_challenge);
+            }
+
+            bool save_poly_to_file(const polynomial_type& combined_Q,
+                                   const boost::filesystem::path &output_file) {
+                using polynomial_marshalling_type = nil::crypto3::marshalling::types::polynomial<
+                    TTypeBase, polynomial_type>;
+
+                BOOST_LOG_TRIVIAL(info) << "Writing polynomial to " << output_file << std::endl;
+
+                polynomial_marshalling_type marshalled_poly = nil::crypto3::marshalling::types::fill_polynomial<Endianness, polynomial_type>(combined_Q);
+
+                return detail::encode_marshalling_to_file<polynomial_marshalling_type>
+                    (output_file, marshalled_poly);
+            }
+
+            bool read_challenge_and_generate_combined_Q_to_file(
+                const boost::filesystem::path &aggregated_challenge_file,
+                std::size_t starting_power,
+                const boost::filesystem::path &output_combined_Q_file) {
+                std::optional<typename BlueprintField::value_type> challenge = read_challenge(
+                    aggregated_challenge_file);
+                if (!challenge) {
+                    return false;
+                }
+                polynomial_type combined_Q = lpc_scheme_->prepare_combined_Q(challenge, starting_power);
+                return save_poly_to_file(combined_Q, output_combined_Q_file); 
             }
 
         private:
diff --git a/proof-producer/bin/proof-producer/src/arg_parser.cpp b/proof-producer/bin/proof-producer/src/arg_parser.cpp
index efd8ced431..1f5e29dee9 100644
--- a/proof-producer/bin/proof-producer/src/arg_parser.cpp
+++ b/proof-producer/bin/proof-producer/src/arg_parser.cpp
@@ -73,7 +73,7 @@ namespace nil {
             // clang-format off
             auto options_appender = config.add_options()
                 ("stage", make_defaulted_option(prover_options.stage),
-                 "Stage of the prover to run, one of (all, preprocess, prove, verify, generate-aggregated-challenge). Defaults to 'all'.")
+                 "Stage of the prover to run, one of (all, preprocess, prove, verify, generate-aggregated-challenge, generate-combined-Q). Defaults to 'all'.")
                 ("proof,p", make_defaulted_option(prover_options.proof_file_path), "Proof file")
                 ("json,j", make_defaulted_option(prover_options.json_file_path), "JSON proof file")
                 ("common-data", make_defaulted_option(prover_options.preprocessed_common_data_path), "Preprocessed common data file")
@@ -92,7 +92,11 @@ namespace nil {
                 ("input-challenge-files,u", po::value<std::vector<boost::filesystem::path>>(&prover_options.input_challenge_files)->multitoken(),
                  "Input challenge files. Used with 'generate-aggregated-challenge' stage.")
                 ("aggregated-challenge-file", po::value<boost::filesystem::path>(&prover_options.aggregated_challenge_file),
-                 "Aggregated challenge file. Used with 'generate-aggregated-challenge' stage");
+                 "Aggregated challenge file. Used with 'generate-aggregated-challenge' stage")
+                ("combined-Q-polynomial-file", po::value<boost::filesystem::path>(&prover_options.combined_Q_polynomial_file),
+                 "File containing the polynomial combined-Q, generated on a single prover.")
+                ("combined-Q-starting-power", po::value<std::size_t>(&prover_options.combined_Q_starting_power),
+                 "The starting power for combined-Q polynomial for the current prover.");
 
             // clang-format on
             po::options_description cmdline_options("nil; Proof Producer");