ecnerwala's competitive programming library
#include "fft/engine.hpp"
| Coverage | Exec / Excl / Total | |
|---|---|---|
| Lines | 75.0% | 3 / 0 / 4 |
| Functions | 100.0% | 4 / 0 / 4 |
| Full report |
#pragma once
#include <concepts>
#include <span>
#include <type_traits>
#include <utility>
#include "fft/common.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 */ }
#include <concepts>
#include <span>
#include <type_traits>
#include <utility>
#include <algorithm>
#include <iterator>
#include <vector>
#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 */ }
// 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;
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
#pragma once
#include <concepts>
#include <span>
#include <type_traits>
#include <utility>
#include "fft/common.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 */ }