github.com/zxy12/go_duplicate_112_new@v0.0.0-20200807091221-747231827200/src/sort/genzfunc.go (about)

     1  // Copyright 2016 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  // +build ignore
     6  
     7  // This program is run via "go generate" (via a directive in sort.go)
     8  // to generate zfuncversion.go.
     9  //
    10  // It copies sort.go to zfuncversion.go, only retaining funcs which
    11  // take a "data Interface" parameter, and renaming each to have a
    12  // "_func" suffix and taking a "data lessSwap" instead. It then rewrites
    13  // each internal function call to the appropriate _func variants.
    14  
    15  package main
    16  
    17  import (
    18  	"bytes"
    19  	"go/ast"
    20  	"go/format"
    21  	"go/parser"
    22  	"go/token"
    23  	"io/ioutil"
    24  	"log"
    25  	"regexp"
    26  )
    27  
    28  var fset = token.NewFileSet()
    29  
    30  func main() {
    31  	af, err := parser.ParseFile(fset, "sort.go", nil, 0)
    32  	if err != nil {
    33  		log.Fatal(err)
    34  	}
    35  	af.Doc = nil
    36  	af.Imports = nil
    37  	af.Comments = nil
    38  
    39  	var newDecl []ast.Decl
    40  	for _, d := range af.Decls {
    41  		fd, ok := d.(*ast.FuncDecl)
    42  		if !ok {
    43  			continue
    44  		}
    45  		if fd.Recv != nil || fd.Name.IsExported() {
    46  			continue
    47  		}
    48  		typ := fd.Type
    49  		if len(typ.Params.List) < 1 {
    50  			continue
    51  		}
    52  		arg0 := typ.Params.List[0]
    53  		arg0Name := arg0.Names[0].Name
    54  		arg0Type := arg0.Type.(*ast.Ident)
    55  		if arg0Name != "data" || arg0Type.Name != "Interface" {
    56  			continue
    57  		}
    58  		arg0Type.Name = "lessSwap"
    59  
    60  		newDecl = append(newDecl, fd)
    61  	}
    62  	af.Decls = newDecl
    63  	ast.Walk(visitFunc(rewriteCalls), af)
    64  
    65  	var out bytes.Buffer
    66  	if err := format.Node(&out, fset, af); err != nil {
    67  		log.Fatalf("format.Node: %v", err)
    68  	}
    69  
    70  	// Get rid of blank lines after removal of comments.
    71  	src := regexp.MustCompile(`\n{2,}`).ReplaceAll(out.Bytes(), []byte("\n"))
    72  
    73  	// Add comments to each func, for the lost reader.
    74  	// This is so much easier than adding comments via the AST
    75  	// and trying to get position info correct.
    76  	src = regexp.MustCompile(`(?m)^func (\w+)`).ReplaceAll(src, []byte("\n// Auto-generated variant of sort.go:$1\nfunc ${1}_func"))
    77  
    78  	// Final gofmt.
    79  	src, err = format.Source(src)
    80  	if err != nil {
    81  		log.Fatalf("format.Source: %v on\n%s", err, src)
    82  	}
    83  
    84  	out.Reset()
    85  	out.WriteString(`// Code generated from sort.go using genzfunc.go; DO NOT EDIT.
    86  
    87  // Copyright 2016 The Go Authors. All rights reserved.
    88  // Use of this source code is governed by a BSD-style
    89  // license that can be found in the LICENSE file.
    90  
    91  `)
    92  	out.Write(src)
    93  
    94  	const target = "zfuncversion.go"
    95  	if err := ioutil.WriteFile(target, out.Bytes(), 0644); err != nil {
    96  		log.Fatal(err)
    97  	}
    98  }
    99  
   100  type visitFunc func(ast.Node) ast.Visitor
   101  
   102  func (f visitFunc) Visit(n ast.Node) ast.Visitor { return f(n) }
   103  
   104  func rewriteCalls(n ast.Node) ast.Visitor {
   105  	ce, ok := n.(*ast.CallExpr)
   106  	if ok {
   107  		rewriteCall(ce)
   108  	}
   109  	return visitFunc(rewriteCalls)
   110  }
   111  
   112  func rewriteCall(ce *ast.CallExpr) {
   113  	ident, ok := ce.Fun.(*ast.Ident)
   114  	if !ok {
   115  		// e.g. skip SelectorExpr (data.Less(..) calls)
   116  		return
   117  	}
   118  	// skip casts
   119  	if ident.Name == "int" || ident.Name == "uint" {
   120  		return
   121  	}
   122  	if len(ce.Args) < 1 {
   123  		return
   124  	}
   125  	ident.Name += "_func"
   126  }