cp-book

ecnerwala's competitive programming library

View the Project on GitHub ecnerwala/cp-book

:heavy_check_mark: #include "fft/multiply.hpp"

View this file on GitHub · Last update: 2026-07-30 21:24:54-07:00

Depends on

Required by

Verified with

Code

Coverage Exec / Excl / Total
Lines 100.0% 190 / 0 / 190
Functions 100.0% 213 / 0 / 213
Branches 71.0% 1351 / 0 / 1903
Full report
#pragma once

#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <span>
#include <utility>
#include <vector>

#include "fft/engine.hpp"

namespace ecnerwala::fft {

// ==== multiply layer ====
// These are free functions to convolve spans.
//
// The interfaces will typically take input spans, an output span, and an Op representing how to fold the result into the output.
// Output spans may alias one of the input spans.
// Output spans may be shorter than expected; the output will just be truncated.
//
// Some functions may also take E::transformed& objects associated with the input
// spans. These will be lazily filled (see E::extend_to) and used if available.

// Circular convolution mod n (power of 2)
template <engine E, typename Op = assign_op>
void multiply_circular(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, int n, Op op = {}) {
	assert(!(n & (n-1)));
	auto ta = E::transform(a, n);
	auto tb = E::transform(b, n);
	E::finish(E::mul(ta, tb, n), out, op);
}

template <engine E, typename Op = assign_op>
void square_circular(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, int n, Op op = {}) {
	assert(!(n & (n-1)));
	auto ta = E::transform(a, n);
	E::finish(E::sq(ta, n), out, op);
}

namespace detail {
// Arrays of length 2^k + 1 are somewhat common, so we will optimize them by
// multiplying mod 2^k, and fixing up the leading coefficient.

// Helpers to detect and perform this optimization.
struct conv_size { int n; bool cut; };
inline conv_size conv_size_for(int s) {
	int n = nextPow2(s);
	bool cut = (n == 2 * (s - 1));
	return {cut ? n / 2 : n, cut};
}

// Call op while lazily applying the correction if necessary
template <typename T, typename Op>
void emit_linear(std::span<T> buf, int n, int s, bool cut, T c0, std::span<T> out, Op op) {
	T cn{};
	if (cut) {
		cn = buf[0] - c0;
		buf[0] = c0;
	}
	int lim = min(sz(out), min(s, n));
	for (int i = 0; i < lim; i++) op(out[i], buf[i]);
	if (cut && sz(out) >= s) op(out[s-1], cn);
}

// Applies op, diverting the wrapped leading coefficient of a cut product:
// out[0] receives c0 and the wraparound term is captured into cn for the
// caller to emit at out[s-1].
template <typename T, typename Op>
struct cut_op {
	Op op;
	T* out0;
	T c0;
	T& cn;
	void operator()(T& x, T v) const {
		if (&x == out0) { cn = v - c0; v = c0; }
		op(x, v);
	}
};

// finish + emit_linear fused: write the finished product directly into out,
// applying the cut correction in place.
template <engine E, typename P, typename Op = assign_op>
void finish_linear(
	P&& p, int n, int s, bool cut,
	typename E::value_type c0, std::span<typename E::value_type> out, Op op = {}
) {
	using T = typename E::value_type;
	if (sz(out) == 0) return;
	int lim = min(sz(out), min(s, n));
	if (!cut) {
		E::finish(std::move(p), out.subspan(0, lim), op);
	} else {
		T cn{};
		E::finish(std::move(p), out.subspan(0, lim), cut_op<T, Op>{op, &out[0], c0, cn});
		if (sz(out) >= s) op(out[s-1], cn);
	}
}

}

template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	int s = sz(a) + sz(b) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * b[0];
	auto buf = buffer_pool<T>::get(n);
	multiply_circular<E>(a, b, buf.span(), n);
	detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	int s = sz(a) + sz(b) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * b[0];
	E::extend_to(ta, n, a);
	E::extend_to(tb, n, b);
	detail::finish_linear<E>(E::mul(ta, tb, n), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void multiply_add2(std::span<const typename E::value_type> a1, transformed<E>& ta1,
		std::span<const typename E::value_type> b1, transformed<E>& tb1,
		std::span<const typename E::value_type> a2, transformed<E>& ta2,
		std::span<const typename E::value_type> b2, transformed<E>& tb2,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
	int s = sz(a1) + sz(b1) - 1;
	assert(sz(a2) + sz(b2) - 1 == s);
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a1[0] * b1[0] + a2[0] * b2[0];
	E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
	E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
	detail::finish_linear<E>(E::mul2(ta1, tb1, ta2, tb2, n), n, s, cut, c0, out, op);
}

// As multiply_add2, but also outputs the summed pointwise product as a reusable
// transform of the (full-length) result, like multiply_cached.
template <engine E>
void multiply_add2_cached(
		std::span<const typename E::value_type> a1, transformed<E>& ta1,
		std::span<const typename E::value_type> b1, transformed<E>& tb1,
		std::span<const typename E::value_type> a2, transformed<E>& ta2,
		std::span<const typename E::value_type> b2, transformed<E>& tb2,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
	int s = sz(a1) + sz(b1) - 1;
	assert(sz(a2) + sz(b2) - 1 == s);
	coeffs.assign(size_t(s), T{});
	t = transformed<E>{};
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a1[0] * b1[0] + a2[0] * b2[0];
		E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
		E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
		auto p = E::mul2(ta1, tb1, ta2, tb2, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		multiply_add2<E>(a1, ta1, b1, tb1, a2, ta2, b2, tb2, std::span<T>(coeffs));
	}
}

// This helper also accepts an output transform which will be populated if it is cheap to do so
template <engine E>
void multiply_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	coeffs.assign(size_t(sz(a) && sz(b) ? sz(a) + sz(b) - 1 : 0), T{});
	t = transformed<E>{};
	if (coeffs.empty()) return;
	int s = sz(coeffs);
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a[0] * b[0];
		E::extend_to(ta, n, a);
		E::extend_to(tb, n, b);
		auto p = E::mul(ta, tb, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		multiply<E>(a, ta, b, tb, std::span<T>(coeffs));
	}
}

template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0) return;
	int s = 2 * sz(a) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * a[0];
	auto buf = buffer_pool<T>::get(n);
	square_circular<E>(a, buf.span(), n);
	detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0) return;
	int s = 2 * sz(a) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * a[0];
	E::extend_to(ta, n, a);
	detail::finish_linear<E>(E::sq(ta, n), n, s, cut, c0, out, op);
}

