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  	"decapsulationKeySize768": "decapsulationKeySize1024",
    26  
    27  	"encryptionKey": "encryptionKey1024",
    28  	"decryptionKey": "decryptionKey1024",
    29  
    30  	"EncapsulationKey768":    "EncapsulationKey1024",
    31  	"NewEncapsulationKey768": "NewEncapsulationKey1024",
    32  	"parseEK":                "parseEK1024",
    33  
    34  	"kemEncaps":  "kemEncaps1024",
    35  	"pkeEncrypt": "pkeEncrypt1024",
    36  
    37  	"DecapsulationKey768":               "DecapsulationKey1024",
    38  	"NewDecapsulationKey768":            "NewDecapsulationKey1024",
    39  	"TestingOnlyNewDecapsulationKey768": "TestingOnlyNewDecapsulationKey1024",
    40  	"newKeyFromSeed":                    "newKeyFromSeed1024",
    41  	"TestingOnlyExpandedBytes768":       "TestingOnlyExpandedBytes1024",
    42  
    43  	"kemDecaps":  "kemDecaps1024",
    44  	"pkeDecrypt": "pkeDecrypt1024",
    45  
    46  	"GenerateKey768":         "GenerateKey1024",
    47  	"GenerateKeyInternal768": "GenerateKeyInternal1024",
    48  	"generateKey":            "generateKey1024",
    49  
    50  	"kemKeyGen": "kemKeyGen1024",
    51  	"kemPCT":    "kemPCT1024",
    52  
    53  	"encodingSize4":             "encodingSize5",
    54  	"encodingSize10":            "encodingSize11",
    55  	"ringCompressAndEncode4":    "ringCompressAndEncode5",
    56  	"ringCompressAndEncode10":   "ringCompressAndEncode11",
    57  	"ringDecodeAndDecompress4":  "ringDecodeAndDecompress5",
    58  	"ringDecodeAndDecompress10": "ringDecodeAndDecompress11",
    59  }
    60  
    61  func main() {
    62  	inputFile := flag.String("input", "", "")
    63  	outputFile := flag.String("output", "", "")
    64  	flag.Parse()
    65  
    66  	fset := token.NewFileSet()
    67  	f, err := parser.ParseFile(fset, *inputFile, nil, parser.SkipObjectResolution|parser.ParseComments)
    68  	if err != nil {
    69  		log.Fatal(err)
    70  	}
    71  	cmap := ast.NewCommentMap(fset, f, f.Comments)
    72  
    73  	// Drop header comments.
    74  	cmap[ast.Node(f)] = nil
    75  
    76  	// Remove top-level consts used across the main and generated files.
    77  	var newDecls []ast.Decl
    78  	for _, decl := range f.Decls {
    79  		switch d := decl.(type) {
    80  		case *ast.GenDecl:
    81  			if d.Tok == token.CONST {
    82  				continue // Skip const declarations
    83  			}
    84  			if d.Tok == token.IMPORT {
    85  				cmap[decl] = nil // Drop pre-import comments.
    86  			}
    87  		}
    88  		newDecls = append(newDecls, decl)
    89  	}
    90  	f.Decls = newDecls
    91  
    92  	// Replace identifiers.
    93  	ast.Inspect(f, func(n ast.Node) bool {
    94  		switch x := n.(type) {
    95  		case *ast.Ident:
    96  			if replacement, ok := replacements[x.Name]; ok {
    97  				x.Name = replacement
    98  			}
    99  		}
   100  		return true
   101  	})
   102  
   103  	// Replace identifiers in comments.
   104  	for _, c := range f.Comments {
   105  		for _, l := range c.List {
   106  			for k, v := range replacements {
   107  				if k == "k" {
   108  					continue
   109  				}
   110  				l.Text = strings.ReplaceAll(l.Text, k, v)
   111  			}
   112  		}
   113  	}
   114  
   115  	out, err := os.Create(*outputFile)
   116  	if err != nil {
   117  		log.Fatal(err)
   118  	}
   119  	defer out.Close()
   120  
   121  	out.WriteString("// Code generated by generate1024.go. DO NOT EDIT.\n\n")
   122  
   123  	f.Comments = cmap.Filter(f).Comments()
   124  	err = format.Node(out, fset, f)
   125  	if err != nil {
   126  		log.Fatal(err)
   127  	}
   128  }
   129
View as plain text