BitPoly.h

Go to the documentation of this file.

#pragma once
// SPDX-FileCopyrightText: 2025 Nessan Fitzmaurice <nzznfitz+gh@icloud.com>
// SPDX-License-Identifier: MIT

 
#include <gf2/BitVec.h>
#include <sstream>
 
namespace gf2 {
 
// Forward declarations to avoid recursive inclusion issues
// clang-format off
template<unsigned_word Word> class BitMat;
// clang-format on

template<unsigned_word Word = usize>
class BitPoly {
private:
    BitVec<Word> m_coeffs;
 
public:
    using coeffs_type = BitVec<Word>;

    using word_type = Word;


    constexpr BitPoly() : m_coeffs{} {}

    constexpr BitPoly(const BitVec<Word>& coeffs) : m_coeffs{coeffs} {}

    constexpr BitPoly(BitVec<Word>&& coeffs) : m_coeffs{std::move(coeffs)} {}


    static constexpr BitPoly zero() { return BitPoly<Word>{}; }

    static constexpr BitPoly one() { return BitPoly<Word>{std::move(coeffs_type::ones(1))}; }

    static constexpr BitPoly constant(bool val) { return BitPoly<Word>{std::move(coeffs_type::constant(val, 1))}; }

    static constexpr BitPoly zeros(usize n) { return BitPoly<Word>{std::move(coeffs_type::zeros(n + 1))}; }

    static constexpr BitPoly ones(usize n) { return BitPoly<Word>{std::move(coeffs_type::ones(n + 1))}; }

    static constexpr BitPoly x_to_the(usize n) { return BitPoly<Word>{std::move(coeffs_type::unit(n + 1, n))}; }

    static constexpr BitPoly from(usize n, std::invocable<usize> auto f) {
        return BitPoly<Word>{std::move(coeffs_type::from(n + 1, f))};
    }


    static constexpr BitPoly random(usize n) {
        // BitPoly of degree `n` has `n + 1` coefficients.
        auto coeffs = coeffs_type::random(n + 1);
 
        // If `n > 0` we want the coefficient of `x^n` to be one for sure. If `n == 0` then a random 0/1 is fine.
        if (n > 0) coeffs.set(n);
        return BitPoly{std::move(coeffs)};
    }

    static constexpr BitPoly seeded_random(usize n, std::uint64_t seed) {
        // BitPoly of degree `n` has `n + 1` coefficients.
        auto coeffs = coeffs_type::seeded_random(n + 1, seed);
 
        // If `n > 0` we want the coefficient of `x^n` to be one for sure. If `n == 0` then a random 0/1 is fine.
        if (n > 0) coeffs.set(n);
        return BitPoly{std::move(coeffs)};
    }


    constexpr usize degree() const { return m_coeffs.last_set().value_or(0); }

    constexpr usize size() const { return m_coeffs.size(); }

    constexpr bool is_zero() const { return m_coeffs.none(); }

    constexpr bool is_non_zero() const { return m_coeffs.any(); }

    constexpr bool is_one() const { return degree() == 0 && m_coeffs.size() >= 1 && m_coeffs.get(0); }

    constexpr bool is_constant() const { return degree() == 0; }

    constexpr bool is_monic() const { return m_coeffs.trailing_zeros() == 0; }

    constexpr bool is_empty() const { return m_coeffs.is_empty(); }


    constexpr bool operator[](usize i) const { return m_coeffs[i]; }

    constexpr BitRef<coeffs_type> operator[](usize i) { return m_coeffs[i]; }

    constexpr const coeffs_type& coefficients() const { return m_coeffs; }

    constexpr coeffs_type& coefficients() { return m_coeffs; }

    constexpr void copy_coefficients(const coeffs_type& coeffs) { m_coeffs = coeffs; }

    constexpr void move_coefficients(coeffs_type&& coeffs) { m_coeffs = std::move(coeffs); }


    constexpr BitPoly& clear() {
        m_coeffs.clear();
        return *this;
    }

    constexpr BitPoly& resize(usize n) {
        m_coeffs.resize(n);
        return *this;
    }

    constexpr BitPoly& shrink_to_fit() {
        m_coeffs.shrink_to_fit();
        return *this;
    }