// As square, but also outputs the pointwise product as a reusable transform of
// the result (empty when the engine's product isn't a transform).
template <engine E>
void square_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	coeffs.assign(size_t(sz(a) ? 2 * sz(a) - 1 : 0), T{});
	t = transformed<E>{};
	if (coeffs.empty()) return;
	int s = sz(coeffs);
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a[0] * a[0];
		E::extend_to(ta, n, a);
		auto p = E::sq(ta, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		square<E>(a, ta, std::span<T>(coeffs));
	}
}

template <engine E> vector<typename E::value_type> multiply(
		const vector<typename E::value_type>& a, const vector<typename E::value_type>& b) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	vector<T> r(sz(a) + sz(b) - 1);
	multiply<E>(std::span<const T>(a), std::span<const T>(b), std::span<T>(r));
	return r;
}

template <engine E> vector<typename E::value_type> square(const vector<typename E::value_type>& a) {
	using T = typename E::value_type;
	if (sz(a) == 0) return {};
	vector<T> r(2 * sz(a) - 1);
	square<E>(std::span<const T>(a), std::span<T>(r));
	return r;
}

namespace detail {
// emit_linear but for middle_product
template <typename T, typename Op>
void emit_middle(std::span<T> buf, bool cut, int la, int lb, T c0, T ctop, std::span<T> out, Op op) {
	int m = la - lb + 1;
	T cn{};
	if (cut) {
		cn = buf[0] - c0; // for lb == 1 these coincide: slot 0 = c_0 + c_n and ctop = c_n
		buf[lb - 1] -= ctop;
	}
	int lim = min(sz(out), cut ? m - 1 : m);
	for (int t = 0; t < lim; t++) op(out[t], buf[lb - 1 + t]);
	if (cut && sz(out) >= m) op(out[m-1], cn);
}
}

// Middle product (the transposed multiplication): takes only coefficients of a * b which include terms from all of b.
// Must have len(a) >= len(b)
template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	assert(sz(a) >= sz(b));
	if (sz(a) == sz(b)) {
		T r{};
		for (int i = 0; i < sz(a); i++) {
			r += a[i] * b[sz(b) - 1 - i];
		}
		if (sz(out) > 0) op(out[0], r);
		return;
	}
	auto [n, cut] = detail::conv_size_for(sz(a));
	auto buf = buffer_pool<T>::get(n);
	multiply_circular<E>(a, b, buf.span(), n);
	detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
			a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}

// TODO: Let's decide whether to keep vector<> returning forms or not; this
// largely depends on whether we think these functions are a public interface or
// merely convenience for value type implementors.
template <engine E> vector<typename E::value_type> middle_product(
		std::span<const typename E::value_type> a, std::span<const typename E::value_type> b) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	assert(sz(a) >= sz(b));
	vector<T> r(size_t(sz(a) - sz(b) + 1));
	middle_product<E>(a, b, std::span<T>(r));
	return r;
}

template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	assert(sz(a) >= sz(b));
	if (sz(a) == sz(b)) {
		T r{};
		for (int i = 0; i < sz(a); i++) {
			r += a[i] * b[sz(b) - 1 - i];
		}
		if (sz(out) > 0) op(out[0], r);
		return;
	}
	auto [n, cut] = detail::conv_size_for(sz(a));
	E::extend_to(ta, n, a);
	E::extend_to(tb, n, b);
	auto buf = buffer_pool<T>::get(n);
	E::finish(E::mul(ta, tb, n), buf.span());
	detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
			a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}

template <engine E>
vector<typename E::value_type> middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	assert(sz(a) >= sz(b));
	vector<T> r(size_t(sz(a) - sz(b) + 1));
	middle_product<E>(a, ta, b, tb, std::span<T>(r));
	return r;
}

/* namespace ecnerwala::fft */ }
#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <span>
#include <utility>
#include <vector>
#include <type_traits>
#include <iterator>
#line 2 "src/fft/multiply.hpp"

#line 10 "src/fft/multiply.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 12 "src/fft/multiply.hpp"

