gmsm/mlkem/generate.go

// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

//go:build ignore

package main

import (
	"flag"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"log"
	"os"
	"strings"
)

var replacements512 = map[string]string{
	"k":  "k512",
	"η1": "η1_512",
	"η2": "η2_512",

	"CiphertextSize768":       "CiphertextSize512",
	"EncapsulationKeySize768": "EncapsulationKeySize512",
	"DecapsulationKeySize768": "DecapsulationKeySize512",

	"encryptionKey": "encryptionKey512",
	"decryptionKey": "decryptionKey512",

	"EncapsulationKey768":    "EncapsulationKey512",
	"NewEncapsulationKey768": "NewEncapsulationKey512",
	"parseEK":                "parseEK512",

	"kemEncaps":  "kemEncaps512",
	"pkeEncrypt": "pkeEncrypt512",

	"DecapsulationKey768":            "DecapsulationKey512",
	"NewDecapsulationKey768":         "NewDecapsulationKey512",
	"NewDecapsulationKeyFromSeed768": "NewDecapsulationKeyFromSeed512",
	"newKeyFromSeed":                 "newKeyFromSeed512",

	"kemDecaps":  "kemDecaps512",
	"pkeDecrypt": "pkeDecrypt512",

	"GenerateKey768":         "GenerateKey512",
	"GenerateKeyInternal768": "GenerateKeyInternal512",
	"generateKey":            "generateKey512",

	"kemKeyGen": "kemKeyGen512",
}

var replacements1024 = map[string]string{
	"k":  "k1024",
	"η1": "η1_1024",
	"η2": "η2_1024",

	"CiphertextSize768":       "CiphertextSize1024",
	"EncapsulationKeySize768": "EncapsulationKeySize1024",
	"DecapsulationKeySize768": "DecapsulationKeySize1024",

	"encryptionKey": "encryptionKey1024",
	"decryptionKey": "decryptionKey1024",

	"EncapsulationKey768":    "EncapsulationKey1024",
	"NewEncapsulationKey768": "NewEncapsulationKey1024",
	"parseEK":                "parseEK1024",

	"kemEncaps":  "kemEncaps1024",
	"pkeEncrypt": "pkeEncrypt1024",

	"DecapsulationKey768":            "DecapsulationKey1024",
	"NewDecapsulationKey768":         "NewDecapsulationKey1024",
	"NewDecapsulationKeyFromSeed768": "NewDecapsulationKeyFromSeed1024",
	"newKeyFromSeed":                 "newKeyFromSeed1024",

	"kemDecaps":  "kemDecaps1024",
	"pkeDecrypt": "pkeDecrypt1024",

	"GenerateKey768":         "GenerateKey1024",
	"GenerateKeyInternal768": "GenerateKeyInternal1024",
	"generateKey":            "generateKey1024",

	"kemKeyGen": "kemKeyGen1024",
	"kemPCT":    "kemPCT1024",

	"encodingSize4":             "encodingSize5",
	"encodingSize10":            "encodingSize11",
	"ringCompressAndEncode4":    "ringCompressAndEncode5",
	"ringCompressAndEncode10":   "ringCompressAndEncode11",
	"ringDecodeAndDecompress4":  "ringDecodeAndDecompress5",
	"ringDecodeAndDecompress10": "ringDecodeAndDecompress11",
}

func replaceIdentifiers(inputFile, outputFile *string, replacements map[string]string) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, *inputFile, nil, parser.SkipObjectResolution|parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}
	cmap := ast.NewCommentMap(fset, f, f.Comments)

	// Drop header comments.
	cmap[ast.Node(f)] = nil

	// Remove top-level consts used across the main and generated files.
	var newDecls []ast.Decl
	for _, decl := range f.Decls {
		switch d := decl.(type) {
		case *ast.GenDecl:
			if d.Tok == token.CONST {
				continue // Skip const declarations
			}
			if d.Tok == token.IMPORT {
				cmap[decl] = nil // Drop pre-import comments.
			}
		}
		newDecls = append(newDecls, decl)
	}
	f.Decls = newDecls

	// Replace identifiers.
	ast.Inspect(f, func(n ast.Node) bool {
		switch x := n.(type) {
		case *ast.Ident:
			if replacement, ok := replacements[x.Name]; ok {
				x.Name = replacement
			}
		}
		return true
	})

	// Replace identifiers in comments.
	for _, c := range f.Comments {
		for _, l := range c.List {
			for k, v := range replacements {
				if k == "k" || k == "l" {
					continue
				}
				l.Text = strings.ReplaceAll(l.Text, k, v)
			}
		}
	}

	out, err := os.Create(*outputFile)
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	out.WriteString(`// Copyright 2025 Sun Yimin. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.
// Code generated by generate.go. DO NOT EDIT.

//go:build go1.24

`)

	f.Comments = cmap.Filter(f).Comments()
	err = format.Node(out, fset, f)
	if err != nil {
		log.Fatal(err)
	}

}

func main() {
	inputFile := flag.String("mlkem768", "", "")
	outputFile512 := flag.String("mlkem512", "", "")
	outputFile1024 := flag.String("mlkem1024", "", "")
	flag.Parse()

	replaceIdentifiers(inputFile, outputFile512, replacements512)
	replaceIdentifiers(inputFile, outputFile1024, replacements1024)
}