    constexpr BitPoly& make_monic() {
        if (is_non_zero()) m_coeffs.resize(degree() + 1);
        return *this;
    }


    constexpr BitPoly& operator+=(const BitPoly& rhs) {
        // Edge case.
        if (rhs.is_zero()) return *this;
 
        // Another edge case.
        if (is_zero()) {
            *this = rhs;
            return *this;
        }
 
        // Perhaps we need to get bigger to accommodate the rhs? Note that added coefficients are zeros.
        auto rhs_degree = rhs.degree();
        if (m_coeffs.size() < rhs_degree + 1) m_coeffs.resize(rhs_degree + 1);
 
        // Add in the active rhs words.
        for (auto i = 0uz; i < rhs.monic_word_count(); ++i) {
            m_coeffs.set_word(i, m_coeffs.word(i) ^ rhs.m_coeffs.word(i));
        }
        return *this;
    }

    constexpr BitPoly& operator-=(const BitPoly& rhs) { return operator+=(rhs); }

    constexpr BitPoly& operator*=(const BitPoly& rhs) {
        // Edge cases: zero BitPolys.
        if (rhs.is_zero()) return clear();
        if (is_zero()) return *this;
 
        // Edge cases: one BitPolys.
        if (rhs.is_one()) return *this;
        if (is_one()) {
            *this = rhs;
            return *this;
        }
 
        // Generally we pass the work to the convolution method for bit-vectors.
        *this = BitPoly{std::move(convolve(m_coeffs, rhs.m_coeffs))};
        return *this;
    }

    constexpr BitPoly<Word> operator+(const BitPoly<Word>& rhs) const {
        // Avoid unnecessary resizing by adding the smaller degree BitPoly to the larger one ...
        if (degree() >= rhs.degree()) {
            BitPoly<Word> result{*this};
            result += rhs;
            return result;
        } else {
            BitPoly<Word> result{rhs};
            result += *this;
            return result;
        }
    }

    constexpr BitPoly<Word> operator-(const BitPoly<Word>& rhs) const {
        // Subtraction is identical to addition in GF(2).
        return operator+(rhs);
    }

    constexpr BitPoly<Word> operator*(const BitPoly<Word>& rhs) const {
        BitPoly<Word> result{*this};
        result *= rhs;
        return result;
    }


    constexpr void square_into(BitPoly& dst) const {
        // Edge case: any constant BitPoly.
        if (is_constant()) {
            dst = *this;
            return;
        }
 
        // In GF(2) if p(x) = a + bx + cx^2 + ... then p(x)^2 = a^2 + b^2x^2 + c^2x^4 + ...
        // This identity means we can use the `riffle` method to square the BitPoly.
        m_coeffs.riffle_into(dst.m_coeffs);
    }

    constexpr BitPoly squared() const {
        BitPoly dst;
        square_into(dst);
        return dst;
    }

    constexpr BitPoly& times_x_to_the(usize n) {
        auto new_degree = degree() + n;
        auto new_size = new_degree + 1;
        if (m_coeffs.size() < new_size) { m_coeffs.resize(new_size); }
        m_coeffs >>= n;
        return *this;
    }


    constexpr bool operator()(bool x) const {
        // Edge case: the zero BitPoly.
        if (is_zero()) { return false; }
 
        // Edge case: `x = false` which is the same as `x = 0` & we always have p(0) = p_0.
        if (!x) { return m_coeffs.get(0); }
 
        // We are evaluating the BitPoly at `x = true` which is the same as `x = 1`: p(1) = p_0 + p_1 + p_2 + ...
        Word sum = 0;
        for (auto i = 0uz; i < m_coeffs.words(); ++i) sum ^= m_coeffs.word(i);
        return gf2::count_ones(sum) % 2 == 1;
    }