namespace ecnerwala::fft {

// ==== multiply layer ====
// These are free functions to convolve spans.
//
// The interfaces will typically take input spans, an output span, and an Op representing how to fold the result into the output.
// Output spans may alias one of the input spans.
// Output spans may be shorter than expected; the output will just be truncated.
//
// Some functions may also take E::transformed& objects associated with the input
// spans. These will be lazily filled (see E::extend_to) and used if available.

// Circular convolution mod n (power of 2)
template <engine E, typename Op = assign_op>
void multiply_circular(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, int n, Op op = {}) {
	assert(!(n & (n-1)));
	auto ta = E::transform(a, n);
	auto tb = E::transform(b, n);
	E::finish(E::mul(ta, tb, n), out, op);
}

template <engine E, typename Op = assign_op>
void square_circular(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, int n, Op op = {}) {
	assert(!(n & (n-1)));
	auto ta = E::transform(a, n);
	E::finish(E::sq(ta, n), out, op);
}

namespace detail {
// Arrays of length 2^k + 1 are somewhat common, so we will optimize them by
// multiplying mod 2^k, and fixing up the leading coefficient.

// Helpers to detect and perform this optimization.
struct conv_size { int n; bool cut; };
inline conv_size conv_size_for(int s) {
	int n = nextPow2(s);
	bool cut = (n == 2 * (s - 1));
	return {cut ? n / 2 : n, cut};
}

// Call op while lazily applying the correction if necessary
template <typename T, typename Op>
void emit_linear(std::span<T> buf, int n, int s, bool cut, T c0, std::span<T> out, Op op) {
	T cn{};
	if (cut) {
		cn = buf[0] - c0;
		buf[0] = c0;
	}
	int lim = min(sz(out), min(s, n));
	for (int i = 0; i < lim; i++) op(out[i], buf[i]);
	if (cut && sz(out) >= s) op(out[s-1], cn);
}

// Applies op, diverting the wrapped leading coefficient of a cut product:
// out[0] receives c0 and the wraparound term is captured into cn for the
// caller to emit at out[s-1].
template <typename T, typename Op>
struct cut_op {
	Op op;
	T* out0;
	T c0;
	T& cn;
	void operator()(T& x, T v) const {
		if (&x == out0) { cn = v - c0; v = c0; }
		op(x, v);
	}
};

// finish + emit_linear fused: write the finished product directly into out,
// applying the cut correction in place.
template <engine E, typename P, typename Op = assign_op>
void finish_linear(
	P&& p, int n, int s, bool cut,
	typename E::value_type c0, std::span<typename E::value_type> out, Op op = {}
) {
	using T = typename E::value_type;
	if (sz(out) == 0) return;
	int lim = min(sz(out), min(s, n));
	if (!cut) {
		E::finish(std::move(p), out.subspan(0, lim), op);
	} else {
		T cn{};
		E::finish(std::move(p), out.subspan(0, lim), cut_op<T, Op>{op, &out[0], c0, cn});
		if (sz(out) >= s) op(out[s-1], cn);
	}
}

}

template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	int s = sz(a) + sz(b) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * b[0];
	auto buf = buffer_pool<T>::get(n);
	multiply_circular<E>(a, b, buf.span(), n);
	detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	int s = sz(a) + sz(b) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * b[0];
	E::extend_to(ta, n, a);
	E::extend_to(tb, n, b);
	detail::finish_linear<E>(E::mul(ta, tb, n), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void multiply_add2(std::span<const typename E::value_type> a1, transformed<E>& ta1,
		std::span<const typename E::value_type> b1, transformed<E>& tb1,
		std::span<const typename E::value_type> a2, transformed<E>& ta2,
		std::span<const typename E::value_type> b2, transformed<E>& tb2,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
	int s = sz(a1) + sz(b1) - 1;
	assert(sz(a2) + sz(b2) - 1 == s);
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a1[0] * b1[0] + a2[0] * b2[0];
	E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
	E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
	detail::finish_linear<E>(E::mul2(ta1, tb1, ta2, tb2, n), n, s, cut, c0, out, op);
}

// As multiply_add2, but also outputs the summed pointwise product as a reusable
// transform of the (full-length) result, like multiply_cached.
template <engine E>
void multiply_add2_cached(
		std::span<const typename E::value_type> a1, transformed<E>& ta1,
		std::span<const typename E::value_type> b1, transformed<E>& tb1,
		std::span<const typename E::value_type> a2, transformed<E>& ta2,
		std::span<const typename E::value_type> b2, transformed<E>& tb2,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
	int s = sz(a1) + sz(b1) - 1;
	assert(sz(a2) + sz(b2) - 1 == s);
	coeffs.assign(size_t(s), T{});
	t = transformed<E>{};
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a1[0] * b1[0] + a2[0] * b2[0];
		E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
		E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
		auto p = E::mul2(ta1, tb1, ta2, tb2, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		multiply_add2<E>(a1, ta1, b1, tb1, a2, ta2, b2, tb2, std::span<T>(coeffs));
	}
}

// This helper also accepts an output transform which will be populated if it is cheap to do so
template <engine E>
void multiply_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	coeffs.assign(size_t(sz(a) && sz(b) ? sz(a) + sz(b) - 1 : 0), T{});
	t = transformed<E>{};
	if (coeffs.empty()) return;
	int s = sz(coeffs);
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a[0] * b[0];
		E::extend_to(ta, n, a);
		E::extend_to(tb, n, b);
		auto p = E::mul(ta, tb, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		multiply<E>(a, ta, b, tb, std::span<T>(coeffs));
	}
}

template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0) return;
	int s = 2 * sz(a) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * a[0];
	auto buf = buffer_pool<T>::get(n);
	square_circular<E>(a, buf.span(), n);
	detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0) return;
	int s = 2 * sz(a) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * a[0];
	E::extend_to(ta, n, a);
	detail::finish_linear<E>(E::sq(ta, n), n, s, cut, c0, out, op);
}

