generate1024.mx raw

   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  	"bytes"
  18  )
  19  
  20  var replacements = map[string][]byte{
  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 = bytes.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