    constexpr BitMat<Word> operator()(const BitMat<Word>& M) const {
        // The bit-matrix argument must be square.
        gf2_assert(M.is_square(), "Matrix must be square -- not {} x {}!", M.rows(), M.cols());
 
        // The returned bit-matrix will be n x n.
        auto n = M.rows();
 
        // Edge case: If the polynomial is zero then the return value is the n x n zero matrix.
        if (is_zero()) return BitMat<Word>{n, n};
 
        // Otherwise we start with the polynomial sum being the n x n identity matrix.
        auto result = BitMat<Word>::identity(n);
 
        // Work backwards a la Horner ...
        auto d = degree();
        while (d > 0) {
 
            // Always multiply ...
            result = M * result;
 
            // Add the identity to the sum if the corresponding polynomial coefficient is 1.
            if (m_coeffs[d - 1]) result.add_identity();
 
            // And count down ...
            d--;
        }
 
        return result;
    }


    BitPoly reduce_x_to_the(usize n, bool n_is_log2 = false) const {
        // Error check: anything mod 0 is not defined.
        if (is_zero()) throw std::invalid_argument("... mod P(x) is not defined for P(x) := 0.");
 
        // Edge case: anything mod 1 = 0.
        if (is_one()) return BitPoly<Word>::zero();
 
        // Edge case: x^0 = 1 and 1 mod P(x) = 1 for any P(x) != 1 (we already handled the case where P(x) := 1).
        if (n == 0 && !n_is_log2) return BitPoly<Word>::one();
 
        // The BitPoly P(x) is non-zero so can be written as P(x) = x^d + p(x) where degree(p) < d.
        auto d = degree();
 
        // Edge case: P(x) = x + c where c is a constant so x = P(x) + c (subtraction in GF(2) is the same as addition).
        // Then x^e = (P(x) + c)^e = terms in powers of P(x) + c^e. Hence, x^e mod P(x) = c^e = c.
        if (d == 1) return BitPoly<Word>::constant(m_coeffs.get(1));
 
        // We can write p(x) = p_0 + p_1 x + ... + p_{d-1} x^{d-1}. All that matters are those coefficients.
        auto p = m_coeffs.sub(0, d);
 
        // A lambda to perform: q(x) <- x*q(x) mod P(x) where degree(q) < d.
        // The lambda works on the coefficients of q(x) passed as a bit-vector q.
        auto times_x_step = [&](auto& q) {
            bool add_p = q[d - 1];
            q >>= 1;
            if (add_p) q ^= p;
        };
 
        // Iteratively precompute x^{d+i} mod P(x) for i = 0, 1, ..., d-1 starting with x^d mod P(x) ~ p.
        // We store all the bit-vectors in a standard `std::vector` of length d.
        std::vector<coeffs_type> power_mod(d, coeffs_type{d});
        power_mod[0] = p;
        for (auto i = 1uz; i < d; ++i) {
            power_mod[i] = power_mod[i - 1];
            times_x_step(power_mod[i]);
        }
 
        // Some workspace we use/reuse below in order to minimize allocations/deallocations.
        coeffs_type s{2 * d}, h{d};
 
        // lambda to perform: q(x) <- q(x)^2 mod P(x) where degree(q) < d.
        // The lambda works on the coefficients of q(x) passed as a bit-vector q.
        auto square_step = [&](auto& q) {
            // Square q(x) storing the result in workspace `s`.
            q.riffle_into(s);
 
            // Split s(x) = l(x) + x^d * h(x) where l(x) & h(x) are both of degree less than d.
            // We reuse q(x) for l(x).
            s.split_at(d, q, h);
 
            // s(x) = q(x) + h(x) so s(x) mod P(x) = q(x) + h(x) mod P(x) which we handle term by term.
            // If h(x) != 0 then at most every second term in h(x) is 1 (nature of BitPoly squares in GF(2)).
            if (auto h_first = h.first_set()) {
                auto h_last = h.last_set();
                for (auto i = *h_first; i <= *h_last; i += 2)
                    if (h[i]) q ^= power_mod[i];
            }
        };
 
        // Our return value r(x) := x^e mod P(x) has degree < d: r(x) = r_0 + r_1 x + ... + r_{d-1} x^{d-1}.
        coeffs_type r{d};
 
        // If `n_is_log2` is `true`, we are reducing x^(2^n) mod P(x) which is done iteratively by squaring.
        if (n_is_log2) {
            // Note that we already handled edge case where P(x) = x + c above.
            // Start with r(x) = x mod P(x) -> x^2 mod P(x) -> x^4 mod P(x) ...
            r[1] = true;
            for (auto i = 0uz; i < n; ++i) square_step(r);
            return BitPoly{std::move(r)};
        }
 
        // Small exponent case: n < d => x^n mod P(x) = x^n.
        if (n < d) return BitPoly::x_to_the(n);
 
        // Matching exponent case: n = d => x^n mod P(x) = x^d mod P(x) = p(x).
        if (n == d) return BitPoly{std::move(p)};
 
        // General case: n > d is handled by a square & multiply algorithm.
        usize n_bit = std::bit_floor(n);
 
        // Start with r(x) = x mod P(x) which takes care of the most significant binary digit in n.
        // TODO: We could start a bit further along with a higher power r(x) := x^? mod P(x) where ? < n but > 1.
        r[1] = 1;
        n_bit >>= 1;
 
        // And off we go ...
        while (n_bit) {
 
            // Always do a square step and then a times_x step if necessary (i.e. if current bit in N is set).
            square_step(r);
            if (n & n_bit) times_x_step(r);
 
            // On to the next bit position in n.
            n_bit >>= 1;
        }
 
        // Made it.
        return BitPoly{std::move(r)};
    }


