Generate Finite field operations using Montgomery reduction, outputs Rust code suitable for use with Zexe

gen-montops-rust.py

"""
Generates verbose code for montgomery operations
"""

def mont_reduce(fp_bits, limb_bits=64):
    assert fp_bits % limb_bits == 0
    n_limbs = fp_bits // limb_bits
    n_limbs2 = n_limbs * 2
    args = ', '.join([f'{"&mut " if _ != 0 else ""}r{_}: u{limb_bits}' for _ in range(n_limbs2)])
    yield "#[inline]"
    yield f"fn mont_reduce(&self, {args}) {{"

    yield "\t// The Montgomery reduction here is based on Algorithm 14.32 in"
    yield "\t// Handbook of Applied Cryptography"
    yield "\t// <http://cacr.uwaterloo.ca/hac/about/chap14.pdf>."
    yield "\t"
 
    for i in range(n_limbs):
        yield f"\tlet k = r{i}.wrapping_mul(P::INV);"
        yield "\tlet mut carry = 0;"
        yield f"\tfa::mac_with_carry(r{i}, k, P::MODULUS.0[0], &mut carry);"
        for j in range(1, n_limbs):
            yield f"\tr{i+j} = fa::mac_with_carry(r{i+j}, k, P::MODULUS.0[{j}], &mut carry);"
        yield f"\tr{i+j+1} = fa::adc(r{i+j+1}, {'carry2' if i != 0 else '0'}, &mut carry);"
        if i != (n_limbs-1):
            yield "\tlet mut carry2 = carry;"
        yield "\t"

    for i in range(n_limbs, n_limbs2):
        a = n_limbs-(n_limbs2-i)
        yield f"\t(self.0).0[{a}] = r{i};"

    yield "\t"
    yield "\tself.reduce();"
    yield "}"


def mul_assign(fp_bits, limb_bits=64):
    assert fp_bits % limb_bits == 0
    n_limbs = fp_bits // limb_bits
    n_limbs2 = n_limbs * 2

    yield '#[inline]'
    yield 'fn mul_assign(&mut self, other: &Self) {'
    for i in range(n_limbs):
        yield "\tlet mut carry = 0;"
        for j in range(n_limbs):
            if i == 0:
                a = '0'
            else:
                a = f'r{i+j}'
            yield f"\tlet r{i+j} = fa::mac_with_carry({a}, (self.0).0[{i}], (other.0).0[{j}], &mut carry);"
        yield f"\tlet r{i+j+1} = carry;"
        yield "\t"

    mont_reduce_args = ', '.join([f"r{_}" for _ in range(n_limbs*2)])
    yield f"\tself.mont_reduce({mont_reduce_args});"

    yield '}'


def square_in_place(fp_bits, limb_bits=64):
    assert fp_bits % limb_bits == 0
    n_limbs = fp_bits // limb_bits
    n_limbs2 = n_limbs * 2

    yield '#[inline]'
    yield 'fn square_in_place(&mut self) -> &mut Self {'

    for i in range(1, n_limbs):
        yield "\tlet mut carry = 0;"
        for j in range(i, n_limbs):
            if i == 1:
                a = '0'
            else:
                a = f'r{i+j-1}'
            yield f"\tlet r{i+j-1} = fa::mac_with_carry({a}, (self.0).0[{i-1}], (self.0).0[{j}], &mut carry);"
        yield f"\tlet r{j+i} = carry;"
        yield "\t"

    yield f"\tlet r{n_limbs2-1} = r{n_limbs2-2} >> {limb_bits-1};"
    for i in range(n_limbs2-2,1,-1):
        yield f"\tlet r{i} = (r{i} << 1) | (r{i-1} >> {limb_bits-1});"
    yield f"\tlet r1 = r1 << 1;"

    yield "\t"

    yield "\tlet mut carry = 0;"
    for i in range(0, n_limbs2, 2):
        yield f"\tlet r{i} = fa::mac_with_carry(0, (self.0).0[{i//2}], (self.0).0[{i//2}], &mut carry);"
        yield f"\tlet r{i+1} = fa::adc(r{i+1}, 0, &mut carry);"
    yield "\t"

    mont_reduce_args = ', '.join([f"r{_}" for _ in range(n_limbs2)])
    yield f"\tself.mont_reduce({mont_reduce_args});"
    yield "\tself"
    yield "}"


if __name__ == "__main__":
    for line in mul_assign(512, 64):
        print(line)