polynomial.hpp

Go to the documentation of this file.

// === AUDIT STATUS ===
// internal:    { status: Complete, auditors: [Nishat], commit: 94f596f8b3bbbc216f9ad7dc33253256141156b2 }
// external_1:  { status: not started, auditors: [], commit: }
// external_2:  { status: not started, auditors: [], commit: }
// =====================
 
#pragma once
#include "barretenberg/common/assert.hpp"
#include "barretenberg/common/bb_bench.hpp"
#include "barretenberg/common/mem.hpp"
#include "barretenberg/common/thread.hpp"
#include "barretenberg/common/throw_or_abort.hpp"
#include "barretenberg/common/type_traits.hpp"
#include "barretenberg/common/zip_view.hpp"
#include "barretenberg/constants.hpp"
#include "barretenberg/polynomials/shared_shifted_virtual_zeroes_array.hpp"
#include "evaluation_domain.hpp"
#include "polynomial_arithmetic.hpp"
#include <cstddef>
#include <fstream>
#include <ranges>
namespace bb {
 
/* Span class with a start index offset.
 * We conceptually have a span like a_0 + a_1 x ... a_n x^n and then multiply by x^start_index.
 * This allows more efficient representation than a fully defined span for 'islands' of zeroes. */
template <typename Fr> struct PolynomialSpan {
    size_t start_index;
    std::span<Fr> span;
    PolynomialSpan(size_t start_index, std::span<Fr> span)
        : start_index(start_index)
        , span(span)
    {}
    size_t end_index() const { return start_index + size(); }
    Fr* data() { return span.data(); }
    size_t size() const { return span.size(); }
    Fr& operator[](size_t index)
    {
        BB_ASSERT_DEBUG(index >= start_index);
        BB_ASSERT_DEBUG(index < end_index());
        return span[index - start_index];
    }
    const Fr& operator[](size_t index) const
    {
        BB_ASSERT_DEBUG(index >= start_index);
        BB_ASSERT_DEBUG(index < end_index());
        return span[index - start_index];
    }
    PolynomialSpan subspan(size_t offset, size_t length)
    {
        if (offset > span.size()) { // Return a null span
            return { 0, span.subspan(span.size()) };
        }
        size_t new_length = std::min(length, span.size() - offset);
        return { start_index + offset, span.subspan(offset, new_length) };
    }
    operator PolynomialSpan<const Fr>() const { return PolynomialSpan<const Fr>(start_index, span); }
};
 
template <typename Fr> class Polynomial {
  public:
    using FF = Fr;
    enum class DontZeroMemory { FLAG };
 
    Polynomial(size_t size, size_t virtual_size, size_t start_index = 0);
    // Intended just for plonk, where size == virtual_size always
    Polynomial(size_t size)
        : Polynomial(size, size) {};
 
    // Constructor that does not initialize values, use with caution to save time.
    Polynomial(size_t size, size_t virtual_size, size_t start_index, DontZeroMemory flag);
    Polynomial(size_t size, size_t virtual_size, DontZeroMemory flag)
        : Polynomial(size, virtual_size, 0, flag)
    {}
    Polynomial(size_t size, DontZeroMemory flag)
        : Polynomial(size, size, flag)
    {}
    Polynomial(const Polynomial& other);
    Polynomial(const Polynomial& other, size_t target_size);
 
    // Explicit move constructor zeroes the source's scalar bookkeeping so that a moved-from
    // polynomial reports size() == 0 / virtual_size() == 0, matching the default-constructed
    // empty state. SharedShiftedVirtualZeroesArray's compiler-synthesized move copies its
    // scalar fields while only the BackingMemory is moved out, leaving size() > 0 with
    // data() == nullptr — UB on next use.
    Polynomial(Polynomial&& other) noexcept
        : coefficients_(std::move(other.coefficients_))
    {
        other.coefficients_.start_ = 0;
        other.coefficients_.end_ = 0;
        other.coefficients_.virtual_size_ = 0;
    }
 
    Polynomial(std::span<const Fr> coefficients, size_t virtual_size);
 
    Polynomial(std::span<const Fr> coefficients)
        : Polynomial(coefficients, coefficients.size())
    {}
 
    static Polynomial shiftable(size_t virtual_size, bool masked = false)
    {
        auto p = Polynomial(
            /*actual size*/ virtual_size - NUM_ZERO_ROWS, virtual_size, /*shiftable offset*/ NUM_ZERO_ROWS);
        if (masked) {
            p.add_masking();
        }
        return p;
    }
    static Polynomial shiftable(size_t size, size_t virtual_size, bool masked = false)
    {
        auto p = Polynomial(/*actual size*/ size - NUM_ZERO_ROWS, virtual_size, /*shiftable offset*/ NUM_ZERO_ROWS);
        if (masked) {
            p.add_masking();
        }
        return p;
    }
    static Polynomial shiftable(size_t size, size_t virtual_size, DontZeroMemory flag)
    {
        return Polynomial(/*actual size*/ size - NUM_ZERO_ROWS, virtual_size, /*shiftable offset*/ NUM_ZERO_ROWS, flag);
    }
    // Allow polynomials to be entirely reset/dormant
    Polynomial() = default;
 
    Polynomial(std::span<const Fr> interpolation_points, std::span<const Fr> evaluations, size_t virtual_size);
 
    // move assignment; mirrors the move constructor above (see comment there).
    Polynomial& operator=(Polynomial&& other) noexcept
    {
        if (this != &other) {
            coefficients_ = std::move(other.coefficients_);
            other.coefficients_.start_ = 0;
            other.coefficients_.end_ = 0;
            other.coefficients_.virtual_size_ = 0;
        }
        return *this;
    }
    Polynomial& operator=(const Polynomial& other);
    ~Polynomial() = default;
 
    Polynomial share() const;
 
    bool is_zero() const
    {
        if (is_empty()) {
            throw_or_abort("Checking is_zero on an empty Polynomial!");
        }
        for (size_t i = 0; i < size(); i++) {
            if (coefficients_.data()[i] != 0) {
                return false;
            }
        }
        return true;
    }
 
    bool operator==(Polynomial const& rhs) const;
 
    const Fr& get(size_t i, size_t virtual_padding = 0) const { return coefficients_.get(i, virtual_padding); };
 
    bool is_empty() const { return coefficients_.size() == 0; }
 
    Polynomial shifted() const;
 
    Polynomial reverse() const;
 
    Fr evaluate_mle(std::span<const Fr> evaluation_points, bool shift = false) const;
 
    void factor_roots(const Fr& root) { polynomial_arithmetic::factor_roots(coeffs(), root); };
 
    Fr evaluate(const Fr& z) const;
 
    void add_scaled(PolynomialSpan<const Fr> other, const Fr& scaling_factor);
 
    void add_scaled_chunk(const ThreadChunk& chunk, PolynomialSpan<const Fr> other, const Fr& scaling_factor);
 
    Polynomial& operator+=(PolynomialSpan<const Fr> other);
 
    Polynomial& operator-=(PolynomialSpan<const Fr> other);
 
    Polynomial& operator*=(const Fr& scaling_factor);
 
    void multiply_chunk(const ThreadChunk& chunk, const Fr& scaling_factor);
 
    std::size_t size() const { return coefficients_.size(); }
    std::size_t virtual_size() const { return coefficients_.virtual_size(); }
    void increase_virtual_size(const size_t size_in) { coefficients_.increase_virtual_size(size_in); };
 
    Fr* data() { return coefficients_.data(); }
    const Fr* data() const { return coefficients_.data(); }
 
    Fr& at(size_t index) { return coefficients_[index]; }
    const Fr& at(size_t index) const { return coefficients_[index]; }
 
    const Fr& operator[](size_t i) { return get(i); }
    const Fr& operator[](size_t i) const { return get(i); }
 
    static Polynomial random(size_t size, size_t start_index = 0)
    {
        BB_BENCH_NAME("generate random polynomial");
 
        BB_ASSERT_GTE(size, start_index);
        return random(size - start_index, size, start_index);
    }
 
    static Polynomial random(size_t size, size_t virtual_size, size_t start_index)
    {
        Polynomial p(size, virtual_size, start_index, DontZeroMemory::FLAG);
        parallel_for_heuristic(
            size,
            [&](size_t i) { p.coefficients_.data()[i] = Fr::random_element(); },
            thread_heuristics::ALWAYS_MULTITHREAD);
        return p;
    }
 
    static Polynomial create_non_parallel_zero_init(size_t size, size_t virtual_size);
 
    void shrink_end_index(const size_t new_end_index);
 
    Polynomial full() const;
 
    // The extents of the actual memory-backed polynomial region
    size_t start_index() const { return coefficients_.start_; }
    size_t end_index() const { return coefficients_.end_; }
    bool is_shiftable() const { return start_index() == NUM_ZERO_ROWS; }
 
    std::span<Fr> coeffs(size_t offset = 0) { return { data() + offset, data() + size() }; }
    std::span<const Fr> coeffs(size_t offset = 0) const { return { data() + offset, data() + size() }; }
    operator PolynomialSpan<Fr>() { return { start_index(), coeffs() }; }
 
    operator PolynomialSpan<const Fr>() const { return { start_index(), coeffs() }; }
 
    auto indices() const { return std::ranges::iota_view(start_index(), end_index()); }
    auto indexed_values() { return zip_view(indices(), coeffs()); }
    auto indexed_values() const { return zip_view(indices(), coeffs()); }
    bool is_valid_set_index(size_t index) const { return (index >= start_index() && index < end_index()); }
    void set_if_valid_index(size_t index, const Fr& value)
    {
        BB_ASSERT_NO_WASM(value.is_zero() || is_valid_set_index(index));
        if (is_valid_set_index(index)) {
            at(index) = value;
        }
    }
 
    template <typename T> void copy_vector(const std::vector<T>& vec)
    {
        BB_ASSERT_LTE(vec.size(), end_index());
        BB_ASSERT_LTE(vec.size() - start_index(), size());
        for (size_t i = start_index(); i < vec.size(); i++) {
            at(i) = vec[i];
        }
    }
 
    void add_masking()
    {
        for (size_t j = 0; j < NUM_MASKED_ROWS; j++) {
            at(NUM_ZERO_ROWS + j) = Fr::random_element();
        }
    }
 
  private:
    // allocate a fresh memory pointer for backing memory
    // DOES NOT initialize memory
    void allocate_backing_memory(size_t size, size_t virtual_size, size_t start_index);
 
    // The underlying memory, with a bespoke (but minimal) shared array struct that fits our needs.
    // Namely, it supports polynomial shifts and 'virtual' zeroes past a size up until a 'virtual' size.
    SharedShiftedVirtualZeroesArray<Fr> coefficients_;
};
 
template <typename Fr>
void add_scaled_batch(Polynomial<Fr>& dst,
                      std::span<const PolynomialSpan<const Fr>> sources,
                      std::span<const Fr> scalars);
 
// NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays)
template <typename Fr> std::shared_ptr<Fr[]> _allocate_aligned_memory(size_t n_elements)
{
    // NOLINTNEXTLINE(cppcoreguidelines-avoid-c-arrays)
    return std::make_shared<Fr[]>(n_elements);
}
 
template <typename Fr_>
Fr_ _evaluate_mle(std::span<const Fr_> evaluation_points,
                  const SharedShiftedVirtualZeroesArray<Fr_>& coefficients,
                  bool shift)
{
    constexpr bool is_native = IsAnyOf<Fr_, bb::fr, grumpkin::fr>;
    // shift ==> native
    BB_ASSERT(!shift || is_native);
 
    if (coefficients.size() == 0) {
        return Fr_(0);
    }
 
    const size_t n = evaluation_points.size();
    // A 0-variable MLE is the constant polynomial; return the single coefficient directly.
    if (n == 0) {
        BB_ASSERT_EQ(coefficients.virtual_size(), 1UL);
        return coefficients.get(0);
    }
    const size_t dim = numeric::get_msb(coefficients.end_ - 1) + 1; // Round up to next power of 2
 
    // To simplify handling of edge cases, we assume that the index space is always a power of 2
    BB_ASSERT_EQ(coefficients.virtual_size(), 1UL << n);
 
    // We first fold over dim rounds l = 0,...,dim-1.
    // in round l, n_l is the size of the buffer containing the Polynomial partially evaluated
    // at u₀,..., u_l.
    // In round 0, this is half the size of dim
    size_t n_l = 1UL << (dim - 1);
 
    // temporary buffer of half the size of the Polynomial
    auto tmp_ptr = _allocate_aligned_memory<Fr_>(n_l);
    auto tmp = tmp_ptr.get();
 
    size_t offset = 0;
    if constexpr (is_native) {
        if (shift) {
            BB_ASSERT_EQ(coefficients.get(0), Fr_::zero());
            offset++;
        }
    }
 
    Fr_ u_l = evaluation_points[0];
 
    // Note below: i * 2 + 1 + offset might equal virtual_size. This used to subtlely be handled by extra capacity
    // padding (and there used to be no assert time checks, which this constant helps with).
    const size_t ALLOW_ONE_PAST_READ = 1;
    for (size_t i = 0; i < n_l; ++i) {
        // curr[i] = (Fr(1) - u_l) * prev[i * 2] + u_l * prev[(i * 2) + 1];
        tmp[i] = coefficients.get(i * 2 + offset) +
                 u_l * (coefficients.get(i * 2 + 1 + offset, ALLOW_ONE_PAST_READ) - coefficients.get(i * 2 + offset));
    }
 
    // partially evaluate the dim-1 remaining points
    for (size_t l = 1; l < dim; ++l) {
        n_l = 1UL << (dim - l - 1);
        u_l = evaluation_points[l];
        for (size_t i = 0; i < n_l; ++i) {
            tmp[i] = tmp[i * 2] + u_l * (tmp[(i * 2) + 1] - tmp[i * 2]);
        }
    }
    auto result = tmp[0];
 
    // We handle the "trivial" dimensions which are full of zeros.
    for (size_t i = dim; i < n; i++) {
        result *= (Fr_(1) - evaluation_points[i]);
    }
 
    return result;
}
 
template <typename Fr_>
Fr_ generic_evaluate_mle(std::span<const Fr_> evaluation_points,
                         const SharedShiftedVirtualZeroesArray<Fr_>& coefficients)
{
    return _evaluate_mle(evaluation_points, coefficients, false);
}
 
template <typename Fr> inline std::ostream& operator<<(std::ostream& os, const Polynomial<Fr>& p)
{
    if (p.size() == 0) {
        return os << "[]";
    }
    if (p.size() == 1) {
        return os << "[ data " << p[0] << "]";
    }
    return os << "[ data\n"
              << "  " << p[0] << ",\n"
              << "  " << p[1] << ",\n"
              << "  ... ,\n"
              << "  " << p[p.size() - 2] << ",\n"
              << "  " << p[p.size() - 1] << ",\n"
              << "]";
}
 
template <typename Poly, typename... Polys> auto zip_polys(Poly&& poly, Polys&&... polys)
{
    // Ensure all polys have the same start_index() and end_index() as poly
    // Use fold expression to check all polys exactly match our size
    // Wrap BB_ASSERT_EQ_RELEASE in a lambda to make it usable in a fold expression
    auto check_indices = [&](const auto& other) {
        BB_ASSERT_EQ(poly.start_index(), other.start_index());
        BB_ASSERT_EQ(poly.end_index(), other.end_index());
    };
    // Apply the lambda to each poly in the parameter pack
    (check_indices(polys), ...);
    return zip_view(poly.indices(), poly.coeffs(), polys.coeffs()...);
}
} // namespace bb