maps.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  package modernize
     6  
     7  // This file defines modernizers that use the "maps" package.
     8  
     9  import (
    10  	"fmt"
    11  	"go/ast"
    12  	"go/token"
    13  	"go/types"
    14  
    15  	"golang.org/x/tools/go/analysis"
    16  	"golang.org/x/tools/go/analysis/passes/inspect"
    17  	"golang.org/x/tools/go/ast/inspector"
    18  	"golang.org/x/tools/internal/analysis/analyzerutil"
    19  	"golang.org/x/tools/internal/astutil"
    20  	"golang.org/x/tools/internal/refactor"
    21  	"golang.org/x/tools/internal/typeparams"
    22  	"golang.org/x/tools/internal/typesinternal"
    23  	"golang.org/x/tools/internal/versions"
    24  )
    25  
    26  var MapsLoopAnalyzer = &analysis.Analyzer{
    27  	Name:     "mapsloop",
    28  	Doc:      analyzerutil.MustExtractDoc(doc, "mapsloop"),
    29  	Requires: []*analysis.Analyzer{inspect.Analyzer},
    30  	Run:      mapsloop,
    31  	URL:      "https://pkg.go.dev/golang.org/x/tools/go/analysis/passes/modernize#mapsloop",
    32  }
    33  
    34  // The mapsloop pass offers to simplify a loop of map insertions:
    35  //
    36  //	for k, v := range x {
    37  //		m[k] = v
    38  //	}
    39  //
    40  // by a call to go1.23's maps package. There are four variants, the
    41  // product of two axes: whether the source x is a map or an iter.Seq2,
    42  // and whether the destination m is a newly created map:
    43  //
    44  //	maps.Copy(m, x)		(x is map)
    45  //	maps.Insert(m, x)       (x is iter.Seq2)
    46  //	m = maps.Clone(x)       (x is a non-nil map, m is a new map)
    47  //	m = maps.Collect(x)     (x is iter.Seq2, m is a new map)
    48  //
    49  // A map is newly created if the preceding statement has one of these
    50  // forms, where M is a map type:
    51  //
    52  //	m = make(M)
    53  //	m = M{}
    54  func mapsloop(pass *analysis.Pass) (any, error) {
    55  	// Skip the analyzer in packages where its
    56  	// fixes would create an import cycle.
    57  	if within(pass, "maps", "bytes", "runtime") {
    58  		return nil, nil
    59  	}
    60  
    61  	info := pass.TypesInfo
    62  
    63  	// check is called for each statement of this form:
    64  	//   for k, v := range x { m[k] = v }
    65  	check := func(file *ast.File, curRange inspector.Cursor, assign *ast.AssignStmt, m, x ast.Expr) {
    66  
    67  		// Is x a map or iter.Seq2?
    68  		tx := types.Unalias(info.TypeOf(x))
    69  		var xmap bool
    70  		switch typeparams.CoreType(tx).(type) {
    71  		case *types.Map:
    72  			xmap = true
    73  
    74  		case *types.Signature:
    75  			k, v, ok := assignableToIterSeq2(tx)
    76  			if !ok {
    77  				return // a named isomer of Seq2
    78  			}
    79  			xmap = false
    80  
    81  			// Record in tx the unnamed map[K]V type
    82  			// derived from the yield function.
    83  			// This is the type of maps.Collect(x).
    84  			tx = types.NewMap(k, v)
    85  
    86  		default:
    87  			return // e.g. slice, channel (or no core type!)
    88  		}
    89  
    90  		// Is the preceding statement of the form
    91  		//    m = make(M) or M{}
    92  		// and can we replace its RHS with slices.{Clone,Collect}?
    93  		//
    94  		// Beware: if x may be nil, we cannot use Clone as it preserves nilness.
    95  		var mrhs ast.Expr       // make(M) or M{}, or nil
    96  		var mAssign token.Token // token used to assign m
    97  		if curPrev, ok := curRange.PrevSibling(); ok {
    98  			if assign, ok := curPrev.Node().(*ast.AssignStmt); ok &&
    99  				len(assign.Lhs) == 1 &&
   100  				len(assign.Rhs) == 1 &&
   101  				astutil.EqualSyntax(assign.Lhs[0], m) {
   102  				mAssign = assign.Tok
   103  				// Have: m = rhs; for k, v := range x { m[k] = v }
   104  				var newMap bool
   105  				rhs := assign.Rhs[0]
   106  				switch rhs := ast.Unparen(rhs).(type) {
   107  				case *ast.CallExpr:
   108  					if id, ok := ast.Unparen(rhs.Fun).(*ast.Ident); ok &&
   109  						info.Uses[id] == builtinMake {
   110  						// Have: m = make(...)
   111  						newMap = true
   112  					}
   113  				case *ast.CompositeLit:
   114  					if len(rhs.Elts) == 0 {
   115  						// Have m = M{}
   116  						newMap = true
   117  					}
   118  				}
   119  
   120  				// Take care not to change type of m's RHS expression.
   121  				if newMap {
   122  					trhs := info.TypeOf(rhs)
   123  
   124  					// Inv: tx is the type of maps.F(x)
   125  					// - maps.Clone(x) has the same type as x.
   126  					// - maps.Collect(x) returns an unnamed map type.
   127  
   128  					if assign.Tok == token.DEFINE {
   129  						// DEFINE (:=): we must not
   130  						// change the type of RHS.
   131  						if types.Identical(tx, trhs) {
   132  							mrhs = rhs
   133  						}
   134  					} else {
   135  						// ASSIGN (=): the types of LHS
   136  						// and RHS may differ in namedness.
   137  						if types.AssignableTo(tx, trhs) {
   138  							mrhs = rhs
   139  						}
   140  					}
   141  
   142  					// Temporarily disable the transformation to the
   143  					// (nil-preserving) maps.Clone until we can prove
   144  					// that x is non-nil. This is rarely possible,
   145  					// and may require control flow analysis
   146  					// (e.g. a dominating "if len(x)" check).
   147  					// See #71844.
   148  					if xmap {
   149  						mrhs = nil
   150  					}
   151  				}
   152  			}
   153  		}
   154  
   155  		// Choose function.
   156  		var funcName string
   157  		if mrhs != nil {
   158  			funcName = cond(xmap, "Clone", "Collect")
   159  		} else {
   160  			funcName = cond(xmap, "Copy", "Insert")
   161  		}
   162  
   163  		// Report diagnostic, and suggest fix.
   164  		rng := curRange.Node()
   165  		prefix, importEdits := refactor.AddImport(info, file, "maps", "maps", funcName, rng.Pos())
   166  		var (
   167  			newText    []byte
   168  			start, end token.Pos
   169  		)
   170  		if mrhs != nil {
   171  			// Replace assignment and loop with expression.
   172  			//
   173  			//   m = make(...)
   174  			//   for k, v := range x { /* comments */ m[k] = v }
   175  			//
   176  			//   ->
   177  			//
   178  			//   /* comments */
   179  			//   m = maps.Collect(x)
   180  			curPrev, _ := curRange.PrevSibling()
   181  			start, end = curPrev.Node().Pos(), rng.End()
   182  			newText = fmt.Appendf(nil, "%s%s %s %s%s(%s)",
   183  				allComments(file, start, end),
   184  				astutil.Format(pass.Fset, m),
   185  				mAssign.String(),
   186  				prefix,
   187  				funcName,
   188  				astutil.Format(pass.Fset, x))
   189  		} else {
   190  			// Replace loop with call statement.
   191  			//
   192  			//   for k, v := range x { /* comments */ m[k] = v }
   193  			//
   194  			//   ->
   195  			//
   196  			//   /* comments */
   197  			//   maps.Copy(m, x)
   198  			start, end = rng.Pos(), rng.End()
   199  			newText = fmt.Appendf(nil, "%s%s%s(%s, %s)",
   200  				allComments(file, start, end),
   201  				prefix,
   202  				funcName,
   203  				astutil.Format(pass.Fset, m),
   204  				astutil.Format(pass.Fset, x))
   205  		}
   206  		pass.Report(analysis.Diagnostic{
   207  			Pos:     assign.Lhs[0].Pos(),
   208  			End:     assign.Lhs[0].End(),
   209  			Message: "Replace m[k]=v loop with maps." + funcName,
   210  			SuggestedFixes: []analysis.SuggestedFix{{
   211  				Message: "Replace m[k]=v loop with maps." + funcName,
   212  				TextEdits: append(importEdits, []analysis.TextEdit{{
   213  					Pos:     start,
   214  					End:     end,
   215  					NewText: newText,
   216  				}}...),
   217  			}},
   218  		})
   219  
   220  	}
   221  
   222  	// Find all range loops around m[k] = v.
   223  	for curFile := range filesUsingGoVersion(pass, versions.Go1_23) {
   224  		file := curFile.Node().(*ast.File)
   225  
   226  		for curRange := range curFile.Preorder((*ast.RangeStmt)(nil)) {
   227  			rng := curRange.Node().(*ast.RangeStmt)
   228  
   229  			if rng.Tok == token.DEFINE &&
   230  				rng.Key != nil &&
   231  				rng.Value != nil &&
   232  				isAssignBlock(rng.Body) {
   233  				// Have: for k, v := range x { lhs = rhs }
   234  
   235  				assign := rng.Body.List[0].(*ast.AssignStmt)
   236  
   237  				// usesKV reports whether e references vars k or v.
   238  				usesKV := func(e ast.Expr) bool {
   239  					k := info.Defs[rng.Key.(*ast.Ident)]
   240  					v := info.Defs[rng.Value.(*ast.Ident)]
   241  					for n := range ast.Preorder(e) {
   242  						if id, ok := n.(*ast.Ident); ok {
   243  							obj := info.Uses[id]
   244  							if obj != nil && // don't rely on k, v being non-nil
   245  								(obj == k || obj == v) {
   246  								return true
   247  							}
   248  						}
   249  					}
   250  					return false
   251  				}
   252  
   253  				if index, ok := assign.Lhs[0].(*ast.IndexExpr); ok &&
   254  					len(assign.Lhs) == 1 &&
   255  					astutil.EqualSyntax(rng.Key, index.Index) &&
   256  					astutil.EqualSyntax(rng.Value, assign.Rhs[0]) &&
   257  					!usesKV(index.X) { // reject (e.g.) f(k, v)[k] = v
   258  					if tmap, ok := typeparams.CoreType(info.TypeOf(index.X)).(*types.Map); ok &&
   259  						types.Identical(info.TypeOf(index), info.TypeOf(rng.Value)) && // m[k], v
   260  						types.Identical(tmap.Key(), info.TypeOf(rng.Key)) {
   261  
   262  						// Have: for k, v := range x { m[k] = v }
   263  						// where there is no implicit conversion
   264  						// of either key or value.
   265  						check(file, curRange, assign, index.X, rng.X)
   266  					}
   267  				}
   268  			}
   269  		}
   270  	}
   271  	return nil, nil
   272  }
   273  
   274  // assignableToIterSeq2 reports whether t is assignable to
   275  // iter.Seq[K, V] and returns K and V if so.
   276  func assignableToIterSeq2(t types.Type) (k, v types.Type, ok bool) {
   277  	// The only named type assignable to iter.Seq2 is iter.Seq2.
   278  	if is[*types.Named](t) {
   279  		if !typesinternal.IsTypeNamed(t, "iter", "Seq2") {
   280  			return
   281  		}
   282  		t = t.Underlying()
   283  	}
   284  
   285  	if t, ok := t.(*types.Signature); ok {
   286  		// func(yield func(K, V) bool)?
   287  		if t.Params().Len() == 1 && t.Results().Len() == 0 {
   288  			if yield, ok := t.Params().At(0).Type().(*types.Signature); ok { // sic, no Underlying/CoreType
   289  				if yield.Params().Len() == 2 &&
   290  					yield.Results().Len() == 1 &&
   291  					types.Identical(yield.Results().At(0).Type(), builtinBool.Type()) {
   292  					return yield.Params().At(0).Type(), yield.Params().At(1).Type(), true
   293  				}
   294  			}
   295  		}
   296  	}
   297  	return
   298  }
   299
View as plain text