Source file src/crypto/internal/fips140/mlkem/generate1024.go

     1  // Copyright 2024 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  //go:build ignore
     6  
     7  package main
     8  
     9  import (
    10  	"flag"
    11  	"go/ast"
    12  	"go/format"
    13  	"go/parser"
    14  	"go/token"
    15  	"log"
    16  	"os"
    17  	"strings"
    18  )
    19  
    20  var replacements = map[string]string{
    21  	"k": "k1024",
    22  
    23  	"CiphertextSize768":       "CiphertextSize1024",
    24  	"EncapsulationKeySize768": "EncapsulationKeySize1024",
    25  
    26  	"encryptionKey": "encryptionKey1024",
    27  	"decryptionKey": "decryptionKey1024",
    28  
    29  	"EncapsulationKey768":    "EncapsulationKey1024",
    30  	"NewEncapsulationKey768": "NewEncapsulationKey1024",
    31  	"parseEK":                "parseEK1024",
    32  
    33  	"kemEncaps":  "kemEncaps1024",
    34  	"pkeEncrypt": "pkeEncrypt1024",
    35  
    36  	"DecapsulationKey768":    "DecapsulationKey1024",
    37  	"NewDecapsulationKey768": "NewDecapsulationKey1024",
    38  	"newKeyFromSeed":         "newKeyFromSeed1024",
    39  
    40  	"kemDecaps":  "kemDecaps1024",
    41  	"pkeDecrypt": "pkeDecrypt1024",
    42  
    43  	"GenerateKey768":         "GenerateKey1024",
    44  	"GenerateKeyInternal768": "GenerateKeyInternal1024",
    45  	"generateKey":            "generateKey1024",
    46  
    47  	"kemKeyGen": "kemKeyGen1024",
    48  	"kemPCT":    "kemPCT1024",
    49  
    50  	"encodingSize4":             "encodingSize5",
    51  	"encodingSize10":            "encodingSize11",
    52  	"ringCompressAndEncode4":    "ringCompressAndEncode5",
    53  	"ringCompressAndEncode10":   "ringCompressAndEncode11",
    54  	"ringDecodeAndDecompress4":  "ringDecodeAndDecompress5",
    55  	"ringDecodeAndDecompress10": "ringDecodeAndDecompress11",
    56  }
    57  
    58  func main() {
    59  	inputFile := flag.String("input", "", "")
    60  	outputFile := flag.String("output", "", "")
    61  	flag.Parse()
    62  
    63  	fset := token.NewFileSet()
    64  	f, err := parser.ParseFile(fset, *inputFile, nil, parser.SkipObjectResolution|parser.ParseComments)
    65  	if err != nil {
    66  		log.Fatal(err)
    67  	}
    68  	cmap := ast.NewCommentMap(fset, f, f.Comments)
    69  
    70  	// Drop header comments.
    71  	cmap[ast.Node(f)] = nil
    72  
    73  	// Remove top-level consts used across the main and generated files.
    74  	var newDecls []ast.Decl
    75  	for _, decl := range f.Decls {
    76  		switch d := decl.(type) {
    77  		case *ast.GenDecl:
    78  			if d.Tok == token.CONST {
    79  				continue // Skip const declarations
    80  			}
    81  			if d.Tok == token.IMPORT {
    82  				cmap[decl] = nil // Drop pre-import comments.
    83  			}
    84  		}
    85  		newDecls = append(newDecls, decl)
    86  	}
    87  	f.Decls = newDecls
    88  
    89  	// Replace identifiers.
    90  	ast.Inspect(f, func(n ast.Node) bool {
    91  		switch x := n.(type) {
    92  		case *ast.Ident:
    93  			if replacement, ok := replacements[x.Name]; ok {
    94  				x.Name = replacement
    95  			}
    96  		}
    97  		return true
    98  	})
    99  
   100  	// Replace identifiers in comments.
   101  	for _, c := range f.Comments {
   102  		for _, l := range c.List {
   103  			for k, v := range replacements {
   104  				if k == "k" {
   105  					continue
   106  				}
   107  				l.Text = strings.ReplaceAll(l.Text, k, v)
   108  			}
   109  		}
   110  	}
   111  
   112  	out, err := os.Create(*outputFile)
   113  	if err != nil {
   114  		log.Fatal(err)
   115  	}
   116  	defer out.Close()
   117  
   118  	out.WriteString("// Code generated by generate1024.go. DO NOT EDIT.\n\n")
   119  
   120  	f.Comments = cmap.Filter(f).Comments()
   121  	err = format.Node(out, fset, f)
   122  	if err != nil {
   123  		log.Fatal(err)
   124  	}
   125  }
   126  

View as plain text