// As square, but also outputs the pointwise product as a reusable transform of
// the result (empty when the engine's product isn't a transform).
template <engine E>
void square_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	coeffs.assign(size_t(sz(a) ? 2 * sz(a) - 1 : 0), T{});
	t = transformed<E>{};
	if (coeffs.empty()) return;
	int s = sz(coeffs);
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a[0] * a[0];
		E::extend_to(ta, n, a);
		auto p = E::sq(ta, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		square<E>(a, ta, std::span<T>(coeffs));
	}
}

template <engine E> vector<typename E::value_type> multiply(
		const vector<typename E::value_type>& a, const vector<typename E::value_type>& b) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	vector<T> r(sz(a) + sz(b) - 1);
	multiply<E>(std::span<const T>(a), std::span<const T>(b), std::span<T>(r));
	return r;
}

template <engine E> vector<typename E::value_type> square(const vector<typename E::value_type>& a) {
	using T = typename E::value_type;
	if (sz(a) == 0) return {};
	vector<T> r(2 * sz(a) - 1);
	square<E>(std::span<const T>(a), std::span<T>(r));
	return r;
}

namespace detail {
// emit_linear but for middle_product
template <typename T, typename Op>
void emit_middle(std::span<T> buf, bool cut, int la, int lb, T c0, T ctop, std::span<T> out, Op op) {
	int m = la - lb + 1;
	T cn{};
	if (cut) {
		cn = buf[0] - c0; // for lb == 1 these coincide: slot 0 = c_0 + c_n and ctop = c_n
		buf[lb - 1] -= ctop;
	}
	int lim = min(sz(out), cut ? m - 1 : m);
	for (int t = 0; t < lim; t++) op(out[t], buf[lb - 1 + t]);
	if (cut && sz(out) >= m) op(out[m-1], cn);
}
}

// Middle product (the transposed multiplication): takes only coefficients of a * b which include terms from all of b.
// Must have len(a) >= len(b)
template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	assert(sz(a) >= sz(b));
	if (sz(a) == sz(b)) {
		T r{};
		for (int i = 0; i < sz(a); i++) {
			r += a[i] * b[sz(b) - 1 - i];
		}
		if (sz(out) > 0) op(out[0], r);
		return;
	}
	auto [n, cut] = detail::conv_size_for(sz(a));
	auto buf = buffer_pool<T>::get(n);
	multiply_circular<E>(a, b, buf.span(), n);
	detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
			a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}

// TODO: Let's decide whether to keep vector<> returning forms or not; this
// largely depends on whether we think these functions are a public interface or
// merely convenience for value type implementors.
template <engine E> vector<typename E::value_type> middle_product(
		std::span<const typename E::value_type> a, std::span<const typename E::value_type> b) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	assert(sz(a) >= sz(b));
	vector<T> r(size_t(sz(a) - sz(b) + 1));
	middle_product<E>(a, b, std::span<T>(r));
	return r;
}

template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	assert(sz(a) >= sz(b));
	if (sz(a) == sz(b)) {
		T r{};
		for (int i = 0; i < sz(a); i++) {
			r += a[i] * b[sz(b) - 1 - i];
		}
		if (sz(out) > 0) op(out[0], r);
		return;
	}
	auto [n, cut] = detail::conv_size_for(sz(a));
	E::extend_to(ta, n, a);
	E::extend_to(tb, n, b);
	auto buf = buffer_pool<T>::get(n);
	E::finish(E::mul(ta, tb, n), buf.span());
	detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
			a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}

template <engine E>
vector<typename E::value_type> middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	assert(sz(a) >= sz(b));
	vector<T> r(size_t(sz(a) - sz(b) + 1));
	middle_product<E>(a, ta, b, tb, std::span<T>(r));
	return r;
}

