Sage script to generate an ark_ff field implementation

gen-ff.sage

import sys

LIMB = 64

mod   = int(sys.argv[1]) if len(sys.argv) > 1 else 28948022309329048855892746252171976963363056481941560715954676764349967630337
name  = sys.argv[2] if len(sys.argv) > 2 else 'Fp'
bits  = int(mod).bit_length()
limbs = (bits + LIMB - 1) // LIMB
size  = limbs * LIMB

F = GF(mod)

fp = 'Fp{size}'.format(name=name, size=size)
params = 'Fp{size}Parameters'.format(size=size)
integer = 'BigInteger{size}'.format(size=size)

def adic2(mul_order):
    n = 0
    o = mul_order
    while o % 2 == 0:
        o //= 2
        n += 1
    return n

def to_integer(v):
    elems = []
    for _ in range(limbs):
        elems.append('0x{:016x}'.format(v % 2^LIMB))
        v //= 2^LIMB

    return elems

adicity = adic2(mod - 1)

G = F.multiplicative_generator()
S = 2^adicity
T = G.multiplicative_order() // S
H = G^T

assert H.multiplicative_order() == S
assert T % 2 == 1

def param(name, value):
    return '    const {name}: {integer} = {integer}([\n        {val}\n    ]);'.format(
        integer=integer,
        name=name,
        val=',\n        '.join(to_integer(value))
    )

code = []
code.append('use ark_ff::{{FftParameters, {fp}, {params}, {integer}}};'.format(fp=fp, params=params, integer=integer))
code.append('')
code.append('pub struct {name}Parameters;'.format(name=name))
code.append('')
code.append('pub type {name} = {fp}<{name}Parameters>;'.format(fp=fp, name=name))
code.append('')
code.append('impl {params} for {name}Parameters {{}}'.format(params=params, name=name))
code.append('')
code.append('impl FftParameters for {name}Parameters {{'.format(name=name))
code.append('    type BigInt = {integer};'.format(integer=integer))
code.append('')
code.append('    const TWO_ADICITY: u32 = {adicity};'.format(adicity=adicity))
code.append('')
code.append('    #[rustfmt::skip]')
code.append(param('TWO_ADIC_ROOT_OF_UNITY', int(H)))
code.append('}')
code.append('')
code.append('impl ark_ff::FpParameters for {name}Parameters {{'.format(name=name))
code.append(param('MODULUS', mod))
code.append('')
code.append(param('R', 2^size % mod))
code.append('')
code.append(param('R2', (2^size)^2 % mod))
code.append('')
code.append(param('MODULUS_MINUS_ONE_DIV_TWO', (mod-1) // 2))
code.append('')
code.append(param('T', T))
code.append('')
code.append(param('GENERATOR', int(G)))
code.append('')
code.append(param('T_MINUS_ONE_DIV_TWO', (T-1) // 2))
code.append('')
code.append('    const INV: u64 = 0x{:016x};'.format(Integer(-mod).inverse_mod(2^LIMB)))
code.append('    const CAPACITY: u32 = Self::MODULUS_BITS - 1;')
code.append('    const REPR_SHAVE_BITS: u32 = 1;')
code.append('    const MODULUS_BITS: u32 = {bits};'.format(bits=bits))
code.append('}')

print('\n'.join(code))