An example on how to evaluate a composition polynomial in Stone PRover

composition_poly_evaluation.cc

#include "starkware/stark/stark.h"

#include <algorithm>
#include <cstddef>
#include <fstream>
#include <iostream>
#include <map>
#include <optional>
#include <string>
#include <utility>

#include "gmock/gmock.h"
#include "gtest/gtest.h"

#include "starkware/air/components/permutation/permutation_dummy_air.h"
#include "starkware/air/components/permutation/permutation_trace_context.h"
#include "starkware/air/degree_three_example/degree_three_example_air.h"
#include "starkware/air/degree_three_example/degree_three_example_trace_context.h"
#include "starkware/air/fibonacci/fibonacci_air.h"
#include "starkware/air/fibonacci/fibonacci_trace_context.h"
#include "starkware/air/test_utils.h"
#include "starkware/air/trace_context.h"
#include "starkware/algebra/fields/extension_field_element.h"
#include "starkware/algebra/fields/test_field_element.h"
#include "starkware/channel/annotation_scope.h"
#include "starkware/channel/noninteractive_prover_channel.h"
#include "starkware/channel/noninteractive_verifier_channel.h"
#include "starkware/commitment_scheme/commitment_scheme_builder.h"
#include "starkware/commitment_scheme/table_prover_impl.h"
#include "starkware/commitment_scheme/table_verifier_impl.h"
#include "starkware/crypt_tools/blake2s.h"
#include "starkware/crypt_tools/keccak_256.h"
#include "starkware/error_handling/test_utils.h"
#include "starkware/proof_system/proof_system.h"
#include "starkware/stark/test_utils.h"
#include "starkware/stark/utils.h"
#include "starkware/stl_utils/containers.h"
#include "starkware/utils/maybe_owned_ptr.h"

using namespace starkware;

using StarkField252Element = PrimeFieldElement<252, 0>;
using FibAirT = FibonacciAir<StarkField252Element>;
using FibTraceContextT = FibonacciTraceContext<StarkField252Element>;

std::unique_ptr<CompositionPolynomial>
CreateCompositionPolynomial(Channel *channel, const Field &field,
                            const FieldElement &trace_generator,
                            const Air &air) {
  size_t num_random_coefficients_required = air.NumRandomCoefficients();
  FieldElementVector random_coefficients = FieldElementVector::Make(field);
  FieldElement alpha = channel->GetRandomFieldElementFromVerifier(
      field, "Constraint polynomial random element");
  std::cout << "alpha: " << alpha << std::endl;
  FieldElement curr = field.One();
  for (size_t i = 0; i < num_random_coefficients_required; ++i) {
    random_coefficients.PushBack(curr);
    curr = curr * alpha;
  }
  return air.CreateCompositionPolynomial(trace_generator, random_coefficients);
}

std::vector<size_t> DrawFriStepsList(const size_t log_degree_bound,
                                     Prng *prng) {
  if (prng == nullptr) {
    return {3, 3, 1, 1};
  }
  std::vector<size_t> res;
  for (size_t curr_sum = 0; curr_sum < log_degree_bound;) {
    res.push_back(prng->UniformInt<size_t>(1, log_degree_bound - curr_sum));
    curr_sum += res.back();
  }
  ASSERT_RELEASE(Sum(res) == log_degree_bound,
                 "The sum of all steps must be the log of the degree bound.");
  return res;
}

template <typename FieldElementT>
StarkParameters GenerateParameters(std::unique_ptr<Air> air, Prng *prng,
                                   size_t proof_of_work_bits = 15,
                                   bool use_extension_field = false) {
  // Field used.
  const Field field(Field::Create<FieldElementT>());
  const size_t trace_length = air->TraceLength();
  uint64_t degree_bound =
      air->GetCompositionPolynomialDegreeBound() / trace_length;

  const auto log_n_cosets =
      (prng == nullptr ? 4 : prng->UniformInt<size_t>(Log2Ceil(degree_bound), 6));
  const size_t log_coset_size = Log2Ceil(trace_length);
  const size_t n_cosets = Pow2(log_n_cosets);
  const size_t log_evaluation_domain_size = log_coset_size + log_n_cosets;

  // Fri steps - hard-coded pattern for this tests suit.
  const size_t log_degree_bound = log_coset_size;
  const std::vector<size_t> fri_step_list =
      DrawFriStepsList(log_degree_bound, prng);

  // Bases used by FRI.
  const FieldElementT offset =
      (prng == nullptr ? FieldElementT::One()
                       : FieldElementT::RandomElement(prng));
  FftBasesImpl<FftMultiplicativeGroup<FieldElementT>> bases =
      MakeFftBases(log_evaluation_domain_size, offset);

  // FRI parameters.
  FriParameters fri_params{/*fri_step_list=*/fri_step_list,
                           /*last_layer_degree_bound=*/1,
                           /*n_queries=*/30,
                           /*fft_bases=*/UseMovedValue(std::move(bases)),
                           /*field=*/field,
                           /*proof_of_work_bits=*/proof_of_work_bits};

  return StarkParameters(field, use_extension_field, n_cosets, trace_length,
                         TakeOwnershipFrom(std::move(air)),
                         UseMovedValue(std::move(fri_params)));
}

