simplify.go

     1  // Copyright 2010 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  package main
     6  
     7  import (
     8  	"go/ast"
     9  	"go/token"
    10  	"reflect"
    11  )
    12  
    13  type simplifier struct{}
    14  
    15  func (s simplifier) Visit(node ast.Node) ast.Visitor {
    16  	switch n := node.(type) {
    17  	case *ast.CompositeLit:
    18  		// array, slice, and map composite literals may be simplified
    19  		outer := n
    20  		var keyType, eltType ast.Expr
    21  		switch typ := outer.Type.(type) {
    22  		case *ast.ArrayType:
    23  			eltType = typ.Elt
    24  		case *ast.MapType:
    25  			keyType = typ.Key
    26  			eltType = typ.Value
    27  		}
    28  
    29  		if eltType != nil {
    30  			var ktyp reflect.Value
    31  			if keyType != nil {
    32  				ktyp = reflect.ValueOf(keyType)
    33  			}
    34  			typ := reflect.ValueOf(eltType)
    35  			for i, x := range outer.Elts {
    36  				px := &outer.Elts[i]
    37  				// look at value of indexed/named elements
    38  				if t, ok := x.(*ast.KeyValueExpr); ok {
    39  					if keyType != nil {
    40  						s.simplifyLiteral(ktyp, keyType, t.Key, &t.Key)
    41  					}
    42  					x = t.Value
    43  					px = &t.Value
    44  				}
    45  				s.simplifyLiteral(typ, eltType, x, px)
    46  			}
    47  			// node was simplified - stop walk (there are no subnodes to simplify)
    48  			return nil
    49  		}
    50  
    51  	case *ast.SliceExpr:
    52  		// a slice expression of the form: s[a:len(s)]
    53  		// can be simplified to: s[a:]
    54  		// if s is "simple enough" (for now we only accept identifiers)
    55  		//
    56  		// Note: This may not be correct because len may have been redeclared in
    57  		//       the same package. However, this is extremely unlikely and so far
    58  		//       (April 2022, after years of supporting this rewrite feature)
    59  		//       has never come up, so let's keep it working as is (see also #15153).
    60  		//
    61  		// Also note that this code used to use go/ast's object tracking,
    62  		// which was removed in exchange for go/parser.Mode.SkipObjectResolution.
    63  		// False positives are extremely unlikely as described above,
    64  		// and go/ast's object tracking is incomplete in any case.
    65  		if n.Max != nil {
    66  			// - 3-index slices always require the 2nd and 3rd index
    67  			break
    68  		}
    69  		if s, _ := n.X.(*ast.Ident); s != nil {
    70  			// the array/slice object is a single identifier
    71  			if call, _ := n.High.(*ast.CallExpr); call != nil && len(call.Args) == 1 && !call.Ellipsis.IsValid() {
    72  				// the high expression is a function call with a single argument
    73  				if fun, _ := call.Fun.(*ast.Ident); fun != nil && fun.Name == "len" {
    74  					// the function called is "len"
    75  					if arg, _ := call.Args[0].(*ast.Ident); arg != nil && arg.Name == s.Name {
    76  						// the len argument is the array/slice object
    77  						n.High = nil
    78  					}
    79  				}
    80  			}
    81  		}
    82  		// Note: We could also simplify slice expressions of the form s[0:b] to s[:b]
    83  		//       but we leave them as is since sometimes we want to be very explicit
    84  		//       about the lower bound.
    85  		// An example where the 0 helps:
    86  		//       x, y, z := b[0:2], b[2:4], b[4:6]
    87  		// An example where it does not:
    88  		//       x, y := b[:n], b[n:]
    89  
    90  	case *ast.RangeStmt:
    91  		// - a range of the form: for x, _ = range v {...}
    92  		// can be simplified to: for x = range v {...}
    93  		// - a range of the form: for _ = range v {...}
    94  		// can be simplified to: for range v {...}
    95  		if isBlank(n.Value) {
    96  			n.Value = nil
    97  		}
    98  		if isBlank(n.Key) && n.Value == nil {
    99  			n.Key = nil
   100  		}
   101  	}
   102  
   103  	return s
   104  }
   105  
   106  func (s simplifier) simplifyLiteral(typ reflect.Value, astType, x ast.Expr, px *ast.Expr) {
   107  	ast.Walk(s, x) // simplify x
   108  
   109  	// if the element is a composite literal and its literal type
   110  	// matches the outer literal's element type exactly, the inner
   111  	// literal type may be omitted
   112  	if inner, ok := x.(*ast.CompositeLit); ok {
   113  		if match(nil, typ, reflect.ValueOf(inner.Type)) {
   114  			inner.Type = nil
   115  		}
   116  	}
   117  	// if the outer literal's element type is a pointer type *T
   118  	// and the element is & of a composite literal of type T,
   119  	// the inner &T may be omitted.
   120  	if ptr, ok := astType.(*ast.StarExpr); ok {
   121  		if addr, ok := x.(*ast.UnaryExpr); ok && addr.Op == token.AND {
   122  			if inner, ok := addr.X.(*ast.CompositeLit); ok {
   123  				if match(nil, reflect.ValueOf(ptr.X), reflect.ValueOf(inner.Type)) {
   124  					inner.Type = nil // drop T
   125  					*px = inner      // drop &
   126  				}
   127  			}
   128  		}
   129  	}
   130  }
   131  
   132  func isBlank(x ast.Expr) bool {
   133  	ident, ok := x.(*ast.Ident)
   134  	return ok && ident.Name == "_"
   135  }
   136  
   137  func simplify(f *ast.File) {
   138  	// remove empty declarations such as "const ()", etc
   139  	removeEmptyDeclGroups(f)
   140  
   141  	var s simplifier
   142  	ast.Walk(s, f)
   143  }
   144  
   145  func removeEmptyDeclGroups(f *ast.File) {
   146  	i := 0
   147  	for _, d := range f.Decls {
   148  		if g, ok := d.(*ast.GenDecl); !ok || !isEmpty(f, g) {
   149  			f.Decls[i] = d
   150  			i++
   151  		}
   152  	}
   153  	f.Decls = f.Decls[:i]
   154  }
   155  
   156  func isEmpty(f *ast.File, g *ast.GenDecl) bool {
   157  	if g.Doc != nil || g.Specs != nil {
   158  		return false
   159  	}
   160  
   161  	for _, c := range f.Comments {
   162  		// if there is a comment in the declaration, it is not considered empty
   163  		if g.Pos() <= c.Pos() && c.End() <= g.End() {
   164  			return false
   165  		}
   166  	}
   167  
   168  	return true
   169  }
   170
View as plain text