cp-book
ecnerwala's competitive programming library
#include "fft/engines/algebras.hpp"
Depends on
fft/common.hppRequired by
Code
| Coverage | Exec / Excl / Total | |
|---|---|---|
| Lines | 64.6% | 113 / 0 / 175 |
| Functions | 100.0% | 129 / 0 / 129 |
| Branches | 99.5% | 405 / 0 / 407 |
| Full report |
#pragma once
#include <array>
#include <concepts>
#include <cstddef>
#include <span>
#include <utility>
#include "fft/engine.hpp"
namespace ecnerwala::fft::engines {
// Small NxN matrix over num, row-major
// ==== wrapper engines ====
template <typename num, int N> struct mat {
std::array<num, size_t(N) * N> a{};
num& operator[](std::array<int, 2> rc) { return a[size_t(rc[0]) * N + rc[1]]; }
const num& operator[](std::array<int, 2> rc) const { return a[size_t(rc[0]) * N + rc[1]]; }
num* data() { return a.data(); }
const num* data() const { return a.data(); }
mat& operator+=(const mat& o) { for (int i = 0; i < N*N; i++) a[i] += o.a[i]; return *this; }
friend mat operator+(mat x, const mat& y) { x += y; return x; }
mat& operator-=(const mat& o) { for (int i = 0; i < N*N; i++) a[i] -= o.a[i]; return *this; }
friend mat operator-(mat x, const mat& y) { x -= y; return x; }
friend mat operator*(const mat& x, const mat& y) {
mat r;
for (int i = 0; i < N; i++) for (int k = 0; k < N; k++) for (int j = 0; j < N; j++)
r[{i, j}] += x[{i, k}] * y[{k, j}];
return r;
}
mat& operator*=(const mat& o) { return *this = *this * o; }
friend bool operator==(const mat&, const mat&) = default;
};
// Truncated polynomial mod x^N over num
template <typename num, int N> struct trunc_series {
std::array<num, size_t(N)> a{};
num& operator[](int i) { return a[size_t(i)]; }
const num& operator[](int i) const { return a[size_t(i)]; }
num* data() { return a.data(); }
const num* data() const { return a.data(); }
trunc_series& operator+=(const trunc_series& o) { for (int i = 0; i < N; i++) a[i] += o.a[i]; return *this; }
friend trunc_series operator+(trunc_series x, const trunc_series& y) { x += y; return x; }
trunc_series& operator-=(const trunc_series& o) { for (int i = 0; i < N; i++) a[i] -= o.a[i]; return *this; }
friend trunc_series operator-(trunc_series x, const trunc_series& y) { x -= y; return x; }
friend trunc_series operator*(const trunc_series& x, const trunc_series& y) {
trunc_series r;
for (int i = 0; i < N; i++) for (int j = 0; j < N - i; j++) r[i + j] += x[i] * y[j];
return r;
}
trunc_series& operator*=(const trunc_series& o) { return *this = *this * o; }
friend bool operator==(const trunc_series&, const trunc_series&) = default;
};
// componentwise
// These are "componentwise engines" which model free modules/algebras over the underlying ring.
// Matrices are the canonical example: we can take entry-wise transforms, then multiply/add in transformed space.
//
// We start with a shared `componentwise` base class which handles all linear ops, i.e. not mul.
//
// If the underlying has unit_scale = 1, we may need to avoid accumulation of sums in transformed space;
// then, the product and the transformed data may have different dimensions.
// We'll represent this by an array Ofs of prefix offsets mapping each input/transform-space dimension to a range of product-space dimensions.
// Specifically, out[c] = sum prod[Ofs[c]:Ofs[c+1]]
template <int L> constexpr std::array<int, size_t(L) + 1> componentwise_iota = [] {
std::array<int, size_t(L) + 1> r{};
for (int i = 0; i <= L; i++) r[size_t(i)] = i;
return r;
}();
template <engine E, typename V, int L, std::array<int, size_t(L) + 1> Ofs = componentwise_iota<L>>
struct componentwise {
using S = typename E::value_type;
using value_type = V;
static constexpr int P = Ofs[size_t(L)]; // total product components
static constexpr int unit_scale = E::unit_scale;
template <int A = unit_scale> struct transformed_t {
std::array<typename E::template transformed_t<A>, size_t(L)> t;
int size() const { return t[0].size(); }
transformed_t() = default;
template <int A2> requires (A2 != A) explicit(A2 > A) transformed_t(transformed_t<A2>&& o) {
for (int c = 0; c < L; c++)
t[c] = typename E::template transformed_t<A>(std::move(o.t[c]));
}
};
using transformed = transformed_t<>;
// TODO: if E::product_t == E::transformed_t, mirror that here
template <int K> struct product_t {
std::array<typename E::template product_t<K>, size_t(P)> t;
int size() const { return t[0].size(); }
product_t() = default;
template <int K2> requires (K2 != K) explicit(K2 > K) product_t(product_t<K2>&& o) {
for (int c = 0; c < P; c++)
t[c] = typename E::template product_t<K>(std::move(o.t[c]));
}
};
static transformed transform(std::span<const V> a, int n) {
transformed r;
auto buf = buffer_pool<S>::get(sz(a));
for (int c = 0; c < L; c++) {
for (int i = 0; i < sz(a); i++) buf[i] = a[i].data()[c];
r.t[c] = E::transform(std::span<const S>(buf.span()), n);
}
return r;
}
static void extend_to(transformed& t, int m, std::span<const V> coeffs) {
if (t.size() >= m) return;
auto buf = buffer_pool<S>::get(sz(coeffs));
for (int c = 0; c < L; c++) {
for (int i = 0; i < sz(coeffs); i++) buf[i] = coeffs[i].data()[c];
E::extend_to(t.t[c], m, std::span<const S>(buf.span()));
}
}
template <int A> static transformed_t<A> downsample(const transformed_t<A>& t, int n, bool odd) {
transformed_t<A> r;
for (int c = 0; c < L; c++) r.t[c] = E::downsample(t.t[c], n, odd);
return r;
}
template <int K> static product_t<K> downsample(const product_t<K>& p, int n, bool odd) {
product_t<K> r;
for (int c = 0; c < P; c++) r.t[c] = E::downsample(p.t[c], n, odd);
return r;
}
template <int A> static transformed_t<A> negate_arg(const transformed_t<A>& t, int n) {
transformed_t<A> r;
for (int c = 0; c < L; c++) r.t[c] = E::negate_arg(t.t[c], n);
return r;
}
template <int A, int B> static transformed_t<A + B> add(transformed_t<A>&& a, const transformed_t<B>& b) {
transformed_t<A + B> r;
for (int c = 0; c < L; c++) r.t[c] = E::add(std::move(a.t[c]), b.t[c]);
return r;
}
template <int K1, int K2> static product_t<K1 + K2> add(product_t<K1>&& a, product_t<K2>&& b) {
product_t<K1 + K2> r;
for (int c = 0; c < P; c++) r.t[c] = E::add(std::move(a.t[c]), std::move(b.t[c]));
return r;
}
template <int K, typename Op = assign_op> static void finish(product_t<K>&& p, std::span<V> out, Op op = {}) {
auto buf = buffer_pool<S>::get(sz(out));
auto emit = [&](std::span<V> dst) {
for (int c = 0; c < L; c++) {
E::finish(std::move(p.t[Ofs[size_t(c)]]), buf.span());
for (int j = Ofs[size_t(c)] + 1; j < Ofs[size_t(c) + 1]; j++)
E::finish(std::move(p.t[j]), buf.span(), add_op{});
for (int i = 0; i < sz(dst); i++) dst[i].data()[c] = buf[i];
}
};
// Op must see each out element whole, exactly once, so compose
// non-assign ops through an element buffer.
if constexpr (std::same_as<Op, assign_op>) {
emit(out);
} else {
auto vbuf = buffer_pool<V>::get(sz(out));
emit(vbuf.span());
for (int i = 0; i < sz(out); i++) op(out[i], vbuf.span()[i]);
}
}
};
// Convolve mat<N> (NxN matrices), with accumulation in product space
template <engine E, int N>
struct matrix : componentwise<E, mat<typename E::value_type, N>, N * N> {
using base = componentwise<E, mat<typename E::value_type, N>, N * N>;
static constexpr bool commutative = false;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K * N>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
// right fold over k so a tracked inner engine's per-addend types line up
template <int A, int B, int k = 0>
static auto entry(const transformed_t<A>& a, const transformed_t<B>& b, int r, int c, int n) {
auto e = E::mul(a.t[size_t(r) * N + k], b.t[size_t(k) * N + c], n);
if constexpr (k + 1 == N) return e;
else return E::add(std::move(e), entry<A, B, k + 1>(a, b, r, c, n));
}
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++)
p.t[size_t(r) * N + c] = entry<A, B>(a, b, r, c, n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2, int k = 0>
static auto entry2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int r, int c, int n
) {
auto e = E::mul2(
a1.t[size_t(r) * N + k], b1.t[size_t(k) * N + c],
a2.t[size_t(r) * N + k], b2.t[size_t(k) * N + c],
n
);
if constexpr (k + 1 == N) return e;
else return E::add(std::move(e), entry2<A1, B1, A2, B2, k + 1>(a1, b1, a2, b2, r, c, n));
}
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++)
p.t[size_t(r) * N + c] = entry2<A1, B1, A2, B2>(a1, b1, a2, b2, r, c, n);
return p;
}
};
// Convolve trunc_series<num, N> (power series truncated at N), with accumulation in product space
template <engine E, int N>
struct trunc : componentwise<E, trunc_series<typename E::value_type, N>, N> {
using base = componentwise<E, trunc_series<typename E::value_type, N>, N>;
static constexpr bool commutative = E::commutative;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K * N>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B, int s, int i = 0>
static auto entry(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
auto e = E::mul(a.t[size_t(i)], b.t[size_t(s - i)], n);
if constexpr (i == s) return e;
else return E::add(std::move(e), entry<A, B, s, i + 1>(a, b, n));
}
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
[&]<size_t... s_>(std::index_sequence<s_...>) {
((p.t[s_] = entry<A, B, int(s_)>(a, b, n)), ...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2, int s, int i = 0>
static auto entry2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
auto e = E::mul2(a1.t[size_t(i)], b1.t[size_t(s - i)], a2.t[size_t(i)], b2.t[size_t(s - i)], n);
if constexpr (i == s) return e;
else return E::add(std::move(e), entry2<A1, B1, A2, B2, s, i + 1>(a1, b1, a2, b2, n));
}
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
[&]<size_t... s_>(std::index_sequence<s_...>) {
((p.t[s_] = entry2<A1, B1, A2, B2, int(s_)>(a1, b1, a2, b2, n)), ...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
};
// Stable variants of the wrapper engines: do not accumulate in product space.
// This costs an extra log factor.
template <int N> constexpr std::array<int, size_t(N) * N + 1> matrix_stable_ofs = [] {
std::array<int, size_t(N) * N + 1> r{};
for (int i = 0; i <= N * N; i++) r[size_t(i)] = i * N;
return r;
}();
template <engine E, int N>
struct matrix_stable
: componentwise<E, mat<typename E::value_type, N>, N * N, matrix_stable_ofs<N>> {
using base = componentwise<E, mat<typename E::value_type, N>, N * N, matrix_stable_ofs<N>>;
static constexpr bool commutative = false;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
// entry (r, c)'s k-th addend a(r,k)*b(k,c), grouped per the offsets
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++) for (int k = 0; k < N; k++)
p.t[(size_t(r) * N + c) * N + k] = E::mul(a.t[size_t(r) * N + k], b.t[size_t(k) * N + c], n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++) for (int k = 0; k < N; k++)
p.t[(size_t(r) * N + c) * N + k] = E::mul2(
a1.t[size_t(r) * N + k], b1.t[size_t(k) * N + c],
a2.t[size_t(r) * N + k], b2.t[size_t(k) * N + c],
n
);
return p;
}
};
template <int N> constexpr std::array<int, size_t(N) + 1> trunc_series_stable_ofs = [] {
std::array<int, size_t(N) + 1> r{};
for (int i = 0; i <= N; i++) r[size_t(i)] = i * (i + 1) / 2;
return r;
}();
template <engine E, int N>
struct trunc_stable
: componentwise<E, trunc_series<typename E::value_type, N>, N, trunc_series_stable_ofs<N>> {
using base = componentwise<E, trunc_series<typename E::value_type, N>, N, trunc_series_stable_ofs<N>>;
static constexpr bool commutative = E::commutative;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
for (int s = 0; s < N; s++) for (int i = 0; i <= s; i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)] + i)] = E::mul(a.t[size_t(i)], b.t[size_t(s - i)], n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int s = 0; s < N; s++) for (int i = 0; i <= s; i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)] + i)] = E::mul2(
a1.t[size_t(i)], b1.t[size_t(s - i)],
a2.t[size_t(i)], b2.t[size_t(s - i)],
n
);
return p;
}
};
/* namespace ecnerwala::fft::engines */ }
#include <array>
#include <concepts>
#include <cstddef>
#include <span>
#include <utility>
#include <type_traits>
#include <algorithm>
#include <iterator>
#include <vector>
#line 2 "src/fft/engines/algebras.hpp"
#line 8 "src/fft/engines/algebras.hpp"
#line 2 "src/fft/engine.hpp"
#line 7 "src/fft/engine.hpp"
#line 2 "src/fft/common.hpp"
#line 8 "src/fft/common.hpp"
/**
* Author: Andrew He
* Source: http://neerc.ifmo.ru/trains/toulouse/2017/fft2.pdf
* Papers about accuracy: http://www.daemonology.net/papers/fft.pdf, http://www.cs.berkeley.edu/~fateman/papers/fftvsothers.pdf
* For integers rounding works if $(|a| + |b|)\max(a, b) < \mathtt{\sim} 10^9$, or in theory maybe $10^6$.
*
* Abstraction layers:
* fft_core<num> FFT itself and other ops on rings with 2^k-th roots of unity. We use bit-reversed indexing in the frequency domain.
*
* engines Engines for packing/unpacking arbitrary rings for convolution (the `engine` concept).
* Still expose (opaque) transform-domain objects for caching/fusion.
*
* multiply layer Wrappers for convolving bounded sequences: track length/truncation.
*
* value types series::vec<E, exact> - R[[x]]
* series::exact<E> - exact (finite-support) power series
* series::trunc<E> - truncated prefix of an (infinite) power series
*
* polynomials - R[x]. Under x -> 1/x a polynomial becomes a Laurent polynomial in 1/x;
* shifting by x^{deg P} (reversal) lands it in R[[x]], and we store that exact series.
* poly::vec<E> - polynomial type, supporting natural indexing
* poly::form<E> - finite-support linear forms, via the pairing <P, S> = [x^0] P(1/x) S(x)
* a linear form is one side of this pairing, applied to the other
*
* online_multiplier<E> - online (relaxed) multiplication of 2 sequences in n log^2 n time
* ap_sampled_poly<E> - a polynomial stored as its evaluations on an arithmetic progression
*/
namespace ecnerwala {
template<class T> int sz(T&& arg) { using std::size; return int(size(std::forward<T>(arg))); }
inline int nextPow2(int s) { return 1 << (s > 1 ? 32 - __builtin_clz(s-1) : 0); }
namespace fft {
using std::swap;
using std::vector;
using std::min;
using std::max;
// Reusable scratch buffers. Not thread-safe by default: this is deliberately plain
// static storage so single-threaded programs pay no TLS indirection; define
// ECNERWALA_FFT_POOL_STORAGE to `thread_local` for multithreaded use.
#ifndef ECNERWALA_FFT_POOL_STORAGE
#define ECNERWALA_FFT_POOL_STORAGE
#endif
template <typename T> struct buffer_pool {
static inline ECNERWALA_FFT_POOL_STORAGE std::vector<std::vector<T>> free_list;
struct handle {
std::vector<T> v;
explicit handle(int n) {
if (!free_list.empty()) {
v = std::move(free_list.back());
free_list.pop_back();
}
v.assign(n, T());
}
handle(const handle&) = delete;
handle& operator=(const handle&) = delete;
handle(handle&& o) noexcept : v(std::move(o.v)) {}
~handle() { if (v.capacity()) free_list.push_back(std::move(v)); }
T& operator[](int i) { return v[i]; }
operator std::span<T>() { return std::span<T>(v); }
std::span<T> span() { return std::span<T>(v); }
};
static handle get(int n) { return handle(n); }
};
/* namespace fft */ }
/* namespace ecnerwala */ }
#line 9 "src/fft/engine.hpp"
namespace ecnerwala::fft {
// ==== engine concept ====
// Output operations for the finish step to express arbitrary fusion into the output buffer.
struct assign_op { template <typename T> void operator()(T& d, T v) const { d = v; } };
struct add_op { template <typename T> void operator()(T& d, T v) const { d += v; } };
struct sub_op { template <typename T> void operator()(T& d, T v) const { d -= v; } };
struct add_twice_op { template <typename T> void operator()(T& d, T v) const { d += v + v; } };
// `engine` contract
// engine represents a way of packing/unpacking sequences over an arbitrary ring into FFT-style transforms.
// We expect transforms/products of transforms to be linear but potentially lossy/imprecise, so we'll track precision
// at compile-time as a template parameter.
//
// E::value_type The ring we operate over
// E::unit_scale 0 or 1 depending on whether there's error that can accumulate
// E::commutative A marker for whether the ring is commutative
//
// transformed_t<A> The transform of a sequence. This object owns its data buffer.
// product_t<A> The product of 2 transforms. May equal transformed_t, particularly when unit_scale = 0.
//
// transformed alias for transformed_t<unit_scale>
// product alias for product_t<unit_scale>
//
// The basic multiplication API is
// transform(span<const value_type> in, int n) -> transformed_t<unit_scale>
// mul(transformed_t<A>, transformed_t<B>, int n) -> product_t<A*B>
// mul2(a1, b1, a2, b2, int n) -> product_t<A1*B1 + A2*B2>, computing a1*b1 + a2*b2 in one pass
// finish(product_t<A>, span<value_type>& out, Op) -> void
//
// Input span can be length up to 2n.
// Output spans can be length up to n; only the prefix that exists is filled.
// finish applies Op exactly once per out element, in index order, to
// value_type targets (so ops may be stateful).
// Transforms can be longer than necessary, and only the relevant prefix is used.
//
// For non-exact engines, there's some subtlety in whether we wrap before or after packing.
// We will choose to wrap *after* packing, which hurts error bounds but makes the prefix condition more uniform.
//
// Additionally, we have APIs to take advantage of linearity in both transformed and product space:
// add(transformed_t<A>, transformed_t<B>) -> transformed_t<A+B>
// add(product_t<K1>, product_t<K2>) -> product_t<K1+K2>
//
// Finally, we expose some additional fast-transform optimization paths.
// extend_to only operates on transformed_t<unit_scale>; the others are scale-generic.
// downsample is also defined on product_t (halving before finish saves inverse-transform work).
// extend_to build (if empty) or grow a transform to size m by repeated doubling; feed the same coefficients (sz <= 2m) every time
// (each doubling step reads only the coefficients that fit, so zero-padded buffers are fine:
// coefficients past twice the existing transform's size must be zero, or it couldn't be a prefix)
// downsample compute the half-sized transform/product of just the even (odd = false) or odd terms of the input
// negate_arg size n transform of A(-x)
template <typename E>
concept engine = requires(
std::span<const typename E::value_type> in,
std::span<typename E::value_type> out,
typename E::transformed& t,
const typename E::transformed& ct,
typename E::product& p,
const typename E::product& cp,
int n
) {
typename E::value_type;
{ E::transform(in, n) } -> std::same_as<typename E::transformed>;
{ ct.size() } -> std::same_as<int>;
E::extend_to(t, n, in);
{ E::downsample(ct, n, false) } -> std::same_as<typename E::transformed>;
{ E::downsample(cp, n, false) } -> std::same_as<typename E::product>;
{ E::negate_arg(ct, n) } -> std::same_as<typename E::transformed>;
{ E::mul(ct, ct, n) } -> std::same_as<typename E::product>;
{ E::sq(ct, n) } -> std::same_as<typename E::product>;
{ E::mul2(ct, ct, ct, ct, n) } -> std::same_as<typename E::template product_t<2 * E::unit_scale>>;
E::finish(std::move(p), out);
E::finish(std::move(p), out, add_op{});
E::finish(E::add(std::move(p), std::move(p)), out);
{ E::add(E::transform(in, n), ct) } -> std::same_as<typename E::template transformed_t<2 * E::unit_scale>>;
{ E::add(std::move(p), std::move(p)) } -> std::same_as<typename E::template product_t<2 * E::unit_scale>>;
requires std::same_as<std::remove_cvref_t<decltype(E::commutative)>, bool>;
requires std::same_as<std::remove_cvref_t<decltype(E::unit_scale)>, int>;
};
// Constrains two engine-parameterized value types to share the same engine.
template <typename A, typename B>
concept same_engine = std::same_as<typename A::engine_t, typename B::engine_t>;
// short spelling for E::transformed at use sites
template <engine E> using transformed = typename E::transformed;
/* namespace ecnerwala::fft */ }
#line 10 "src/fft/engines/algebras.hpp"
namespace ecnerwala::fft::engines {
// Small NxN matrix over num, row-major
// ==== wrapper engines ====
template <typename num, int N> struct mat {
std::array<num, size_t(N) * N> a{};
num& operator[](std::array<int, 2> rc) { return a[size_t(rc[0]) * N + rc[1]]; }
const num& operator[](std::array<int, 2> rc) const { return a[size_t(rc[0]) * N + rc[1]]; }
num* data() { return a.data(); }
const num* data() const { return a.data(); }
mat& operator+=(const mat& o) { for (int i = 0; i < N*N; i++) a[i] += o.a[i]; return *this; }
friend mat operator+(mat x, const mat& y) { x += y; return x; }
mat& operator-=(const mat& o) { for (int i = 0; i < N*N; i++) a[i] -= o.a[i]; return *this; }
friend mat operator-(mat x, const mat& y) { x -= y; return x; }
friend mat operator*(const mat& x, const mat& y) {
mat r;
for (int i = 0; i < N; i++) for (int k = 0; k < N; k++) for (int j = 0; j < N; j++)
r[{i, j}] += x[{i, k}] * y[{k, j}];
return r;
}
mat& operator*=(const mat& o) { return *this = *this * o; }
friend bool operator==(const mat&, const mat&) = default;
};
// Truncated polynomial mod x^N over num
template <typename num, int N> struct trunc_series {
std::array<num, size_t(N)> a{};
num& operator[](int i) { return a[size_t(i)]; }
const num& operator[](int i) const { return a[size_t(i)]; }
num* data() { return a.data(); }
const num* data() const { return a.data(); }
trunc_series& operator+=(const trunc_series& o) { for (int i = 0; i < N; i++) a[i] += o.a[i]; return *this; }
friend trunc_series operator+(trunc_series x, const trunc_series& y) { x += y; return x; }
trunc_series& operator-=(const trunc_series& o) { for (int i = 0; i < N; i++) a[i] -= o.a[i]; return *this; }
friend trunc_series operator-(trunc_series x, const trunc_series& y) { x -= y; return x; }
friend trunc_series operator*(const trunc_series& x, const trunc_series& y) {
trunc_series r;
for (int i = 0; i < N; i++) for (int j = 0; j < N - i; j++) r[i + j] += x[i] * y[j];
return r;
}
trunc_series& operator*=(const trunc_series& o) { return *this = *this * o; }
friend bool operator==(const trunc_series&, const trunc_series&) = default;
};
// componentwise
// These are "componentwise engines" which model free modules/algebras over the underlying ring.
// Matrices are the canonical example: we can take entry-wise transforms, then multiply/add in transformed space.
//
// We start with a shared `componentwise` base class which handles all linear ops, i.e. not mul.
//
// If the underlying has unit_scale = 1, we may need to avoid accumulation of sums in transformed space;
// then, the product and the transformed data may have different dimensions.
// We'll represent this by an array Ofs of prefix offsets mapping each input/transform-space dimension to a range of product-space dimensions.
// Specifically, out[c] = sum prod[Ofs[c]:Ofs[c+1]]
template <int L> constexpr std::array<int, size_t(L) + 1> componentwise_iota = [] {
std::array<int, size_t(L) + 1> r{};
for (int i = 0; i <= L; i++) r[size_t(i)] = i;
return r;
}();
template <engine E, typename V, int L, std::array<int, size_t(L) + 1> Ofs = componentwise_iota<L>>
struct componentwise {
using S = typename E::value_type;
using value_type = V;
static constexpr int P = Ofs[size_t(L)]; // total product components
static constexpr int unit_scale = E::unit_scale;
template <int A = unit_scale> struct transformed_t {
std::array<typename E::template transformed_t<A>, size_t(L)> t;
int size() const { return t[0].size(); }
transformed_t() = default;
template <int A2> requires (A2 != A) explicit(A2 > A) transformed_t(transformed_t<A2>&& o) {
for (int c = 0; c < L; c++)
t[c] = typename E::template transformed_t<A>(std::move(o.t[c]));
}
};
using transformed = transformed_t<>;
// TODO: if E::product_t == E::transformed_t, mirror that here
template <int K> struct product_t {
std::array<typename E::template product_t<K>, size_t(P)> t;
int size() const { return t[0].size(); }
product_t() = default;
template <int K2> requires (K2 != K) explicit(K2 > K) product_t(product_t<K2>&& o) {
for (int c = 0; c < P; c++)
t[c] = typename E::template product_t<K>(std::move(o.t[c]));
}
};
static transformed transform(std::span<const V> a, int n) {
transformed r;
auto buf = buffer_pool<S>::get(sz(a));
for (int c = 0; c < L; c++) {
for (int i = 0; i < sz(a); i++) buf[i] = a[i].data()[c];
r.t[c] = E::transform(std::span<const S>(buf.span()), n);
}
return r;
}
static void extend_to(transformed& t, int m, std::span<const V> coeffs) {
if (t.size() >= m) return;
auto buf = buffer_pool<S>::get(sz(coeffs));
for (int c = 0; c < L; c++) {
for (int i = 0; i < sz(coeffs); i++) buf[i] = coeffs[i].data()[c];
E::extend_to(t.t[c], m, std::span<const S>(buf.span()));
}
}
template <int A> static transformed_t<A> downsample(const transformed_t<A>& t, int n, bool odd) {
transformed_t<A> r;
for (int c = 0; c < L; c++) r.t[c] = E::downsample(t.t[c], n, odd);
return r;
}
template <int K> static product_t<K> downsample(const product_t<K>& p, int n, bool odd) {
product_t<K> r;
for (int c = 0; c < P; c++) r.t[c] = E::downsample(p.t[c], n, odd);
return r;
}
template <int A> static transformed_t<A> negate_arg(const transformed_t<A>& t, int n) {
transformed_t<A> r;
for (int c = 0; c < L; c++) r.t[c] = E::negate_arg(t.t[c], n);
return r;
}
template <int A, int B> static transformed_t<A + B> add(transformed_t<A>&& a, const transformed_t<B>& b) {
transformed_t<A + B> r;
for (int c = 0; c < L; c++) r.t[c] = E::add(std::move(a.t[c]), b.t[c]);
return r;
}
template <int K1, int K2> static product_t<K1 + K2> add(product_t<K1>&& a, product_t<K2>&& b) {
product_t<K1 + K2> r;
for (int c = 0; c < P; c++) r.t[c] = E::add(std::move(a.t[c]), std::move(b.t[c]));
return r;
}
template <int K, typename Op = assign_op> static void finish(product_t<K>&& p, std::span<V> out, Op op = {}) {
auto buf = buffer_pool<S>::get(sz(out));
auto emit = [&](std::span<V> dst) {
for (int c = 0; c < L; c++) {
E::finish(std::move(p.t[Ofs[size_t(c)]]), buf.span());
for (int j = Ofs[size_t(c)] + 1; j < Ofs[size_t(c) + 1]; j++)
E::finish(std::move(p.t[j]), buf.span(), add_op{});
for (int i = 0; i < sz(dst); i++) dst[i].data()[c] = buf[i];
}
};
// Op must see each out element whole, exactly once, so compose
// non-assign ops through an element buffer.
if constexpr (std::same_as<Op, assign_op>) {
emit(out);
} else {
auto vbuf = buffer_pool<V>::get(sz(out));
emit(vbuf.span());
for (int i = 0; i < sz(out); i++) op(out[i], vbuf.span()[i]);
}
}
};
// Convolve mat<N> (NxN matrices), with accumulation in product space
template <engine E, int N>
struct matrix : componentwise<E, mat<typename E::value_type, N>, N * N> {
using base = componentwise<E, mat<typename E::value_type, N>, N * N>;
static constexpr bool commutative = false;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K * N>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
// right fold over k so a tracked inner engine's per-addend types line up
template <int A, int B, int k = 0>
static auto entry(const transformed_t<A>& a, const transformed_t<B>& b, int r, int c, int n) {
auto e = E::mul(a.t[size_t(r) * N + k], b.t[size_t(k) * N + c], n);
if constexpr (k + 1 == N) return e;
else return E::add(std::move(e), entry<A, B, k + 1>(a, b, r, c, n));
}
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++)
p.t[size_t(r) * N + c] = entry<A, B>(a, b, r, c, n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2, int k = 0>
static auto entry2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int r, int c, int n
) {
auto e = E::mul2(
a1.t[size_t(r) * N + k], b1.t[size_t(k) * N + c],
a2.t[size_t(r) * N + k], b2.t[size_t(k) * N + c],
n
);
if constexpr (k + 1 == N) return e;
else return E::add(std::move(e), entry2<A1, B1, A2, B2, k + 1>(a1, b1, a2, b2, r, c, n));
}
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++)
p.t[size_t(r) * N + c] = entry2<A1, B1, A2, B2>(a1, b1, a2, b2, r, c, n);
return p;
}
};
// Convolve trunc_series<num, N> (power series truncated at N), with accumulation in product space
template <engine E, int N>
struct trunc : componentwise<E, trunc_series<typename E::value_type, N>, N> {
using base = componentwise<E, trunc_series<typename E::value_type, N>, N>;
static constexpr bool commutative = E::commutative;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K * N>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B, int s, int i = 0>
static auto entry(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
auto e = E::mul(a.t[size_t(i)], b.t[size_t(s - i)], n);
if constexpr (i == s) return e;
else return E::add(std::move(e), entry<A, B, s, i + 1>(a, b, n));
}
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
[&]<size_t... s_>(std::index_sequence<s_...>) {
((p.t[s_] = entry<A, B, int(s_)>(a, b, n)), ...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2, int s, int i = 0>
static auto entry2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
auto e = E::mul2(a1.t[size_t(i)], b1.t[size_t(s - i)], a2.t[size_t(i)], b2.t[size_t(s - i)], n);
if constexpr (i == s) return e;
else return E::add(std::move(e), entry2<A1, B1, A2, B2, s, i + 1>(a1, b1, a2, b2, n));
}
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
[&]<size_t... s_>(std::index_sequence<s_...>) {
((p.t[s_] = entry2<A1, B1, A2, B2, int(s_)>(a1, b1, a2, b2, n)), ...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
};
// Stable variants of the wrapper engines: do not accumulate in product space.
// This costs an extra log factor.
template <int N> constexpr std::array<int, size_t(N) * N + 1> matrix_stable_ofs = [] {
std::array<int, size_t(N) * N + 1> r{};
for (int i = 0; i <= N * N; i++) r[size_t(i)] = i * N;
return r;
}();
template <engine E, int N>
struct matrix_stable
: componentwise<E, mat<typename E::value_type, N>, N * N, matrix_stable_ofs<N>> {
using base = componentwise<E, mat<typename E::value_type, N>, N * N, matrix_stable_ofs<N>>;
static constexpr bool commutative = false;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
// entry (r, c)'s k-th addend a(r,k)*b(k,c), grouped per the offsets
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++) for (int k = 0; k < N; k++)
p.t[(size_t(r) * N + c) * N + k] = E::mul(a.t[size_t(r) * N + k], b.t[size_t(k) * N + c], n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++) for (int k = 0; k < N; k++)
p.t[(size_t(r) * N + c) * N + k] = E::mul2(
a1.t[size_t(r) * N + k], b1.t[size_t(k) * N + c],
a2.t[size_t(r) * N + k], b2.t[size_t(k) * N + c],
n
);
return p;
}
};
template <int N> constexpr std::array<int, size_t(N) + 1> trunc_series_stable_ofs = [] {
std::array<int, size_t(N) + 1> r{};
for (int i = 0; i <= N; i++) r[size_t(i)] = i * (i + 1) / 2;
return r;
}();
template <engine E, int N>
struct trunc_stable
: componentwise<E, trunc_series<typename E::value_type, N>, N, trunc_series_stable_ofs<N>> {
using base = componentwise<E, trunc_series<typename E::value_type, N>, N, trunc_series_stable_ofs<N>>;
static constexpr bool commutative = E::commutative;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
for (int s = 0; s < N; s++) for (int i = 0; i <= s; i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)] + i)] = E::mul(a.t[size_t(i)], b.t[size_t(s - i)], n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int s = 0; s < N; s++) for (int i = 0; i <= s; i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)] + i)] = E::mul2(
a1.t[size_t(i)], b1.t[size_t(s - i)],
a2.t[size_t(i)], b2.t[size_t(s - i)],
n
);
return p;
}
};
/* namespace ecnerwala::fft::engines */ }
// clang-format off
// @formatter:off
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wpragmas"
#pragma GCC diagnostic ignored "-Wunknown-warning-option"
#pragma GCC diagnostic ignored "-Wmisleading-indentation"
#pragma GCC diagnostic ignored "-Wmultistatement-macros"
#include <bits/stdc++.h>
// src/fft/common.hpp
namespace ecnerwala{
template<class T>int sz(T&&arg){using std::size;return int(size(std::forward<T>(arg)));}
inline int nextPow2(int s){return 1<<(s>1?32-__builtin_clz(s-1):0);}
namespace fft{
using std::swap;
using std::vector;
using std::min;
using std::max;
#ifndef ECNERWALA_FFT_POOL_STORAGE
#define ECNERWALA_FFT_POOL_STORAGE
#endif
template<typename T>struct buffer_pool{
static inline ECNERWALA_FFT_POOL_STORAGE std::vector<std::vector<T>>free_list;
struct handle{
std::vector<T>v;
explicit handle(int n){
if(!free_list.empty()){
v=std::move(free_list.back());
free_list.pop_back();
}
v.assign(n,T());
}
handle(const handle&)=delete;
handle&operator=(const handle&)=delete;
handle(handle&&o)noexcept:v(std::move(o.v)){}
~handle(){if(v.capacity())free_list.push_back(std::move(v));}
T&operator[](int i){return v[i];}
operator std::span<T>(){return std::span<T>(v);}
std::span<T>span(){return std::span<T>(v);}
};
static handle get(int n){return handle(n);}
};
}
}
// src/fft/engine.hpp
namespace ecnerwala::fft{
struct assign_op{template<typename T>void operator()(T&d,T v)const{d=v;}};
struct add_op{template<typename T>void operator()(T&d,T v)const{d+=v;}};
struct sub_op{template<typename T>void operator()(T&d,T v)const{d-=v;}};
struct add_twice_op{template<typename T>void operator()(T&d,T v)const{d+=v+v;}};
template<typename E>
concept engine=requires(
std::span<const typename E::value_type>in,
std::span<typename E::value_type>out,
typename E::transformed&t,
const typename E::transformed&ct,
typename E::product&p,
const typename E::product&cp,
int n
){
typename E::value_type;
{E::transform(in,n)}->std::same_as<typename E::transformed>;
{ct.size()}->std::same_as<int>;
E::extend_to(t,n,in);
{E::downsample(ct,n,false)}->std::same_as<typename E::transformed>;
{E::downsample(cp,n,false)}->std::same_as<typename E::product>;
{E::negate_arg(ct,n)}->std::same_as<typename E::transformed>;
{E::mul(ct,ct,n)}->std::same_as<typename E::product>;
{E::sq(ct,n)}->std::same_as<typename E::product>;
{E::mul2(ct,ct,ct,ct,n)}->std::same_as<typename E::template product_t<2*E::unit_scale>>;
E::finish(std::move(p),out);
E::finish(std::move(p),out,add_op{});
E::finish(E::add(std::move(p),std::move(p)),out);
{E::add(E::transform(in,n),ct)}->std::same_as<typename E::template transformed_t<2*E::unit_scale>>;
{E::add(std::move(p),std::move(p))}->std::same_as<typename E::template product_t<2*E::unit_scale>>;
requires std::same_as<std::remove_cvref_t<decltype(E::commutative)>,bool>;
requires std::same_as<std::remove_cvref_t<decltype(E::unit_scale)>,int>;
};
template<typename A,typename B>
concept same_engine=std::same_as<typename A::engine_t,typename B::engine_t>;
template<engine E>using transformed=typename E::transformed;
}
// src/fft/engines/algebras.hpp
namespace ecnerwala::fft::engines{
template<typename num,int N>struct mat{
std::array<num,size_t(N)*N>a{};
num&operator[](std::array<int,2>rc){return a[size_t(rc[0])*N+rc[1]];}
const num&operator[](std::array<int,2>rc)const{return a[size_t(rc[0])*N+rc[1]];}
num*data(){return a.data();}
const num*data()const{return a.data();}
mat&operator+=(const mat&o){for(int i=0;i<N*N;i++)a[i]+=o.a[i];return*this;}
friend mat operator+(mat x,const mat&y){x+=y;return x;}
mat&operator-=(const mat&o){for(int i=0;i<N*N;i++)a[i]-=o.a[i];return*this;}
friend mat operator-(mat x,const mat&y){x-=y;return x;}
friend mat operator*(const mat&x,const mat&y){
mat r;
for(int i=0;i<N;i++)for(int k=0;k<N;k++)for(int j=0;j<N;j++)
r[{i,j}]+=x[{i,k}]*y[{k,j}];
return r;
}
mat&operator*=(const mat&o){return*this=*this*o;}
friend bool operator==(const mat&,const mat&)=default;
};
template<typename num,int N>struct trunc_series{
std::array<num,size_t(N)>a{};
num&operator[](int i){return a[size_t(i)];}
const num&operator[](int i)const{return a[size_t(i)];}
num*data(){return a.data();}
const num*data()const{return a.data();}
trunc_series&operator+=(const trunc_series&o){for(int i=0;i<N;i++)a[i]+=o.a[i];return*this;}
friend trunc_series operator+(trunc_series x,const trunc_series&y){x+=y;return x;}
trunc_series&operator-=(const trunc_series&o){for(int i=0;i<N;i++)a[i]-=o.a[i];return*this;}
friend trunc_series operator-(trunc_series x,const trunc_series&y){x-=y;return x;}
friend trunc_series operator*(const trunc_series&x,const trunc_series&y){
trunc_series r;
for(int i=0;i<N;i++)for(int j=0;j<N-i;j++)r[i+j]+=x[i]*y[j];
return r;
}
trunc_series&operator*=(const trunc_series&o){return*this=*this*o;}
friend bool operator==(const trunc_series&,const trunc_series&)=default;
};
template<int L>constexpr std::array<int,size_t(L)+1>componentwise_iota=[]{
std::array<int,size_t(L)+1>r{};
for(int i=0;i<=L;i++)r[size_t(i)]=i;
return r;
}();
template<engine E,typename V,int L,std::array<int,size_t(L)+1>Ofs=componentwise_iota<L>>
struct componentwise{
using S=typename E::value_type;
using value_type=V;
static constexpr int P=Ofs[size_t(L)];
static constexpr int unit_scale=E::unit_scale;
template<int A=unit_scale>struct transformed_t{
std::array<typename E::template transformed_t<A>,size_t(L)>t;
int size()const{return t[0].size();}
transformed_t()=default;
template<int A2>requires(A2!=A)explicit(A2>A)transformed_t(transformed_t<A2>&&o){
for(int c=0;c<L;c++)
t[c]=typename E::template transformed_t<A>(std::move(o.t[c]));
}
};
using transformed=transformed_t<>;
template<int K>struct product_t{
std::array<typename E::template product_t<K>,size_t(P)>t;
int size()const{return t[0].size();}
product_t()=default;
template<int K2>requires(K2!=K)explicit(K2>K)product_t(product_t<K2>&&o){
for(int c=0;c<P;c++)
t[c]=typename E::template product_t<K>(std::move(o.t[c]));
}
};
static transformed transform(std::span<const V>a,int n){
transformed r;
auto buf=buffer_pool<S>::get(sz(a));
for(int c=0;c<L;c++){
for(int i=0;i<sz(a);i++)buf[i]=a[i].data()[c];
r.t[c]=E::transform(std::span<const S>(buf.span()),n);
}
return r;
}
static void extend_to(transformed&t,int m,std::span<const V>coeffs){
if(t.size()>=m)return;
auto buf=buffer_pool<S>::get(sz(coeffs));
for(int c=0;c<L;c++){
for(int i=0;i<sz(coeffs);i++)buf[i]=coeffs[i].data()[c];
E::extend_to(t.t[c],m,std::span<const S>(buf.span()));
}
}
template<int A>static transformed_t<A>downsample(const transformed_t<A>&t,int n,bool odd){
transformed_t<A>r;
for(int c=0;c<L;c++)r.t[c]=E::downsample(t.t[c],n,odd);
return r;
}
template<int K>static product_t<K>downsample(const product_t<K>&p,int n,bool odd){
product_t<K>r;
for(int c=0;c<P;c++)r.t[c]=E::downsample(p.t[c],n,odd);
return r;
}
template<int A>static transformed_t<A>negate_arg(const transformed_t<A>&t,int n){
transformed_t<A>r;
for(int c=0;c<L;c++)r.t[c]=E::negate_arg(t.t[c],n);
return r;
}
template<int A,int B>static transformed_t<A+B>add(transformed_t<A>&&a,const transformed_t<B>&b){
transformed_t<A+B>r;
for(int c=0;c<L;c++)r.t[c]=E::add(std::move(a.t[c]),b.t[c]);
return r;
}
template<int K1,int K2>static product_t<K1+K2>add(product_t<K1>&&a,product_t<K2>&&b){
product_t<K1+K2>r;
for(int c=0;c<P;c++)r.t[c]=E::add(std::move(a.t[c]),std::move(b.t[c]));
return r;
}
template<int K,typename Op=assign_op>static void finish(product_t<K>&&p,std::span<V>out,Op op={}){
auto buf=buffer_pool<S>::get(sz(out));
auto emit=[&](std::span<V>dst){
for(int c=0;c<L;c++){
E::finish(std::move(p.t[Ofs[size_t(c)]]),buf.span());
for(int j=Ofs[size_t(c)]+1;j<Ofs[size_t(c)+1];j++)
E::finish(std::move(p.t[j]),buf.span(),add_op{});
for(int i=0;i<sz(dst);i++)dst[i].data()[c]=buf[i];
}
};
if constexpr(std::same_as<Op,assign_op>){
emit(out);
}else{
auto vbuf=buffer_pool<V>::get(sz(out));
emit(vbuf.span());
for(int i=0;i<sz(out);i++)op(out[i],vbuf.span()[i]);
}
}
};
template<engine E,int N>
struct matrix:componentwise<E,mat<typename E::value_type,N>,N*N>{
using base=componentwise<E,mat<typename E::value_type,N>,N*N>;
static constexpr bool commutative=false;
static constexpr int unit_scale=base::unit_scale;
template<int A=unit_scale>using transformed_t=typename base::template transformed_t<A>;
template<int K>using product_t=typename base::template product_t<K*N>;
using transformed=typename base::transformed;
using product=product_t<unit_scale*unit_scale>;
template<int A,int B,int k=0>
static auto entry(const transformed_t<A>&a,const transformed_t<B>&b,int r,int c,int n){
auto e=E::mul(a.t[size_t(r)*N+k],b.t[size_t(k)*N+c],n);
if constexpr(k+1==N)return e;
else return E::add(std::move(e),entry<A,B,k+1>(a,b,r,c,n));
}
template<int A,int B>
static product_t<A*B>mul(const transformed_t<A>&a,const transformed_t<B>&b,int n){
product_t<A*B>p;
for(int r=0;r<N;r++)for(int c=0;c<N;c++)
p.t[size_t(r)*N+c]=entry<A,B>(a,b,r,c,n);
return p;
}
template<int A>static auto sq(const transformed_t<A>&a,int n){return mul(a,a,n);}
template<int A1,int B1,int A2,int B2,int k=0>
static auto entry2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int r,int c,int n
){
auto e=E::mul2(
a1.t[size_t(r)*N+k],b1.t[size_t(k)*N+c],
a2.t[size_t(r)*N+k],b2.t[size_t(k)*N+c],
n
);
if constexpr(k+1==N)return e;
else return E::add(std::move(e),entry2<A1,B1,A2,B2,k+1>(a1,b1,a2,b2,r,c,n));
}
template<int A1,int B1,int A2,int B2>
static product_t<A1*B1+A2*B2>mul2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
product_t<A1*B1+A2*B2>p;
for(int r=0;r<N;r++)for(int c=0;c<N;c++)
p.t[size_t(r)*N+c]=entry2<A1,B1,A2,B2>(a1,b1,a2,b2,r,c,n);
return p;
}
};
template<engine E,int N>
struct trunc:componentwise<E,trunc_series<typename E::value_type,N>,N>{
using base=componentwise<E,trunc_series<typename E::value_type,N>,N>;
static constexpr bool commutative=E::commutative;
static constexpr int unit_scale=base::unit_scale;
template<int A=unit_scale>using transformed_t=typename base::template transformed_t<A>;
template<int K>using product_t=typename base::template product_t<K*N>;
using transformed=typename base::transformed;
using product=product_t<unit_scale*unit_scale>;
template<int A,int B,int s,int i=0>
static auto entry(const transformed_t<A>&a,const transformed_t<B>&b,int n){
auto e=E::mul(a.t[size_t(i)],b.t[size_t(s-i)],n);
if constexpr(i==s)return e;
else return E::add(std::move(e),entry<A,B,s,i+1>(a,b,n));
}
template<int A,int B>
static product_t<A*B>mul(const transformed_t<A>&a,const transformed_t<B>&b,int n){
product_t<A*B>p;
[&]<size_t...s_>(std::index_sequence<s_...>){
((p.t[s_]=entry<A,B,int(s_)>(a,b,n)),...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
template<int A>static auto sq(const transformed_t<A>&a,int n){return mul(a,a,n);}
template<int A1,int B1,int A2,int B2,int s,int i=0>
static auto entry2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
auto e=E::mul2(a1.t[size_t(i)],b1.t[size_t(s-i)],a2.t[size_t(i)],b2.t[size_t(s-i)],n);
if constexpr(i==s)return e;
else return E::add(std::move(e),entry2<A1,B1,A2,B2,s,i+1>(a1,b1,a2,b2,n));
}
template<int A1,int B1,int A2,int B2>
static product_t<A1*B1+A2*B2>mul2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
product_t<A1*B1+A2*B2>p;
[&]<size_t...s_>(std::index_sequence<s_...>){
((p.t[s_]=entry2<A1,B1,A2,B2,int(s_)>(a1,b1,a2,b2,n)),...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
};
template<int N>constexpr std::array<int,size_t(N)*N+1>matrix_stable_ofs=[]{
std::array<int,size_t(N)*N+1>r{};
for(int i=0;i<=N*N;i++)r[size_t(i)]=i*N;
return r;
}();
template<engine E,int N>
struct matrix_stable
:componentwise<E,mat<typename E::value_type,N>,N*N,matrix_stable_ofs<N>>{
using base=componentwise<E,mat<typename E::value_type,N>,N*N,matrix_stable_ofs<N>>;
static constexpr bool commutative=false;
static constexpr int unit_scale=base::unit_scale;
template<int A=unit_scale>using transformed_t=typename base::template transformed_t<A>;
template<int K>using product_t=typename base::template product_t<K>;
using transformed=typename base::transformed;
using product=product_t<unit_scale*unit_scale>;
template<int A,int B>
static product_t<A*B>mul(const transformed_t<A>&a,const transformed_t<B>&b,int n){
product_t<A*B>p;
for(int r=0;r<N;r++)for(int c=0;c<N;c++)for(int k=0;k<N;k++)
p.t[(size_t(r)*N+c)*N+k]=E::mul(a.t[size_t(r)*N+k],b.t[size_t(k)*N+c],n);
return p;
}
template<int A>static auto sq(const transformed_t<A>&a,int n){return mul(a,a,n);}
template<int A1,int B1,int A2,int B2>
static product_t<A1*B1+A2*B2>mul2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
product_t<A1*B1+A2*B2>p;
for(int r=0;r<N;r++)for(int c=0;c<N;c++)for(int k=0;k<N;k++)
p.t[(size_t(r)*N+c)*N+k]=E::mul2(
a1.t[size_t(r)*N+k],b1.t[size_t(k)*N+c],
a2.t[size_t(r)*N+k],b2.t[size_t(k)*N+c],
n
);
return p;
}
};
template<int N>constexpr std::array<int,size_t(N)+1>trunc_series_stable_ofs=[]{
std::array<int,size_t(N)+1>r{};
for(int i=0;i<=N;i++)r[size_t(i)]=i*(i+1)/2;
return r;
}();
template<engine E,int N>
struct trunc_stable
:componentwise<E,trunc_series<typename E::value_type,N>,N,trunc_series_stable_ofs<N>>{
using base=componentwise<E,trunc_series<typename E::value_type,N>,N,trunc_series_stable_ofs<N>>;
static constexpr bool commutative=E::commutative;
static constexpr int unit_scale=base::unit_scale;
template<int A=unit_scale>using transformed_t=typename base::template transformed_t<A>;
template<int K>using product_t=typename base::template product_t<K>;
using transformed=typename base::transformed;
using product=product_t<unit_scale*unit_scale>;
template<int A,int B>
static product_t<A*B>mul(const transformed_t<A>&a,const transformed_t<B>&b,int n){
product_t<A*B>p;
for(int s=0;s<N;s++)for(int i=0;i<=s;i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)]+i)]=E::mul(a.t[size_t(i)],b.t[size_t(s-i)],n);
return p;
}
template<int A>static auto sq(const transformed_t<A>&a,int n){return mul(a,a,n);}
template<int A1,int B1,int A2,int B2>
static product_t<A1*B1+A2*B2>mul2(
const transformed_t<A1>&a1,const transformed_t<B1>&b1,
const transformed_t<A2>&a2,const transformed_t<B2>&b2,
int n
){
product_t<A1*B1+A2*B2>p;
for(int s=0;s<N;s++)for(int i=0;i<=s;i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)]+i)]=E::mul2(
a1.t[size_t(i)],b1.t[size_t(s-i)],
a2.t[size_t(i)],b2.t[size_t(s-i)],
n
);
return p;
}
};
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
#pragma once
#include <array>
#include <concepts>
#include <cstddef>
#include <span>
#include <utility>
#include "fft/engine.hpp"
namespace ecnerwala::fft::engines {
// Small NxN matrix over num, row-major
// ==== wrapper engines ====
template <typename num, int N> struct mat {
std::array<num, size_t(N) * N> a{};
num& operator[](std::array<int, 2> rc) { return a[size_t(rc[0]) * N + rc[1]]; }
const num& operator[](std::array<int, 2> rc) const { return a[size_t(rc[0]) * N + rc[1]]; }
num* data() { return a.data(); }
const num* data() const { return a.data(); }
mat& operator+=(const mat& o) { for (int i = 0; i < N*N; i++) a[i] += o.a[i]; return *this; }
friend mat operator+(mat x, const mat& y) { x += y; return x; }
mat& operator-=(const mat& o) { for (int i = 0; i < N*N; i++) a[i] -= o.a[i]; return *this; }
friend mat operator-(mat x, const mat& y) { x -= y; return x; }
friend mat operator*(const mat& x, const mat& y) {
mat r;
for (int i = 0; i < N; i++) for (int k = 0; k < N; k++) for (int j = 0; j < N; j++)
r[{i, j}] += x[{i, k}] * y[{k, j}];
return r;
}
mat& operator*=(const mat& o) { return *this = *this * o; }
friend bool operator==(const mat&, const mat&) = default;
};
// Truncated polynomial mod x^N over num
template <typename num, int N> struct trunc_series {
std::array<num, size_t(N)> a{};
num& operator[](int i) { return a[size_t(i)]; }
const num& operator[](int i) const { return a[size_t(i)]; }
num* data() { return a.data(); }
const num* data() const { return a.data(); }
trunc_series& operator+=(const trunc_series& o) { for (int i = 0; i < N; i++) a[i] += o.a[i]; return *this; }
friend trunc_series operator+(trunc_series x, const trunc_series& y) { x += y; return x; }
trunc_series& operator-=(const trunc_series& o) { for (int i = 0; i < N; i++) a[i] -= o.a[i]; return *this; }
friend trunc_series operator-(trunc_series x, const trunc_series& y) { x -= y; return x; }
friend trunc_series operator*(const trunc_series& x, const trunc_series& y) {
trunc_series r;
for (int i = 0; i < N; i++) for (int j = 0; j < N - i; j++) r[i + j] += x[i] * y[j];
return r;
}
trunc_series& operator*=(const trunc_series& o) { return *this = *this * o; }
friend bool operator==(const trunc_series&, const trunc_series&) = default;
};
// componentwise
// These are "componentwise engines" which model free modules/algebras over the underlying ring.
// Matrices are the canonical example: we can take entry-wise transforms, then multiply/add in transformed space.
//
// We start with a shared `componentwise` base class which handles all linear ops, i.e. not mul.
//
// If the underlying has unit_scale = 1, we may need to avoid accumulation of sums in transformed space;
// then, the product and the transformed data may have different dimensions.
// We'll represent this by an array Ofs of prefix offsets mapping each input/transform-space dimension to a range of product-space dimensions.
// Specifically, out[c] = sum prod[Ofs[c]:Ofs[c+1]]
template <int L> constexpr std::array<int, size_t(L) + 1> componentwise_iota = [] {
std::array<int, size_t(L) + 1> r{};
for (int i = 0; i <= L; i++) r[size_t(i)] = i;
return r;
}();
template <engine E, typename V, int L, std::array<int, size_t(L) + 1> Ofs = componentwise_iota<L>>
struct componentwise {
using S = typename E::value_type;
using value_type = V;
static constexpr int P = Ofs[size_t(L)]; // total product components
static constexpr int unit_scale = E::unit_scale;
template <int A = unit_scale> struct transformed_t {
std::array<typename E::template transformed_t<A>, size_t(L)> t;
int size() const { return t[0].size(); }
transformed_t() = default;
template <int A2> requires (A2 != A) explicit(A2 > A) transformed_t(transformed_t<A2>&& o) {
for (int c = 0; c < L; c++)
t[c] = typename E::template transformed_t<A>(std::move(o.t[c]));
}
};
using transformed = transformed_t<>;
// TODO: if E::product_t == E::transformed_t, mirror that here
template <int K> struct product_t {
std::array<typename E::template product_t<K>, size_t(P)> t;
int size() const { return t[0].size(); }
product_t() = default;
template <int K2> requires (K2 != K) explicit(K2 > K) product_t(product_t<K2>&& o) {
for (int c = 0; c < P; c++)
t[c] = typename E::template product_t<K>(std::move(o.t[c]));
}
};
static transformed transform(std::span<const V> a, int n) {
transformed r;
auto buf = buffer_pool<S>::get(sz(a));
for (int c = 0; c < L; c++) {
for (int i = 0; i < sz(a); i++) buf[i] = a[i].data()[c];
r.t[c] = E::transform(std::span<const S>(buf.span()), n);
}
return r;
}
static void extend_to(transformed& t, int m, std::span<const V> coeffs) {
if (t.size() >= m) return;
auto buf = buffer_pool<S>::get(sz(coeffs));
for (int c = 0; c < L; c++) {
for (int i = 0; i < sz(coeffs); i++) buf[i] = coeffs[i].data()[c];
E::extend_to(t.t[c], m, std::span<const S>(buf.span()));
}
}
template <int A> static transformed_t<A> downsample(const transformed_t<A>& t, int n, bool odd) {
transformed_t<A> r;
for (int c = 0; c < L; c++) r.t[c] = E::downsample(t.t[c], n, odd);
return r;
}
template <int K> static product_t<K> downsample(const product_t<K>& p, int n, bool odd) {
product_t<K> r;
for (int c = 0; c < P; c++) r.t[c] = E::downsample(p.t[c], n, odd);
return r;
}
template <int A> static transformed_t<A> negate_arg(const transformed_t<A>& t, int n) {
transformed_t<A> r;
for (int c = 0; c < L; c++) r.t[c] = E::negate_arg(t.t[c], n);
return r;
}
template <int A, int B> static transformed_t<A + B> add(transformed_t<A>&& a, const transformed_t<B>& b) {
transformed_t<A + B> r;
for (int c = 0; c < L; c++) r.t[c] = E::add(std::move(a.t[c]), b.t[c]);
return r;
}
template <int K1, int K2> static product_t<K1 + K2> add(product_t<K1>&& a, product_t<K2>&& b) {
product_t<K1 + K2> r;
for (int c = 0; c < P; c++) r.t[c] = E::add(std::move(a.t[c]), std::move(b.t[c]));
return r;
}
template <int K, typename Op = assign_op> static void finish(product_t<K>&& p, std::span<V> out, Op op = {}) {
auto buf = buffer_pool<S>::get(sz(out));
auto emit = [&](std::span<V> dst) {
for (int c = 0; c < L; c++) {
E::finish(std::move(p.t[Ofs[size_t(c)]]), buf.span());
for (int j = Ofs[size_t(c)] + 1; j < Ofs[size_t(c) + 1]; j++)
E::finish(std::move(p.t[j]), buf.span(), add_op{});
for (int i = 0; i < sz(dst); i++) dst[i].data()[c] = buf[i];
}
};
// Op must see each out element whole, exactly once, so compose
// non-assign ops through an element buffer.
if constexpr (std::same_as<Op, assign_op>) {
emit(out);
} else {
auto vbuf = buffer_pool<V>::get(sz(out));
emit(vbuf.span());
for (int i = 0; i < sz(out); i++) op(out[i], vbuf.span()[i]);
}
}
};
// Convolve mat<N> (NxN matrices), with accumulation in product space
template <engine E, int N>
struct matrix : componentwise<E, mat<typename E::value_type, N>, N * N> {
using base = componentwise<E, mat<typename E::value_type, N>, N * N>;
static constexpr bool commutative = false;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K * N>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
// right fold over k so a tracked inner engine's per-addend types line up
template <int A, int B, int k = 0>
static auto entry(const transformed_t<A>& a, const transformed_t<B>& b, int r, int c, int n) {
auto e = E::mul(a.t[size_t(r) * N + k], b.t[size_t(k) * N + c], n);
if constexpr (k + 1 == N) return e;
else return E::add(std::move(e), entry<A, B, k + 1>(a, b, r, c, n));
}
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++)
p.t[size_t(r) * N + c] = entry<A, B>(a, b, r, c, n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2, int k = 0>
static auto entry2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int r, int c, int n
) {
auto e = E::mul2(
a1.t[size_t(r) * N + k], b1.t[size_t(k) * N + c],
a2.t[size_t(r) * N + k], b2.t[size_t(k) * N + c],
n
);
if constexpr (k + 1 == N) return e;
else return E::add(std::move(e), entry2<A1, B1, A2, B2, k + 1>(a1, b1, a2, b2, r, c, n));
}
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++)
p.t[size_t(r) * N + c] = entry2<A1, B1, A2, B2>(a1, b1, a2, b2, r, c, n);
return p;
}
};
// Convolve trunc_series<num, N> (power series truncated at N), with accumulation in product space
template <engine E, int N>
struct trunc : componentwise<E, trunc_series<typename E::value_type, N>, N> {
using base = componentwise<E, trunc_series<typename E::value_type, N>, N>;
static constexpr bool commutative = E::commutative;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K * N>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B, int s, int i = 0>
static auto entry(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
auto e = E::mul(a.t[size_t(i)], b.t[size_t(s - i)], n);
if constexpr (i == s) return e;
else return E::add(std::move(e), entry<A, B, s, i + 1>(a, b, n));
}
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
[&]<size_t... s_>(std::index_sequence<s_...>) {
((p.t[s_] = entry<A, B, int(s_)>(a, b, n)), ...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2, int s, int i = 0>
static auto entry2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
auto e = E::mul2(a1.t[size_t(i)], b1.t[size_t(s - i)], a2.t[size_t(i)], b2.t[size_t(s - i)], n);
if constexpr (i == s) return e;
else return E::add(std::move(e), entry2<A1, B1, A2, B2, s, i + 1>(a1, b1, a2, b2, n));
}
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
[&]<size_t... s_>(std::index_sequence<s_...>) {
((p.t[s_] = entry2<A1, B1, A2, B2, int(s_)>(a1, b1, a2, b2, n)), ...);
}(std::make_index_sequence<size_t(N)>{});
return p;
}
};
// Stable variants of the wrapper engines: do not accumulate in product space.
// This costs an extra log factor.
template <int N> constexpr std::array<int, size_t(N) * N + 1> matrix_stable_ofs = [] {
std::array<int, size_t(N) * N + 1> r{};
for (int i = 0; i <= N * N; i++) r[size_t(i)] = i * N;
return r;
}();
template <engine E, int N>
struct matrix_stable
: componentwise<E, mat<typename E::value_type, N>, N * N, matrix_stable_ofs<N>> {
using base = componentwise<E, mat<typename E::value_type, N>, N * N, matrix_stable_ofs<N>>;
static constexpr bool commutative = false;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
// entry (r, c)'s k-th addend a(r,k)*b(k,c), grouped per the offsets
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++) for (int k = 0; k < N; k++)
p.t[(size_t(r) * N + c) * N + k] = E::mul(a.t[size_t(r) * N + k], b.t[size_t(k) * N + c], n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int r = 0; r < N; r++) for (int c = 0; c < N; c++) for (int k = 0; k < N; k++)
p.t[(size_t(r) * N + c) * N + k] = E::mul2(
a1.t[size_t(r) * N + k], b1.t[size_t(k) * N + c],
a2.t[size_t(r) * N + k], b2.t[size_t(k) * N + c],
n
);
return p;
}
};
template <int N> constexpr std::array<int, size_t(N) + 1> trunc_series_stable_ofs = [] {
std::array<int, size_t(N) + 1> r{};
for (int i = 0; i <= N; i++) r[size_t(i)] = i * (i + 1) / 2;
return r;
}();
template <engine E, int N>
struct trunc_stable
: componentwise<E, trunc_series<typename E::value_type, N>, N, trunc_series_stable_ofs<N>> {
using base = componentwise<E, trunc_series<typename E::value_type, N>, N, trunc_series_stable_ofs<N>>;
static constexpr bool commutative = E::commutative;
static constexpr int unit_scale = base::unit_scale;
template <int A = unit_scale> using transformed_t = typename base::template transformed_t<A>;
template <int K> using product_t = typename base::template product_t<K>;
using transformed = typename base::transformed;
using product = product_t<unit_scale * unit_scale>;
template <int A, int B>
static product_t<A * B> mul(const transformed_t<A>& a, const transformed_t<B>& b, int n) {
product_t<A * B> p;
for (int s = 0; s < N; s++) for (int i = 0; i <= s; i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)] + i)] = E::mul(a.t[size_t(i)], b.t[size_t(s - i)], n);
return p;
}
template <int A> static auto sq(const transformed_t<A>& a, int n) { return mul(a, a, n); }
template <int A1, int B1, int A2, int B2>
static product_t<A1 * B1 + A2 * B2> mul2(
const transformed_t<A1>& a1, const transformed_t<B1>& b1,
const transformed_t<A2>& a2, const transformed_t<B2>& b2,
int n
) {
product_t<A1 * B1 + A2 * B2> p;
for (int s = 0; s < N; s++) for (int i = 0; i <= s; i++)
p.t[size_t(trunc_series_stable_ofs<N>[size_t(s)] + i)] = E::mul2(
a1.t[size_t(i)], b1.t[size_t(s - i)],
a2.t[size_t(i)], b2.t[size_t(s - i)],
n
);
return p;
}
};
/* namespace ecnerwala::fft::engines */ }