github.com/SagerNet/gvisor@v0.0.0-20210707092255-7731c139d75c/tools/go_generics/remove.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"go/ast"
    19  	"go/token"
    20  )
    21  
    22  type typeSet map[string]struct{}
    23  
    24  // isTypeOrPointerToType determines if the given AST expression represents a
    25  // type or a pointer to a type that exists in the provided type set.
    26  func isTypeOrPointerToType(set typeSet, expr ast.Expr, starCount int) bool {
    27  	switch e := expr.(type) {
    28  	case *ast.Ident:
    29  		_, ok := set[e.Name]
    30  		return ok
    31  	case *ast.StarExpr:
    32  		if starCount > 1 {
    33  			return false
    34  		}
    35  		return isTypeOrPointerToType(set, e.X, starCount+1)
    36  	case *ast.ParenExpr:
    37  		return isTypeOrPointerToType(set, e.X, starCount)
    38  	default:
    39  		return false
    40  	}
    41  }
    42  
    43  // isMethodOf determines if the given function declaration is a method of one
    44  // of the types in the provided type set. To do that, it checks if the function
    45  // has a receiver and that its type is either T or *T, where T is a type that
    46  // exists in the set. This is per the spec:
    47  //
    48  // That parameter section must declare a single parameter, the receiver. Its
    49  // type must be of the form T or *T (possibly using parentheses) where T is a
    50  // type name. The type denoted by T is called the receiver base type; it must
    51  // not be a pointer or interface type and it must be declared in the same
    52  // package as the method.
    53  func isMethodOf(set typeSet, f *ast.FuncDecl) bool {
    54  	// If the function doesn't have exactly one receiver, then it's
    55  	// definitely not a method.
    56  	if f.Recv == nil || len(f.Recv.List) != 1 {
    57  		return false
    58  	}
    59  
    60  	return isTypeOrPointerToType(set, f.Recv.List[0].Type, 0)
    61  }
    62  
    63  // removeTypeDefinitions removes the definition of all types contained in the
    64  // provided type set.
    65  func removeTypeDefinitions(set typeSet, d *ast.GenDecl) {
    66  	if d.Tok != token.TYPE {
    67  		return
    68  	}
    69  
    70  	i := 0
    71  	for _, gs := range d.Specs {
    72  		s := gs.(*ast.TypeSpec)
    73  		if _, ok := set[s.Name.Name]; !ok {
    74  			d.Specs[i] = gs
    75  			i++
    76  		}
    77  	}
    78  
    79  	d.Specs = d.Specs[:i]
    80  }
    81  
    82  // removeTypes removes from the AST the definition of all types and their
    83  // method sets that are contained in the provided type set.
    84  func removeTypes(set typeSet, f *ast.File) {
    85  	// Go through the top-level declarations.
    86  	i := 0
    87  	for _, decl := range f.Decls {
    88  		keep := true
    89  		switch d := decl.(type) {
    90  		case *ast.GenDecl:
    91  			countBefore := len(d.Specs)
    92  			removeTypeDefinitions(set, d)
    93  			keep = countBefore == 0 || len(d.Specs) > 0
    94  		case *ast.FuncDecl:
    95  			keep = !isMethodOf(set, d)
    96  		}
    97  
    98  		if keep {
    99  			f.Decls[i] = decl
   100  			i++
   101  		}
   102  	}
   103  
   104  	f.Decls = f.Decls[:i]
   105  }