#!/usr/bin/env python3

import sys
import datetime

# (C) 2018,2023,2025 Jack Lloyd
# Botan is released under the Simplified BSD License (see license.txt)

# Used to generate src/lib/math/mp/mp_monty_n.cpp

def monty_redc_code(n, p_dash1=False):

    lines = []

    fn_name = "bigint_monty_redc_pdash1" if p_dash1 else "bigint_monty_redc"

    if p_dash1:
        hdr = "void %s_%d(word r[%d], const word z[%d], const word p[%d], word ws[%d]) {\n" % (fn_name, n, n, 2*n, n, n)
    else:
        hdr = "void %s_%d(word r[%d], const word z[%d], const word p[%d], word p_dash, word ws[%d]) {\n" % (fn_name, n, n, 2*n, n, n)

    ftr = "\n}\n\n"

    lines.append("word3<word> accum;")
    lines.append("accum.add(z[0]);")
    if p_dash1:
        lines.append("ws[0] = accum.monty_step_pdash1();")
    else:
        lines.append("ws[0] = accum.monty_step(p[0], p_dash);")

    for i in range(1, n):
        for j in range(0, i):
            lines.append("accum.mul(ws[%d], p[%d]);" % (j, i-j))

        lines.append("accum.add(z[%d]);" % (i))
        if p_dash1:
            lines.append("ws[%d] = accum.monty_step_pdash1();" % (i))
        else:
            lines.append("ws[%d] = accum.monty_step(p[0], p_dash);" % (i))

    for i in range(0, n - 1):
        for j in range(i + 1, n):
            lines.append("accum.mul(ws[%d], p[%d]);" % (j, n + i-j))

        lines.append("accum.add(z[%d]);" % (n+i))
        lines.append("ws[%d] = accum.extract();" % (i))

    lines.append("accum.add(z[%d]);" % (2*n-1))

    lines.append("ws[%d] = accum.extract();" % (n - 1))
    lines.append("const word w1 = accum.extract();")

    lines.append("bigint_monty_maybe_sub<%d>(r, w1, ws, p);" % (n))

    return hdr + "\n".join(["   %s" % (line) for line in lines]) + ftr

def main(args = None):
    if args is None:
        args = sys.argv

    if len(args) <= 1:
        sizes = [4, 6, 8, 12, 16, 24, 32]
    else:
        sizes = map(int, args[1:])

    header = """/*
* This file was automatically generated by %s on %s
* All manual changes will be lost. Edit the script instead.
*
* Botan is released under the Simplified BSD License (see license.txt)
*/

#include <botan/internal/mp_core.h>

namespace Botan {

""" % (sys.argv[0], datetime.date.today().strftime("%Y-%m-%d"))

    footer = "}  // namespace Botan\n"

    with open('./src/lib/math/mp/mp_monty_n.cpp', encoding='utf8', mode='w') as monty_n:
        monty_n.write(header)

        for n in sizes:
            monty_n.write(monty_redc_code(n))

        monty_n.write(footer)

    return 0

if __name__ == '__main__':
    sys.exit(main())