template <typename HashT, typename FieldElementT>
std::unique_ptr<TableProver>
MakeTableProver(uint64_t n_segments, uint64_t n_rows_per_segment,
                size_t n_columns, ProverChannel *channel,
                size_t n_tasks_per_segment, size_t n_layer,
                size_t n_verifier_friendly_commitment_layers) {
  return GetTableProverFactory<HashT>(channel, FieldElementT::SizeInBytes(),
                                      n_tasks_per_segment, n_layer,
                                      n_verifier_friendly_commitment_layers,
                                      CommitmentHashes(HashT::HashName()))(
      n_segments, n_rows_per_segment, n_columns);
}

int main() {
  // Fibonacci AIR
  StarkField252Element secret = StarkField252Element::One(); // secret
  const uint64_t fibonacci_claim_index = 251;
  const uint64_t trace_length = 256;

  StarkField252Element claimed_fib =
      FibAirT::PublicInputFromPrivateInput(secret, fibonacci_claim_index);
  auto air = FibAirT(trace_length, fibonacci_claim_index, claimed_fib);

  // Prover channel / Transcript
  const Prng channel_prng = Prng(MakeByteArray<0xca, 0xfe, 0xca, 0xfe>());
  NoninteractiveProverChannel prover_channel =
      NoninteractiveProverChannel(channel_prng.Clone());

  // Stark config
  Prng prng = Prng(MakeByteArray<0xca, 0xfe, 0xca, 0xfe>());
  StarkProverConfig stark_config = StarkProverConfig::InRam();
  auto stark_params = GenerateParameters<StarkField252Element>(
      std::make_unique<FibAirT>(trace_length, fibonacci_claim_index,
                                claimed_fib),
      &prng);

  //  Table Prover
  TableProverFactory table_prover_factory =
      [&](uint64_t n_segments, uint64_t n_rows_per_segment, size_t n_columns) {
        const size_t n_verifier_friendly_commitment_layers = 0;
        return MakeTableProver<Keccak256, StarkField252Element>(
            n_segments, n_rows_per_segment, n_columns, &prover_channel,
            stark_config.table_prover_n_tasks_per_segment,
            stark_config.n_out_of_memory_merkle_layers,
            n_verifier_friendly_commitment_layers);
      };

  auto trace_context = std::make_unique<FibTraceContextT>(
      UseOwned(&air), secret, fibonacci_claim_index);

  // First trace.
  Trace trace = trace_context->GetTrace();

  std::vector<MaybeOwnedPtr<CommittedTraceProverBase>> traces;
  // Add first committed trace.
  CommittedTraceProver committed_trace(
      stark_config.cached_lde_config, UseOwned(&stark_params.evaluation_domain),
      trace.Width(), table_prover_factory);
  committed_trace.Commit(std::move(trace),
                         stark_params.evaluation_domain.Bases(), true);

  traces.emplace_back(UseMovedValue(std::move(committed_trace)));

  //  Prepare for interaction.
  const Air *current_air = (stark_params.air).get();

  // Create composition polynomial from AIR.
  std::unique_ptr<CompositionPolynomial> composition_polynomial;
  composition_polynomial = CreateCompositionPolynomial(
      &prover_channel, stark_params.field,
      stark_params.evaluation_domain.TraceGenerator(), *current_air);

  std::cout << stark_params.evaluation_domain.CosetsOffsets() << std::endl;

  CompositionOracleProver composition_oracle(
      UseOwned(&stark_params.evaluation_domain), std::move(traces),
      current_air->GetMask(), UseOwned(current_air),
      UseOwned(composition_polynomial), &prover_channel);

  FieldElementVector composition_poly_evaluations =
      composition_oracle.EvalComposition(
          stark_config.constraint_polynomial_task_size);

  std::cout << composition_poly_evaluations << std::endl;

  return 0;
}