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  		if curPrev, ok := curRange.PrevSibling(); ok {
    97  			if assign, ok := curPrev.Node().(*ast.AssignStmt); ok &&
    98  				len(assign.Lhs) == 1 &&
    99  				len(assign.Rhs) == 1 &&
   100  				astutil.EqualSyntax(assign.Lhs[0], m) {
   101  
   102  				// Have: m = rhs; for k, v := range x { m[k] = v }
   103  				var newMap bool
   104  				rhs := assign.Rhs[0]
   105  				switch rhs := ast.Unparen(rhs).(type) {
   106  				case *ast.CallExpr:
   107  					if id, ok := ast.Unparen(rhs.Fun).(*ast.Ident); ok &&
   108  						info.Uses[id] == builtinMake {
   109  						// Have: m = make(...)
   110  						newMap = true
   111  					}
   112  				case *ast.CompositeLit:
   113  					if len(rhs.Elts) == 0 {
   114  						// Have m = M{}
   115  						newMap = true
   116  					}
   117  				}
   118  
   119  				// Take care not to change type of m's RHS expression.
   120  				if newMap {
   121  					trhs := info.TypeOf(rhs)
   122  
   123  					// Inv: tx is the type of maps.F(x)
   124  					// - maps.Clone(x) has the same type as x.
   125  					// - maps.Collect(x) returns an unnamed map type.
   126  
   127  					if assign.Tok == token.DEFINE {
   128  						// DEFINE (:=): we must not
   129  						// change the type of RHS.
   130  						if types.Identical(tx, trhs) {
   131  							mrhs = rhs
   132  						}
   133  					} else {
   134  						// ASSIGN (=): the types of LHS
   135  						// and RHS may differ in namedness.
   136  						if types.AssignableTo(tx, trhs) {
   137  							mrhs = rhs
   138  						}
   139  					}
   140  
   141  					// Temporarily disable the transformation to the
   142  					// (nil-preserving) maps.Clone until we can prove
   143  					// that x is non-nil. This is rarely possible,
   144  					// and may require control flow analysis
   145  					// (e.g. a dominating "if len(x)" check).
   146  					// See #71844.
   147  					if xmap {
   148  						mrhs = nil
   149  					}
   150  				}
   151  			}
   152  		}
   153  
   154  		// Choose function.
   155  		var funcName string
   156  		if mrhs != nil {
   157  			funcName = cond(xmap, "Clone", "Collect")
   158  		} else {
   159  			funcName = cond(xmap, "Copy", "Insert")
   160  		}
   161  
   162  		// Report diagnostic, and suggest fix.
   163  		rng := curRange.Node()
   164  		prefix, importEdits := refactor.AddImport(info, file, "maps", "maps", funcName, rng.Pos())
   165  		var (
   166  			newText    []byte
   167  			start, end token.Pos
   168  		)
   169  		if mrhs != nil {
   170  			// Replace assignment and loop with expression.
   171  			//
   172  			//   m = make(...)
   173  			//   for k, v := range x { /* comments */ m[k] = v }
   174  			//
   175  			//   ->
   176  			//
   177  			//   /* comments */
   178  			//   m = maps.Copy(x)
   179  			curPrev, _ := curRange.PrevSibling()
   180  			start, end = curPrev.Node().Pos(), rng.End()
   181  			newText = fmt.Appendf(nil, "%s%s = %s%s(%s)",
   182  				allComments(file, start, end),
   183  				astutil.Format(pass.Fset, m),
   184  				prefix,
   185  				funcName,
   186  				astutil.Format(pass.Fset, x))
   187  		} else {
   188  			// Replace loop with call statement.
   189  			//
   190  			//   for k, v := range x { /* comments */ m[k] = v }
   191  			//
   192  			//   ->
   193  			//
   194  			//   /* comments */
   195  			//   maps.Copy(m, x)
   196  			start, end = rng.Pos(), rng.End()
   197  			newText = fmt.Appendf(nil, "%s%s%s(%s, %s)",
   198  				allComments(file, start, end),
   199  				prefix,
   200  				funcName,
   201  				astutil.Format(pass.Fset, m),
   202  				astutil.Format(pass.Fset, x))
   203  		}
   204  		pass.Report(analysis.Diagnostic{
   205  			Pos:     assign.Lhs[0].Pos(),
   206  			End:     assign.Lhs[0].End(),
   207  			Message: "Replace m[k]=v loop with maps." + funcName,
   208  			SuggestedFixes: []analysis.SuggestedFix{{
   209  				Message: "Replace m[k]=v loop with maps." + funcName,
   210  				TextEdits: append(importEdits, []analysis.TextEdit{{
   211  					Pos:     start,
   212  					End:     end,
   213  					NewText: newText,
   214  				}}...),
   215  			}},
   216  		})
   217  
   218  	}
   219  
   220  	// Find all range loops around m[k] = v.
   221  	for curFile := range filesUsingGoVersion(pass, versions.Go1_23) {
   222  		file := curFile.Node().(*ast.File)
   223  
   224  		for curRange := range curFile.Preorder((*ast.RangeStmt)(nil)) {
   225  			rng := curRange.Node().(*ast.RangeStmt)
   226  
   227  			if rng.Tok == token.DEFINE &&
   228  				rng.Key != nil &&
   229  				rng.Value != nil &&
   230  				isAssignBlock(rng.Body) {
   231  				// Have: for k, v := range x { lhs = rhs }
   232  
   233  				assign := rng.Body.List[0].(*ast.AssignStmt)
   234  
   235  				// usesKV reports whether e references vars k or v.
   236  				usesKV := func(e ast.Expr) bool {
   237  					k := info.Defs[rng.Key.(*ast.Ident)]
   238  					v := info.Defs[rng.Value.(*ast.Ident)]
   239  					for n := range ast.Preorder(e) {
   240  						if id, ok := n.(*ast.Ident); ok {
   241  							obj := info.Uses[id]
   242  							if obj != nil && // don't rely on k, v being non-nil
   243  								(obj == k || obj == v) {
   244  								return true
   245  							}
   246  						}
   247  					}
   248  					return false
   249  				}
   250  
   251  				if index, ok := assign.Lhs[0].(*ast.IndexExpr); ok &&
   252  					len(assign.Lhs) == 1 &&
   253  					astutil.EqualSyntax(rng.Key, index.Index) &&
   254  					astutil.EqualSyntax(rng.Value, assign.Rhs[0]) &&
   255  					!usesKV(index.X) { // reject (e.g.) f(k, v)[k] = v
   256  					if tmap, ok := typeparams.CoreType(info.TypeOf(index.X)).(*types.Map); ok &&
   257  						types.Identical(info.TypeOf(index), info.TypeOf(rng.Value)) && // m[k], v
   258  						types.Identical(tmap.Key(), info.TypeOf(rng.Key)) {
   259  
   260  						// Have: for k, v := range x { m[k] = v }
   261  						// where there is no implicit conversion
   262  						// of either key or value.
   263  						check(file, curRange, assign, index.X, rng.X)
   264  					}
   265  				}
   266  			}
   267  		}
   268  	}
   269  	return nil, nil
   270  }
   271  
   272  // assignableToIterSeq2 reports whether t is assignable to
   273  // iter.Seq[K, V] and returns K and V if so.
   274  func assignableToIterSeq2(t types.Type) (k, v types.Type, ok bool) {
   275  	// The only named type assignable to iter.Seq2 is iter.Seq2.
   276  	if is[*types.Named](t) {
   277  		if !typesinternal.IsTypeNamed(t, "iter", "Seq2") {
   278  			return
   279  		}
   280  		t = t.Underlying()
   281  	}
   282  
   283  	if t, ok := t.(*types.Signature); ok {
   284  		// func(yield func(K, V) bool)?
   285  		if t.Params().Len() == 1 && t.Results().Len() == 0 {
   286  			if yield, ok := t.Params().At(0).Type().(*types.Signature); ok { // sic, no Underlying/CoreType
   287  				if yield.Params().Len() == 2 &&
   288  					yield.Results().Len() == 1 &&
   289  					types.Identical(yield.Results().At(0).Type(), builtinBool.Type()) {
   290  					return yield.Params().At(0).Type(), yield.Params().At(1).Type(), true
   291  				}
   292  			}
   293  		}
   294  	}
   295  	return
   296  }
   297
View as plain text