#!/usr/bin/env python3 # ============================================================================ # seq2cpp v2 -- sequence -> C++ transpiler, Kitamasa/Fiduccia mode. # # Input (stdin): p a0 a1 a2 a3 ... # Flag (argv): --mul naive|karatsuba|ntt|auto (default: auto) # Output (stdout): a self-contained C++20 file computing a(n) mod p for any # 64-bit n, via X^n mod f(X) and a dot product with the # initial window (Fiduccia/Kitamasa). # # Pipeline: Berlekamp-Massey -> minimal char poly f (order d) # strip the x^m preperiod by shifting (emit a lookup table for # the first shift terms; needed when f would have x | f) # verify in Python against ALL given terms # emit C++ with Montgomery arithmetic (R = 2^32) and one of: # naive schoolbook O(d^2) # karatsuba recursive, schoolbook base case for n <= 4 # ntt direct NTT over F_p; requires an NTT length # N = 2^k >= 2d-1 with N | p-1. If the modulus is # bad, --mul ntt aborts with a message, and the # auto heuristic never selects ntt. # # Auto heuristic (rough crossovers, tunable below): # d <= NAIVE_MAX_D -> naive # NAIVE_MAX_D < d <= KARA_MAX_D -> karatsuba # d > KARA_MAX_D and NTT supported -> ntt # d > KARA_MAX_D and NTT unsupported -> karatsuba # # Limitations: p must be prime, 3 <= p < 2^31 (Montgomery R = 2^32); # repeated factors of f are FINE in this mode (no factorization is done); # fewer than 2d terms triggers a warning (BM may under-determine). # ============================================================================ import sys NAIVE_MAX_D = 40 # auto: naive for d up to this KARA_MAX_D = 300 # auto: karatsuba up to this, ntt beyond (if supported) USAGE = ("usage: seq2cpp.py [--mul naive|karatsuba|ntt|auto] < sequence\n" " (sequence on stdin: 'p a0 a1 a2 ...')") def die(msg, code=1): sys.stderr.write("seq2cpp: ERROR: " + msg + "\n") sys.exit(code) # ------------------------------- Berlekamp-Massey --------------------------- def berlekamp_massey(s, p): """returns (L, C) with C[0]=1 and s[n] + sum_{j=1..L} C[j] s[n-j] = 0 (n>=L).""" n = len(s); L = 0; m = 0; b = 1 C = [0]*n; B = [0]*n; C[0] = B[0] = 1 for i in range(n): m += 1 d = s[i] % p for j in range(1, L+1): d = (d + C[j]*s[i-j]) % p if d == 0: continue T = C[:] coef = d * pow(b, p-2, p) % p for j in range(m, n): C[j] = (C[j] - coef*B[j-m]) % p if 2*L > i: continue L = i+1-L; B = T; b = d; m = 0 return L, [C[j+1] % p for j in range(L)] def char_poly(L, delta, p): """f = x^L + delta[0] x^{L-1} + ... + delta[L-1], monic, low -> high.""" f = [0]*(L+1); f[L] = 1 for j in range(L): f[L-1-j] = (f[L-1-j] + delta[j]) % p # C: s_i + sum C_j s_{i-j} = 0 return f # ------------------------------- primality --------------------------------- def is_prime(n): if n < 2: return False for sp in (2,3,5,7,11,13,17,19,23,29,31,37): if n % sp == 0: return n == sp d = n-1; r = 0 while d % 2 == 0: d //= 2; r += 1 for a in (2,3,5,7,11,13,17,19,23,29,31,37): x = pow(a, d, n) if x in (1, n-1): continue for _ in range(r-1): x = x*x % n if x == n-1: break else: return False return True def factorize(n): fac = []; d = 2 while d*d <= n: while n % d == 0: fac.append(d); n //= d d += 1 if n > 1: fac.append(n) return fac def primitive_root(p): """smallest primitive root mod prime p.""" fac = sorted(set(factorize(p-1))) for g in range(2, p): if all(pow(g, (p-1)//q, p) != 1 for q in fac): return g die("internal: no primitive root found") def ntt_support(d, p): """(ok, N, w, wi, ni): direct NTT over F_p for degree-(d-1) products.""" N = 2 while N < 2*d - 1: N <<= 1 if N > p - 1 or (p - 1) % N != 0: return False, None, None, None, None g = primitive_root(p) w = pow(g, (p-1)//N, p) wi = pow(w, -1, p) ni = pow(N, -1, p) # sanity: w must have exact order N if pow(w, N, p) != 1 or pow(w, N//2, p) == 1: return False, None, None, None, None return True, N, w, wi, ni # ------------------------------- verification ------------------------------- def kitamasa_check(p, f, window, shift, seq): """Simulate the emitted algorithm in Python (plain F_p arithmetic): a(n) = dot( X^{n-shift} mod f , window ) for n >= shift.""" d = len(f) - 1 # iterate r = X^m mod f, m = 0, 1, ... r = [1] + [0]*(d-1) if d > 1 else [(-f[0]) % p] for n in range(shift, len(seq)): val = sum(r[i]*window[i] for i in range(d)) % p if val != seq[n] % p: return False, n # step: r := r * X mod f top = r[d-1] r = [0] + r[:d-1] if top: for i in range(d): r[i] = (r[i] - top*f[i]) % p return True, None # ------------------------------- C++ emission ------------------------------- CPP_HEAD = """// ============================================================ // GENERATED by seq2cpp v2 (Kitamasa/Fiduccia mode). // Input: p = @P@, @N@ given terms, minimal recurrence order d = @D@. // Multiplication backend: @BACKEND@@WHYBACKEND@ // Minimal polynomial of the (shifted) sequence: @FPOLY@ // a(n) = sum_i c_i * a(@SHIFT@+i) where X^(n-@SHIFT@) mod f(X) = sum c_i X^i @PRENOTE@// Compile: g++ -std=c++20 -O2 // Benchmark: g++ -std=c++20 -O2 -DBENCH=100000 ... ; echo n | ./bench // The static_asserts verify the sequence at COMPILE TIME. // ============================================================ #include #include #ifdef BENCH #include #endif using u64 = std::uint64_t; // hint for the compiler to unroll the (compile-time constant sized) loops below #if defined(__GNUC__) || defined(__clang__) #define UNROLL _Pragma("GCC unroll 64") #else #define UNROLL #endif // --------------- Montgomery arithmetic mod P, R = 2^32 --------------------- inline constexpr u64 P = @P@; inline constexpr u64 RMASK = (1ull << 32) - 1; static consteval u64 neg_inverse() { // -P^{-1} mod 2^32, by Newton iteration u64 x = 1; for (int i = 0; i < 5; ++i) x *= 2 - P * x; return -x; } inline constexpr u64 PINV = neg_inverse(); inline constexpr u64 R2 = []{ u64 r = (1ull << 32) % P; return r * r % P; }(); static constexpr u64 REDC(u64 T) { // T < P*R; returns T*R^{-1} mod P u64 m = (T & RMASK) * PINV & RMASK; u64 t = (T + m * P) >> 32; return t >= P ? t - P : t; } static constexpr u64 toM(u64 x) { return REDC((x % P) * R2); } static constexpr u64 fromM(u64 a) { return REDC(a); } // ---------------- recurrence data ------------------------------------------ inline constexpr int D = @D@; // order of the recurrence inline constexpr u64 RED[D] = { @RED@ }; // X^D = sum RED[i] X^i (mod P, f(X)); toM() at compile time inline constexpr u64 AW[D] = { @AW@ }; // window a(@SHIFT@..@SHIFT@+D-1); toM() at compile time struct Poly { u64 c[D]; }; // coefficients in [0, P), Montgomery @BACKENDBLOCK@ // ---------------- Fiduccia: X^n mod f(X), then the dot product -------------- static constexpr Poly pone() { Poly q{}; q.c[0] = toM(1); return q; } static constexpr Poly px() { Poly q{}; if constexpr (D == 1) q.c[0] = toM(@XMODF@); // X mod f(X) else q.c[1] = toM(1); return q; } static constexpr Poly xpow(u64 n) { // X^n mod f(X), square-and-multiply Poly r = pone(), b = px(); while (n) { if (n & 1) r = pmul(r, b); b = pmul(b, b); n >>= 1; } return r; } """ def fold_lines(T2, D): return f""" for (int k = {T2-1}; k >= D; --k) {{ // fold X^k = X^(k-D) * X^D u64 c = t[k] % P; if (c) {{ t[k] = 0; UNROLL for (int i = 0; i < D; ++i) t[k - D + i] += REDC(c * RED[i]); }} }} Poly r{{}}; UNROLL for (int i = 0; i < D; ++i) r.c[i] = t[i] % P; return r;""" PMUL_NAIVE = """// ---- backend: NAIVE schoolbook (D^2 REDCs per product) ---- static constexpr Poly pmul(const Poly& a, const Poly& b) { u64 t[@T2@] = {}; UNROLL for (int i = 0; i < D; ++i) if (a.c[i]) UNROLL for (int j = 0; j < D; ++j) t[i + j] += REDC(a.c[i] * b.c[j]); // each < P; sums < D*P @FOLD@ }""" PMUL_KARA = """// ---- backend: KARATSUBA (template recursion on the length, base n <= 4) ---- // every recursive call has a compile-time known length N, so all loops are of // fixed size (unrollable) and the scratch buffers are sized exactly. template static constexpr void kmul(const u64* a, const u64* b, u64* out) { // a, b length N, entries in [0, P); writes product of length 2N-1, [0, P) if constexpr (N <= 4) { UNROLL for (int i = 0; i < 2*N - 1; ++i) out[i] = 0; UNROLL for (int i = 0; i < N; ++i) UNROLL for (int j = 0; j < N; ++j) out[i + j] = (out[i + j] + REDC(a[i] * b[j])) % P; } else { constexpr int h = N / 2, g = N - h; // lo = a[0..h), hi = a[h..N) (len g) u64 z0[2*h - 1] = {}, z1[2*g - 1] = {}, z2[2*g - 1] = {}, m[2*g - 1] = {}; u64 sa[g] = {}, sb[g] = {}; UNROLL for (int i = 0; i < h; ++i) { sa[i] = (a[i] + a[h + i]) % P; // sa = lo + hi (length g) sb[i] = (b[i] + b[h + i]) % P; } if constexpr (g > h) { // odd N: pass the top term sa[h] = a[2*h]; sb[h] = b[2*h]; } kmul(sa, sb, z1); // (lo+hi)(lo'+hi') kmul(a, b, z0); // lo*lo' kmul(a + h, b + h, z2); // hi*hi' UNROLL for (int i = 0; i < 2*g - 1; ++i) { // m = z1 - z0 - z2 (mod P) u64 v = z1[i]; if (i < 2*h - 1) v = (v + P - z0[i]) % P; v = (v + P - z2[i]) % P; m[i] = v; } UNROLL for (int i = 0; i < 2*N - 1; ++i) { // out = z0 + X^h m + X^(2h) z2 u64 v = 0; if (i < 2*h - 1) v = z0[i]; if (i >= h && i < h + 2*g - 1) v = (v + m[i - h]) % P; if (i >= 2*h) v = (v + z2[i - 2*h]) % P; out[i] = v; } } } static constexpr Poly pmul(const Poly& a, const Poly& b) { u64 t[@T2@]; kmul<@D@>(a.c, b.c, t); @FOLD@ }""" PMUL_NTT = """// ---- backend: NTT over F_p (direct; N = 2^k >= 2D-1, N | P-1) ---- inline constexpr int NB = @NNTT@; inline constexpr u64 NW = @W@; // primitive NB-th root of unity mod P inline constexpr u64 NWI = @WI@; // its inverse inline constexpr u64 NNI = @NI@; // NB^{-1} mod P static constexpr u64 mpow(u64 b, u64 e) { u64 r = 1; b %= P; while (e) { if (e & 1) r = r * b % P; b = b * b % P; e >>= 1; } return r; } static constexpr void ntt(u64* a, u64 root) { // iterative Cooley-Tukey mod P for (int i = 1, j = 0; i < NB; ++i) { int bit = NB >> 1; for (; j & bit; bit >>= 1) j ^= bit; j ^= bit; if (i < j) { u64 t = a[i]; a[i] = a[j]; a[j] = t; } } for (int len = 2; len <= NB; len <<= 1) { u64 wl = mpow(root, NB / len); // primitive len-th root for (int i = 0; i < NB; i += len) { u64 w = 1; UNROLL for (int j = 0; j < len / 2; ++j) { u64 u = a[i + j]; u64 v = a[i + j + len/2] * w % P; a[i + j] = (u + v) % P; a[i + j + len / 2] = (u + P - v) % P; w = w * wl % P; } } } } static constexpr Poly pmul(const Poly& a, const Poly& b) { u64 fa[@NNTT@] = {}, fb[@NNTT@] = {}; UNROLL for (int i = 0; i < D; ++i) { fa[i] = fromM(a.c[i]); fb[i] = fromM(b.c[i]); } ntt(fa, NW); ntt(fb, NW); UNROLL for (int i = 0; i < NB; ++i) fa[i] = fa[i] * fb[i] % P; ntt(fa, NWI); // inverse transform (unscaled) u64 t[@T2@]; UNROLL for (int k = 0; k < @T2@; ++k) t[k] = toM(fa[k] * NNI % P); @FOLD@ }""" CPP_TAIL = """ @PREDECL@ // the n-th term of the sequence, mod P (plain residue in [0, P)) static constexpr u64 a_n(u64 n) { @PREBRANCH@ Poly q = xpow(n - @SHIFT@ULL); u64 s = 0; UNROLL for (int i = 0; i < D; ++i) s += REDC(q.c[i] * AW[i]); return fromM(s); } """ MAIN_TAIL = """ int main() { unsigned long long n; #ifdef BENCH if (std::scanf("%llu", &n) != 1) return 1; u64 sink = 0; std::fprintf(stderr, "bench: %d iterations\\n", (int)BENCH); std::timespec t0, t1; std::timespec_get(&t0, TIME_UTC); for (int i = 0; i < BENCH; ++i) sink += a_n(n); std::timespec_get(&t1, TIME_UTC); double ns = (t1.tv_sec - t0.tv_sec) * 1e9 + (t1.tv_nsec - t0.tv_nsec); std::fprintf(stderr, "%.1f ns per a_n (sink %llu)\\n", ns / BENCH, (unsigned long long)sink); return 0; #else while (std::scanf("%llu", &n) == 1) std::printf("%llu\\n", (unsigned long long)a_n(n)); return 0; #endif } """ ZERO_PROG = """// ============================================================ // GENERATED by seq2cpp v2 (Kitamasa/Fiduccia mode). // Input: p = @P@, @N@ given terms. // The sequence is zero from n = @SHIFT@ on; no recurrence needed. // ============================================================ #include #include #ifdef BENCH #include #endif using u64 = std::uint64_t; inline constexpr u64 P = @P@; @PREDECL@ static constexpr u64 a_n(u64 n) { @PREBRANCH@ return 0; } @ASSERTS@ int main() { unsigned long long n; while (std::scanf("%llu", &n) == 1) std::printf("%llu\\n", (unsigned long long)a_n(n)); return 0; } """ def fmt_poly_cpp(f, p): d = len(f) - 1 terms = [] for i in range(d, -1, -1): c = f[i] % p if c == 0: continue if i == 0: terms.append(str(c)) elif i == 1: terms.append("X" if c == 1 else f"{c}X") else: terms.append(f"X^{i}" if c == 1 else f"{c}X^{i}") return "// f(X) = " + (" + ".join(terms).replace("+ -", "- ") if terms else "0") def main(): # ---- flag ---- backend = "auto" args = sys.argv[1:] if args: if args[0] == "--mul" and len(args) >= 2: backend = args[1]; args = args[2:] elif args[0].startswith("--mul="): backend = args[0][6:]; args = args[1:] else: die("unknown argument: %r\n%s" % (args[0], USAGE)) if args: die("unexpected extra arguments\n" + USAGE) if backend not in ("naive", "karatsuba", "ntt", "auto"): die("unknown backend %r (naive|karatsuba|ntt|auto)" % backend) # ---- input ---- data = sys.stdin.read().split() if not data: die("no input\n" + USAGE) p = int(data[0]) seq = [int(x) % p for x in data[1:]] if not is_prime(p): die("%d is not prime" % p) if p == 2: die("p = 2 is not supported (Montgomery needs odd P)") if p >= 2**31: die("p >= 2^31 is not supported (Montgomery R = 2^32 width)") if len(seq) < 2: die("need at least 2 terms") # ---- Berlekamp-Massey with preperiod stripping ---- shift = 0 while True: L, delta = berlekamp_massey(seq[shift:], p) if L == 0: emit_zero_program(p, seq, shift) return f = char_poly(L, delta, p) if f[0] % p != 0: break if shift + 2*L > len(seq): die("cannot strip preperiod: not enough terms " "(min poly starts with X^%d but only %d terms given)" % (L, len(seq))) shift += 1 # strip one unit of preperiod, retry d = L if len(seq) - shift < 2*d: sys.stderr.write("seq2cpp: WARNING: only %d usable terms for order %d; " "BM may under-determine the recurrence\n" % (len(seq) - shift, d)) # ---- verify the recurrence against ALL given terms ---- window = [seq[shift+i] % p for i in range(d)] ok, bad = kitamasa_check(p, f, window, shift, seq) if not ok: die("verification failed at n=%d (internal error)" % bad) # ---- backend selection ---- ntt_ok, N, w, wi, ni = ntt_support(d, p) if backend == "auto": if d <= NAIVE_MAX_D: backend = "naive"; why = " (auto: order %d <= %d)" % (d, NAIVE_MAX_D) elif d <= KARA_MAX_D or not ntt_ok: backend = "karatsuba" why = (" (auto: order %d in (%d, %d])" % (d, NAIVE_MAX_D, KARA_MAX_D) if d <= KARA_MAX_D else " (auto: order %d > %d but NTT unsupported mod %d)" % (d, KARA_MAX_D, p)) else: backend = "ntt"; why = " (auto: order %d > %d and NTT supported)" % (d, KARA_MAX_D) else: why = "" if backend == "ntt" and not ntt_ok: need = 2 while need < 2*d - 1: need <<= 1 v2 = 0; q = p - 1 while q % 2 == 0: q //= 2; v2 += 1 die("--mul ntt: modulus %d is bad for direct NTT: need NTT length " "N = %d (>= 2*%d-1) with N | p-1, but p-1 = 2^%d * %d only " "supports N <= %d" % (p, need, d, v2, q, 1 << v2)) # ---- emit ---- T2 = 2*d - 1 red = ", ".join("toM(%d)" % ((p - f[i]) % p) for i in range(d)) aw = ", ".join("toM(%d)" % (window[i] % p) for i in range(d)) predecl = "" prenote = "" if shift: pre = ", ".join(str(seq[i] % p) for i in range(shift)) predecl = ("inline constexpr u64 PRE[%d] = { %s }; // a(0..%d), plain values\n" % (shift, pre, shift-1)) prenote = ("// NOTE: the minimal polynomial has an X^%d preperiod factor; terms\n" "// n = 0..%d are emitted as a lookup table (shift = %d).\n" % (shift, shift-1, shift)) xmodf = (p - f[0]) % p if d == 1 else 0 if backend == "naive": block = PMUL_NAIVE elif backend == "karatsuba": block = PMUL_KARA else: block = PMUL_NTT block = (block.replace("@T2@", str(T2)).replace("@D@", str(d)) .replace("@NNTT@", str(N)).replace("@W@", str(w)) .replace("@WI@", str(wi)).replace("@NI@", str(ni))) block = block.replace("@FOLD@", fold_lines(T2, d)) out = (CPP_HEAD.replace("@P@", str(p)).replace("@N@", str(len(seq))) .replace("@D@", str(d)).replace("@SHIFT@", str(shift)) .replace("@BACKEND@", backend).replace("@WHYBACKEND@", why) .replace("@FPOLY@", fmt_poly_cpp(f, p)) .replace("@PRENOTE@", prenote) .replace("@RED@", red).replace("@AW@", aw) .replace("@XMODF@", str(xmodf)) .replace("@BACKENDBLOCK@", block)) print(out) prebranch = (" if (n < %dULL) return PRE[n]; // preperiod lookup\n" % shift) if shift else "" print(CPP_TAIL.replace("@SHIFT@", str(shift)).replace("@PREDECL@", predecl) .replace("@PREBRANCH@", prebranch)) # static asserts on up to 100 given terms cap = min(len(seq), 100) for n in range(cap): print("static_assert(a_n(%d) == %d);" % (n, seq[n] % p)) if len(seq) > cap: print("// (%d further terms verified at generation time, not asserted)" % (len(seq) - cap)) print(MAIN_TAIL) def emit_zero_program(p, seq, shift): predecl = "" if shift: pre = ", ".join(str(seq[i] % p) for i in range(shift)) predecl = "inline constexpr u64 PRE[%d] = { %s };\n" % (shift, pre) asserts = "" for n in range(min(len(seq), 100)): asserts += "static_assert(a_n(%d) == %d);\n" % (n, seq[n] % p) prebranch = (" if (n < %dULL) return PRE[n];\n" % shift) if shift else "" print(ZERO_PROG.replace("@P@", str(p)).replace("@N@", str(len(seq))) .replace("@PREDECL@", predecl).replace("@PREBRANCH@", prebranch) .replace("@ASSERTS@", asserts)) if __name__ == "__main__": main()