    std::string to_string(std::string_view var = "x") const {
        // Edge case: the zero BitPoly.
        if (is_zero()) return "0";
 
        // Otherwise we construct the string ...
        std::ostringstream ss;
 
        // All terms other than the first one are preceded by a "+"
        bool first_term = true;
        for (auto i = 0uz; i <= degree(); ++i) {
            if (m_coeffs.get(i)) {
                if (i == 0) {
                    ss << "1";
                } else {
                    if (!first_term) ss << " + ";
                    if (i == 1)
                        ss << var;
                    else
                        ss << var << "^" << i;
                }
                first_term = false;
            }
        }
        return ss.str();
    }

    std::string to_full_string(std::string_view var = "x") const {
        // Edge case: the zero BitPoly.
        if (is_empty()) return "0";
 
        // Otherwise we construct the string ...
        std::ostringstream ss;
 
        // Start with the first term ...
        if (m_coeffs.get(0))
            ss << "1";
        else
            ss << "0";
 
        // Now all the other terms, each preceded by a "+"
        for (auto i = 1uz; i < size(); ++i) {
            ss << " + ";
            if (m_coeffs.get(i) == false) ss << "0";
            if (i == 1)
                ss << var;
            else
                ss << var << "^" << i;
        }
        return ss.str();
    }

    std::ostream& operator<<(std::ostream& s) const { return s << to_string(); }


    friend constexpr bool operator==(const BitPoly& lhs, const BitPoly& rhs) {
        // Edge case.
        if (&lhs == &rhs) return true;
 
        // Check the active words for equality.
        auto count = lhs.monic_word_count();
        if (rhs.monic_word_count() != count) return false;
        for (auto i = 0uz; i < count; ++i)
            if (lhs.m_coeffs.word(i) != rhs.m_coeffs.word(i)) return false;
 
        // Made it
        return true;
    }

 
private:
    // Returns the number of "active" words underlying the coefficient bit-vector.
    // There may be high order trailing zero coefficients that have no real impact on most bit-polynomial calculations.
    // This method returns the number of "words" that matter in the coefficient bit-vector store of words.
    constexpr usize monic_word_count() const {
        if (auto deg = m_coeffs.last_set()) return word_index<Word>(*deg) + 1;
        return 0;
    }
};
 
} // namespace gf2
 
// --------------------------------------------------------------------------------------------------------------------
// Specialises `std::formatter` to handle bit-polynomials ...
// -------------------------------------------------------------------------------------------------------------------

template<gf2::unsigned_word Word>
struct std::formatter<gf2::BitPoly<Word>> {

    constexpr auto parse(const std::format_parse_context& ctx) {
        auto it = ctx.begin();
        if (*it != '}') m_var.clear();
        while (it != ctx.end() && *it != '}') {
            m_var += *it;
            ++it;
        }
        return it;
    }

    template<class FormatContext>
    auto format(const gf2::BitPoly<Word>& rhs, FormatContext& ctx) const {
        return std::format_to(ctx.out(), "{}", rhs.to_string(m_var));
    }
 
    std::string m_var = "x";
};