/* 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>
#include <cassert>
// 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/multiply.hpp
namespace ecnerwala::fft{
template<engine E,typename Op=assign_op>
void multiply_circular(std::span<const typename E::value_type>a,std::span<const typename E::value_type>b,
std::span<typename E::value_type>out,int n,Op op={}){
assert(!(n&(n-1)));
auto ta=E::transform(a,n);
auto tb=E::transform(b,n);
E::finish(E::mul(ta,tb,n),out,op);
}
template<engine E,typename Op=assign_op>
void square_circular(std::span<const typename E::value_type>a,std::span<typename E::value_type>out,int n,Op op={}){
assert(!(n&(n-1)));
auto ta=E::transform(a,n);
E::finish(E::sq(ta,n),out,op);
}
namespace detail{
struct conv_size{int n;bool cut;};
inline conv_size conv_size_for(int s){
int n=nextPow2(s);
bool cut=(n==2*(s-1));
return{cut?n/2:n,cut};
}
template<typename T,typename Op>
void emit_linear(std::span<T>buf,int n,int s,bool cut,T c0,std::span<T>out,Op op){
T cn{};
if(cut){
cn=buf[0]-c0;
buf[0]=c0;
}
int lim=min(sz(out),min(s,n));
for(int i=0;i<lim;i++)op(out[i],buf[i]);
if(cut&&sz(out)>=s)op(out[s-1],cn);
}
template<typename T,typename Op>
struct cut_op{
Op op;
T*out0;
T c0;
T&cn;
void operator()(T&x,T v)const{
if(&x==out0){cn=v-c0;v=c0;}
op(x,v);
}
};
template<engine E,typename P,typename Op=assign_op>
void finish_linear(
P&&p,int n,int s,bool cut,
typename E::value_type c0,std::span<typename E::value_type>out,Op op={}
){
using T=typename E::value_type;
if(sz(out)==0)return;
int lim=min(sz(out),min(s,n));
if(!cut){
E::finish(std::move(p),out.subspan(0,lim),op);
}else{
T cn{};
E::finish(std::move(p),out.subspan(0,lim),cut_op<T,Op>{op,&out[0],c0,cn});
if(sz(out)>=s)op(out[s-1],cn);
}
}
}
template<engine E,typename Op=assign_op>
void multiply(std::span<const typename E::value_type>a,std::span<const typename E::value_type>b,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
int s=sz(a)+sz(b)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*b[0];
auto buf=buffer_pool<T>::get(n);
multiply_circular<E>(a,b,buf.span(),n);
detail::emit_linear<T>(buf.span(),n,s,cut,c0,out,op);
}
template<engine E,typename Op=assign_op>
void multiply(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
int s=sz(a)+sz(b)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*b[0];
E::extend_to(ta,n,a);
E::extend_to(tb,n,b);
detail::finish_linear<E>(E::mul(ta,tb,n),n,s,cut,c0,out,op);
}
template<engine E,typename Op=assign_op>
void multiply_add2(std::span<const typename E::value_type>a1,transformed<E>&ta1,
std::span<const typename E::value_type>b1,transformed<E>&tb1,
std::span<const typename E::value_type>a2,transformed<E>&ta2,
std::span<const typename E::value_type>b2,transformed<E>&tb2,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
assert(sz(a1)>0&&sz(b1)>0&&sz(a2)>0&&sz(b2)>0);
int s=sz(a1)+sz(b1)-1;
assert(sz(a2)+sz(b2)-1==s);
auto[n,cut]=detail::conv_size_for(s);
T c0=a1[0]*b1[0]+a2[0]*b2[0];
E::extend_to(ta1,n,a1);E::extend_to(tb1,n,b1);
E::extend_to(ta2,n,a2);E::extend_to(tb2,n,b2);
detail::finish_linear<E>(E::mul2(ta1,tb1,ta2,tb2,n),n,s,cut,c0,out,op);
}
template<engine E>
void multiply_add2_cached(
std::span<const typename E::value_type>a1,transformed<E>&ta1,
std::span<const typename E::value_type>b1,transformed<E>&tb1,
std::span<const typename E::value_type>a2,transformed<E>&ta2,
std::span<const typename E::value_type>b2,transformed<E>&tb2,
std::vector<typename E::value_type>&coeffs,transformed<E>&t){
using T=typename E::value_type;
assert(sz(a1)>0&&sz(b1)>0&&sz(a2)>0&&sz(b2)>0);
int s=sz(a1)+sz(b1)-1;
assert(sz(a2)+sz(b2)-1==s);
coeffs.assign(size_t(s),T{});
t=transformed<E>{};
if constexpr(std::same_as<typename E::product,transformed<E>>){
auto[n,cut]=detail::conv_size_for(s);
T c0=a1[0]*b1[0]+a2[0]*b2[0];
E::extend_to(ta1,n,a1);E::extend_to(tb1,n,b1);
E::extend_to(ta2,n,a2);E::extend_to(tb2,n,b2);
auto p=E::mul2(ta1,tb1,ta2,tb2,n);
auto tp=p;
detail::finish_linear<E>(std::move(p),n,s,cut,c0,std::span<T>(coeffs));
t=std::move(tp);
}else{
multiply_add2<E>(a1,ta1,b1,tb1,a2,ta2,b2,tb2,std::span<T>(coeffs));
}
}
template<engine E>
void multiply_cached(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb,
std::vector<typename E::value_type>&coeffs,transformed<E>&t){
using T=typename E::value_type;
coeffs.assign(size_t(sz(a)&&sz(b)?sz(a)+sz(b)-1:0),T{});
t=transformed<E>{};
if(coeffs.empty())return;
int s=sz(coeffs);
if constexpr(std::same_as<typename E::product,transformed<E>>){
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*b[0];
E::extend_to(ta,n,a);
E::extend_to(tb,n,b);
auto p=E::mul(ta,tb,n);
auto tp=p;
detail::finish_linear<E>(std::move(p),n,s,cut,c0,std::span<T>(coeffs));
t=std::move(tp);
}else{
multiply<E>(a,ta,b,tb,std::span<T>(coeffs));
}
}
template<engine E,typename Op=assign_op>
void square(std::span<const typename E::value_type>a,std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0)return;
int s=2*sz(a)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*a[0];
auto buf=buffer_pool<T>::get(n);
square_circular<E>(a,buf.span(),n);
detail::emit_linear<T>(buf.span(),n,s,cut,c0,out,op);
}
template<engine E,typename Op=assign_op>
void square(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0)return;
int s=2*sz(a)-1;
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*a[0];
E::extend_to(ta,n,a);
detail::finish_linear<E>(E::sq(ta,n),n,s,cut,c0,out,op);
}
template<engine E>
void square_cached(std::span<const typename E::value_type>a,transformed<E>&ta,
std::vector<typename E::value_type>&coeffs,transformed<E>&t){
using T=typename E::value_type;
coeffs.assign(size_t(sz(a)?2*sz(a)-1:0),T{});
t=transformed<E>{};
if(coeffs.empty())return;
int s=sz(coeffs);
if constexpr(std::same_as<typename E::product,transformed<E>>){
auto[n,cut]=detail::conv_size_for(s);
T c0=a[0]*a[0];
E::extend_to(ta,n,a);
auto p=E::sq(ta,n);
auto tp=p;
detail::finish_linear<E>(std::move(p),n,s,cut,c0,std::span<T>(coeffs));
t=std::move(tp);
}else{
square<E>(a,ta,std::span<T>(coeffs));
}
}
template<engine E>vector<typename E::value_type>multiply(
const vector<typename E::value_type>&a,const vector<typename E::value_type>&b){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return{};
vector<T>r(sz(a)+sz(b)-1);
multiply<E>(std::span<const T>(a),std::span<const T>(b),std::span<T>(r));
return r;
}
template<engine E>vector<typename E::value_type>square(const vector<typename E::value_type>&a){
using T=typename E::value_type;
if(sz(a)==0)return{};
vector<T>r(2*sz(a)-1);
square<E>(std::span<const T>(a),std::span<T>(r));
return r;
}
namespace detail{
template<typename T,typename Op>
void emit_middle(std::span<T>buf,bool cut,int la,int lb,T c0,T ctop,std::span<T>out,Op op){
int m=la-lb+1;
T cn{};
if(cut){
cn=buf[0]-c0;
buf[lb-1]-=ctop;
}
int lim=min(sz(out),cut?m-1:m);
for(int t=0;t<lim;t++)op(out[t],buf[lb-1+t]);
if(cut&&sz(out)>=m)op(out[m-1],cn);
}
}
template<engine E,typename Op=assign_op>
void middle_product(std::span<const typename E::value_type>a,std::span<const typename E::value_type>b,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
assert(sz(a)>=sz(b));
if(sz(a)==sz(b)){
T r{};
for(int i=0;i<sz(a);i++){
r+=a[i]*b[sz(b)-1-i];
}
if(sz(out)>0)op(out[0],r);
return;
}
auto[n,cut]=detail::conv_size_for(sz(a));
auto buf=buffer_pool<T>::get(n);
multiply_circular<E>(a,b,buf.span(),n);
detail::emit_middle<T>(buf.span(),cut,sz(a),sz(b),
a[0]*b[0],a[sz(a)-1]*b[sz(b)-1],out,op);
}
template<engine E>vector<typename E::value_type>middle_product(
std::span<const typename E::value_type>a,std::span<const typename E::value_type>b){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return{};
assert(sz(a)>=sz(b));
vector<T>r(size_t(sz(a)-sz(b)+1));
middle_product<E>(a,b,std::span<T>(r));
return r;
}
template<engine E,typename Op=assign_op>
void middle_product(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb,
std::span<typename E::value_type>out,Op op={}){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return;
assert(sz(a)>=sz(b));
if(sz(a)==sz(b)){
T r{};
for(int i=0;i<sz(a);i++){
r+=a[i]*b[sz(b)-1-i];
}
if(sz(out)>0)op(out[0],r);
return;
}
auto[n,cut]=detail::conv_size_for(sz(a));
E::extend_to(ta,n,a);
E::extend_to(tb,n,b);
auto buf=buffer_pool<T>::get(n);
E::finish(E::mul(ta,tb,n),buf.span());
detail::emit_middle<T>(buf.span(),cut,sz(a),sz(b),
a[0]*b[0],a[sz(a)-1]*b[sz(b)-1],out,op);
}
template<engine E>
vector<typename E::value_type>middle_product(std::span<const typename E::value_type>a,transformed<E>&ta,
std::span<const typename E::value_type>b,transformed<E>&tb){
using T=typename E::value_type;
if(sz(a)==0||sz(b)==0)return{};
assert(sz(a)>=sz(b));
vector<T>r(size_t(sz(a)-sz(b)+1));
middle_product<E>(a,ta,b,tb,std::span<T>(r));
return r;
}
}
#pragma GCC diagnostic pop
// clang-format on
// @formatter:on
#pragma once

#include <algorithm>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <span>
#include <utility>
#include <vector>

#include "fft/engine.hpp"

namespace ecnerwala::fft {

// ==== multiply layer ====
// These are free functions to convolve spans.
//
// The interfaces will typically take input spans, an output span, and an Op representing how to fold the result into the output.
// Output spans may alias one of the input spans.
// Output spans may be shorter than expected; the output will just be truncated.
//
// Some functions may also take E::transformed& objects associated with the input
// spans. These will be lazily filled (see E::extend_to) and used if available.

// Circular convolution mod n (power of 2)
template <engine E, typename Op = assign_op>
void multiply_circular(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, int n, Op op = {}) {
	assert(!(n & (n-1)));
	auto ta = E::transform(a, n);
	auto tb = E::transform(b, n);
	E::finish(E::mul(ta, tb, n), out, op);
}

template <engine E, typename Op = assign_op>
void square_circular(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, int n, Op op = {}) {
	assert(!(n & (n-1)));
	auto ta = E::transform(a, n);
	E::finish(E::sq(ta, n), out, op);
}

namespace detail {
// Arrays of length 2^k + 1 are somewhat common, so we will optimize them by
// multiplying mod 2^k, and fixing up the leading coefficient.

// Helpers to detect and perform this optimization.
struct conv_size { int n; bool cut; };
inline conv_size conv_size_for(int s) {
	int n = nextPow2(s);
	bool cut = (n == 2 * (s - 1));
	return {cut ? n / 2 : n, cut};
}

// Call op while lazily applying the correction if necessary
template <typename T, typename Op>
void emit_linear(std::span<T> buf, int n, int s, bool cut, T c0, std::span<T> out, Op op) {
	T cn{};
	if (cut) {
		cn = buf[0] - c0;
		buf[0] = c0;
	}
	int lim = min(sz(out), min(s, n));
	for (int i = 0; i < lim; i++) op(out[i], buf[i]);
	if (cut && sz(out) >= s) op(out[s-1], cn);
}

// Applies op, diverting the wrapped leading coefficient of a cut product:
// out[0] receives c0 and the wraparound term is captured into cn for the
// caller to emit at out[s-1].
template <typename T, typename Op>
struct cut_op {
	Op op;
	T* out0;
	T c0;
	T& cn;
	void operator()(T& x, T v) const {
		if (&x == out0) { cn = v - c0; v = c0; }
		op(x, v);
	}
};

// finish + emit_linear fused: write the finished product directly into out,
// applying the cut correction in place.
template <engine E, typename P, typename Op = assign_op>
void finish_linear(
	P&& p, int n, int s, bool cut,
	typename E::value_type c0, std::span<typename E::value_type> out, Op op = {}
) {
	using T = typename E::value_type;
	if (sz(out) == 0) return;
	int lim = min(sz(out), min(s, n));
	if (!cut) {
		E::finish(std::move(p), out.subspan(0, lim), op);
	} else {
		T cn{};
		E::finish(std::move(p), out.subspan(0, lim), cut_op<T, Op>{op, &out[0], c0, cn});
		if (sz(out) >= s) op(out[s-1], cn);
	}
}

}

template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	int s = sz(a) + sz(b) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * b[0];
	auto buf = buffer_pool<T>::get(n);
	multiply_circular<E>(a, b, buf.span(), n);
	detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void multiply(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	int s = sz(a) + sz(b) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * b[0];
	E::extend_to(ta, n, a);
	E::extend_to(tb, n, b);
	detail::finish_linear<E>(E::mul(ta, tb, n), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void multiply_add2(std::span<const typename E::value_type> a1, transformed<E>& ta1,
		std::span<const typename E::value_type> b1, transformed<E>& tb1,
		std::span<const typename E::value_type> a2, transformed<E>& ta2,
		std::span<const typename E::value_type> b2, transformed<E>& tb2,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
	int s = sz(a1) + sz(b1) - 1;
	assert(sz(a2) + sz(b2) - 1 == s);
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a1[0] * b1[0] + a2[0] * b2[0];
	E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
	E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
	detail::finish_linear<E>(E::mul2(ta1, tb1, ta2, tb2, n), n, s, cut, c0, out, op);
}

// As multiply_add2, but also outputs the summed pointwise product as a reusable
// transform of the (full-length) result, like multiply_cached.
template <engine E>
void multiply_add2_cached(
		std::span<const typename E::value_type> a1, transformed<E>& ta1,
		std::span<const typename E::value_type> b1, transformed<E>& tb1,
		std::span<const typename E::value_type> a2, transformed<E>& ta2,
		std::span<const typename E::value_type> b2, transformed<E>& tb2,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	assert(sz(a1) > 0 && sz(b1) > 0 && sz(a2) > 0 && sz(b2) > 0);
	int s = sz(a1) + sz(b1) - 1;
	assert(sz(a2) + sz(b2) - 1 == s);
	coeffs.assign(size_t(s), T{});
	t = transformed<E>{};
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a1[0] * b1[0] + a2[0] * b2[0];
		E::extend_to(ta1, n, a1); E::extend_to(tb1, n, b1);
		E::extend_to(ta2, n, a2); E::extend_to(tb2, n, b2);
		auto p = E::mul2(ta1, tb1, ta2, tb2, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		multiply_add2<E>(a1, ta1, b1, tb1, a2, ta2, b2, tb2, std::span<T>(coeffs));
	}
}

// This helper also accepts an output transform which will be populated if it is cheap to do so
template <engine E>
void multiply_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	coeffs.assign(size_t(sz(a) && sz(b) ? sz(a) + sz(b) - 1 : 0), T{});
	t = transformed<E>{};
	if (coeffs.empty()) return;
	int s = sz(coeffs);
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a[0] * b[0];
		E::extend_to(ta, n, a);
		E::extend_to(tb, n, b);
		auto p = E::mul(ta, tb, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		multiply<E>(a, ta, b, tb, std::span<T>(coeffs));
	}
}

template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0) return;
	int s = 2 * sz(a) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * a[0];
	auto buf = buffer_pool<T>::get(n);
	square_circular<E>(a, buf.span(), n);
	detail::emit_linear<T>(buf.span(), n, s, cut, c0, out, op);
}

template <engine E, typename Op = assign_op>
void square(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0) return;
	int s = 2 * sz(a) - 1;
	auto [n, cut] = detail::conv_size_for(s);
	T c0 = a[0] * a[0];
	E::extend_to(ta, n, a);
	detail::finish_linear<E>(E::sq(ta, n), n, s, cut, c0, out, op);
}

// As square, but also outputs the pointwise product as a reusable transform of
// the result (empty when the engine's product isn't a transform).
template <engine E>
void square_cached(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::vector<typename E::value_type>& coeffs, transformed<E>& t) {
	using T = typename E::value_type;
	coeffs.assign(size_t(sz(a) ? 2 * sz(a) - 1 : 0), T{});
	t = transformed<E>{};
	if (coeffs.empty()) return;
	int s = sz(coeffs);
	if constexpr (std::same_as<typename E::product, transformed<E>>) {
		auto [n, cut] = detail::conv_size_for(s);
		T c0 = a[0] * a[0];
		E::extend_to(ta, n, a);
		auto p = E::sq(ta, n);
		auto tp = p;
		detail::finish_linear<E>(std::move(p), n, s, cut, c0, std::span<T>(coeffs));
		t = std::move(tp);
	} else {
		square<E>(a, ta, std::span<T>(coeffs));
	}
}

template <engine E> vector<typename E::value_type> multiply(
		const vector<typename E::value_type>& a, const vector<typename E::value_type>& b) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	vector<T> r(sz(a) + sz(b) - 1);
	multiply<E>(std::span<const T>(a), std::span<const T>(b), std::span<T>(r));
	return r;
}

template <engine E> vector<typename E::value_type> square(const vector<typename E::value_type>& a) {
	using T = typename E::value_type;
	if (sz(a) == 0) return {};
	vector<T> r(2 * sz(a) - 1);
	square<E>(std::span<const T>(a), std::span<T>(r));
	return r;
}

namespace detail {
// emit_linear but for middle_product
template <typename T, typename Op>
void emit_middle(std::span<T> buf, bool cut, int la, int lb, T c0, T ctop, std::span<T> out, Op op) {
	int m = la - lb + 1;
	T cn{};
	if (cut) {
		cn = buf[0] - c0; // for lb == 1 these coincide: slot 0 = c_0 + c_n and ctop = c_n
		buf[lb - 1] -= ctop;
	}
	int lim = min(sz(out), cut ? m - 1 : m);
	for (int t = 0; t < lim; t++) op(out[t], buf[lb - 1 + t]);
	if (cut && sz(out) >= m) op(out[m-1], cn);
}
}

// Middle product (the transposed multiplication): takes only coefficients of a * b which include terms from all of b.
// Must have len(a) >= len(b)
template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, std::span<const typename E::value_type> b,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	assert(sz(a) >= sz(b));
	if (sz(a) == sz(b)) {
		T r{};
		for (int i = 0; i < sz(a); i++) {
			r += a[i] * b[sz(b) - 1 - i];
		}
		if (sz(out) > 0) op(out[0], r);
		return;
	}
	auto [n, cut] = detail::conv_size_for(sz(a));
	auto buf = buffer_pool<T>::get(n);
	multiply_circular<E>(a, b, buf.span(), n);
	detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
			a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}

// TODO: Let's decide whether to keep vector<> returning forms or not; this
// largely depends on whether we think these functions are a public interface or
// merely convenience for value type implementors.
template <engine E> vector<typename E::value_type> middle_product(
		std::span<const typename E::value_type> a, std::span<const typename E::value_type> b) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	assert(sz(a) >= sz(b));
	vector<T> r(size_t(sz(a) - sz(b) + 1));
	middle_product<E>(a, b, std::span<T>(r));
	return r;
}

template <engine E, typename Op = assign_op>
void middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb,
		std::span<typename E::value_type> out, Op op = {}) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return;
	assert(sz(a) >= sz(b));
	if (sz(a) == sz(b)) {
		T r{};
		for (int i = 0; i < sz(a); i++) {
			r += a[i] * b[sz(b) - 1 - i];
		}
		if (sz(out) > 0) op(out[0], r);
		return;
	}
	auto [n, cut] = detail::conv_size_for(sz(a));
	E::extend_to(ta, n, a);
	E::extend_to(tb, n, b);
	auto buf = buffer_pool<T>::get(n);
	E::finish(E::mul(ta, tb, n), buf.span());
	detail::emit_middle<T>(buf.span(), cut, sz(a), sz(b),
			a[0] * b[0], a[sz(a) - 1] * b[sz(b) - 1], out, op);
}

template <engine E>
vector<typename E::value_type> middle_product(std::span<const typename E::value_type> a, transformed<E>& ta,
		std::span<const typename E::value_type> b, transformed<E>& tb) {
	using T = typename E::value_type;
	if (sz(a) == 0 || sz(b) == 0) return {};
	assert(sz(a) >= sz(b));
	vector<T> r(size_t(sz(a) - sz(b) + 1));
	middle_product<E>(a, ta, b, tb, std::span<T>(r));
	return r;
}

/* namespace ecnerwala::fft */ }
Back to top page