rangeint.go

     1  // Copyright 2025 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  import (
     8  	"fmt"
     9  	"go/ast"
    10  	"go/token"
    11  	"go/types"
    12  
    13  	"golang.org/x/tools/go/analysis"
    14  	"golang.org/x/tools/go/analysis/passes/inspect"
    15  	"golang.org/x/tools/go/ast/edge"
    16  	"golang.org/x/tools/go/ast/inspector"
    17  	"golang.org/x/tools/go/types/typeutil"
    18  	"golang.org/x/tools/internal/analysis/analyzerutil"
    19  	typeindexanalyzer "golang.org/x/tools/internal/analysis/typeindex"
    20  	"golang.org/x/tools/internal/astutil"
    21  	"golang.org/x/tools/internal/moreiters"
    22  	"golang.org/x/tools/internal/typesinternal"
    23  	"golang.org/x/tools/internal/typesinternal/typeindex"
    24  	"golang.org/x/tools/internal/versions"
    25  )
    26  
    27  var RangeIntAnalyzer = &analysis.Analyzer{
    28  	Name: "rangeint",
    29  	Doc:  analyzerutil.MustExtractDoc(doc, "rangeint"),
    30  	Requires: []*analysis.Analyzer{
    31  		inspect.Analyzer,
    32  		typeindexanalyzer.Analyzer,
    33  	},
    34  	Run: rangeint,
    35  	URL: "https://pkg.go.dev/golang.org/x/tools/go/analysis/passes/modernize#rangeint",
    36  }
    37  
    38  // rangeint offers a fix to replace a 3-clause 'for' loop:
    39  //
    40  //	for i := 0; i < limit; i++ {}
    41  //
    42  // by a range loop with an integer operand:
    43  //
    44  //	for i := range limit {}
    45  //
    46  // Variants:
    47  //   - The ':=' may be replaced by '='.
    48  //   - The fix may remove "i :=" if it would become unused.
    49  //
    50  // Restrictions:
    51  //   - The variable i must not be assigned or address-taken within the
    52  //     loop, because a "for range int" loop does not respect assignments
    53  //     to the loop index.
    54  //   - The limit must not be b.N, to avoid redundancy with bloop's fixes.
    55  //
    56  // Caveats:
    57  //
    58  // The fix causes the limit expression to be evaluated exactly once,
    59  // instead of once per iteration. So, to avoid changing the
    60  // cardinality of side effects, the limit expression must not involve
    61  // function calls (e.g. seq.Len()) or channel receives. Moreover, the
    62  // value of the limit expression must be loop invariant, which in
    63  // practice means it must take one of the following forms:
    64  //
    65  //   - a local variable that is assigned only once and not address-taken;
    66  //   - a constant; or
    67  //   - len(s), where s has the above properties.
    68  func rangeint(pass *analysis.Pass) (any, error) {
    69  	var (
    70  		info      = pass.TypesInfo
    71  		typeindex = pass.ResultOf[typeindexanalyzer.Analyzer].(*typeindex.Index)
    72  	)
    73  
    74  	for curFile := range filesUsingGoVersion(pass, versions.Go1_22) {
    75  	nextLoop:
    76  		for curLoop := range curFile.Preorder((*ast.ForStmt)(nil)) {
    77  			loop := curLoop.Node().(*ast.ForStmt)
    78  			if init, ok := loop.Init.(*ast.AssignStmt); ok &&
    79  				isSimpleAssign(init) &&
    80  				is[*ast.Ident](init.Lhs[0]) &&
    81  				isZeroIntConst(info, init.Rhs[0]) {
    82  				// Have: for i = 0; ... (or i := 0)
    83  				index := init.Lhs[0].(*ast.Ident)
    84  
    85  				if compare, ok := loop.Cond.(*ast.BinaryExpr); ok &&
    86  					compare.Op == token.LSS &&
    87  					astutil.EqualSyntax(compare.X, init.Lhs[0]) {
    88  					// Have: for i = 0; i < limit; ... {}
    89  
    90  					limit := compare.Y
    91  
    92  					// If limit is "len(slice)", simplify it to "slice".
    93  					//
    94  					// (Don't replace "for i := 0; i < len(map); i++"
    95  					// with "for range m" because it's too hard to prove
    96  					// that len(m) is loop-invariant).
    97  					if call, ok := limit.(*ast.CallExpr); ok &&
    98  						typeutil.Callee(info, call) == builtinLen &&
    99  						is[*types.Slice](info.TypeOf(call.Args[0]).Underlying()) {
   100  						limit = call.Args[0]
   101  					}
   102  
   103  					// Check the form of limit: must be a constant,
   104  					// or a local var that is not assigned or address-taken.
   105  					limitOK := false
   106  					if info.Types[limit].Value != nil {
   107  						limitOK = true // constant
   108  					} else if id, ok := limit.(*ast.Ident); ok {
   109  						if v, ok := info.Uses[id].(*types.Var); ok &&
   110  							!(v.Exported() && typesinternal.IsPackageLevel(v)) {
   111  							// limit is a local or unexported global var.
   112  							// (An exported global may have uses we can't see.)
   113  							for cur := range typeindex.Uses(v) {
   114  								if isScalarLvalue(info, cur) {
   115  									// Limit var is assigned or address-taken.
   116  									continue nextLoop
   117  								}
   118  							}
   119  							limitOK = true
   120  						}
   121  					}
   122  					if !limitOK {
   123  						continue nextLoop
   124  					}
   125  
   126  					validIncrement := false
   127  					if inc, ok := loop.Post.(*ast.IncDecStmt); ok &&
   128  						inc.Tok == token.INC &&
   129  						astutil.EqualSyntax(compare.X, inc.X) {
   130  						// Have: i++
   131  						validIncrement = true
   132  					} else if assign, ok := loop.Post.(*ast.AssignStmt); ok &&
   133  						assign.Tok == token.ADD_ASSIGN &&
   134  						len(assign.Rhs) == 1 && isIntLiteral(info, assign.Rhs[0], 1) &&
   135  						len(assign.Lhs) == 1 && astutil.EqualSyntax(compare.X, assign.Lhs[0]) {
   136  						// Have: i += 1
   137  						validIncrement = true
   138  					}
   139  
   140  					if validIncrement {
   141  						// Have: for i = 0; i < limit; i++ {}
   142  
   143  						// Find references to i within the loop body.
   144  						v := info.ObjectOf(index).(*types.Var)
   145  						// TODO(adonovan): use go1.25 v.Kind() == types.PackageVar
   146  						if typesinternal.IsPackageLevel(v) {
   147  							continue nextLoop
   148  						}
   149  
   150  						// If v is a named result, it is implicitly
   151  						// used after the loop (go.dev/issue/76880).
   152  						// TODO(adonovan): use go1.25 v.Kind() == types.ResultVar.
   153  						if moreiters.Contains(enclosingSignature(curLoop, info).Results().Variables(), v) {
   154  							continue nextLoop
   155  						}
   156  
   157  						used := false
   158  						for curId := range curLoop.Child(loop.Body).Preorder((*ast.Ident)(nil)) {
   159  							id := curId.Node().(*ast.Ident)
   160  							if info.Uses[id] == v {
   161  								used = true
   162  
   163  								// Reject if any is an l-value (assigned or address-taken):
   164  								// a "for range int" loop does not respect assignments to
   165  								// the loop variable.
   166  								if isScalarLvalue(info, curId) {
   167  									continue nextLoop
   168  								}
   169  							}
   170  						}
   171  
   172  						// If i is no longer used, delete "i := ".
   173  						var edits []analysis.TextEdit
   174  						if !used && init.Tok == token.DEFINE {
   175  							edits = append(edits, analysis.TextEdit{
   176  								Pos: index.Pos(),
   177  								End: init.Rhs[0].Pos(),
   178  							})
   179  						}
   180  
   181  						// If i is used after the loop,
   182  						// don't offer a fix, as a range loop
   183  						// leaves i with a different final value (limit-1).
   184  						if init.Tok == token.ASSIGN {
   185  							// Find the nearest ancestor that is not a label.
   186  							// Otherwise, checking for i usage outside of a for
   187  							// loop might not function properly further below.
   188  							// This is because the i usage might be a child of
   189  							// the loop's parent's parent, for example:
   190  							//     var i int
   191  							// Loop:
   192  							//     for i = 0; i < 10; i++ { break loop }
   193  							//     // i is in the sibling of the label, not the loop
   194  							//     fmt.Println(i)
   195  							//
   196  							ancestor := curLoop.Parent()
   197  							for is[*ast.LabeledStmt](ancestor.Node()) {
   198  								ancestor = ancestor.Parent()
   199  							}
   200  							for curId := range ancestor.Preorder((*ast.Ident)(nil)) {
   201  								id := curId.Node().(*ast.Ident)
   202  								if info.Uses[id] == v {
   203  									// Is i used after loop?
   204  									if id.Pos() > loop.End() {
   205  										continue nextLoop
   206  									}
   207  									// Is i used within a defer statement
   208  									// that is within the scope of i?
   209  									//     var i int
   210  									//     defer func() { print(i)}
   211  									//     for i = ... { ... }
   212  									for curDefer := range curId.Enclosing((*ast.DeferStmt)(nil)) {
   213  										if curDefer.Node().Pos() > v.Pos() {
   214  											continue nextLoop
   215  										}
   216  									}
   217  								}
   218  							}
   219  						}
   220  
   221  						// If limit is len(slice),
   222  						// simplify "range len(slice)" to "range slice".
   223  						if call, ok := limit.(*ast.CallExpr); ok &&
   224  							typeutil.Callee(info, call) == builtinLen &&
   225  							is[*types.Slice](info.TypeOf(call.Args[0]).Underlying()) {
   226  							limit = call.Args[0]
   227  						}
   228  
   229  						// If the limit is a untyped constant of non-integer type,
   230  						// such as "const limit = 1e3", its effective type may
   231  						// differ between the two forms.
   232  						// In a for loop, it must be comparable with int i,
   233  						//    for i := 0; i < limit; i++
   234  						// but in a range loop it would become a float,
   235  						//    for i := range limit {}
   236  						// which is a type error. We need to convert it to int
   237  						// in this case.
   238  						//
   239  						// Unfortunately go/types discards the untyped type
   240  						// (but see Untyped in golang/go#70638) so we must
   241  						// re-type check the expression to detect this case.
   242  						var beforeLimit, afterLimit string
   243  						if v := info.Types[limit].Value; v != nil {
   244  							tVar := info.TypeOf(init.Rhs[0])
   245  							file := curFile.Node().(*ast.File)
   246  							// TODO(mkalil): use a types.Qualifier that respects the existing
   247  							// imports of this file that are visible (not shadowed) at the current position.
   248  							qual := typesinternal.FileQualifier(file, pass.Pkg)
   249  							beforeLimit, afterLimit = fmt.Sprintf("%s(", types.TypeString(tVar, qual)), ")"
   250  							info2 := &types.Info{Types: make(map[ast.Expr]types.TypeAndValue)}
   251  							if types.CheckExpr(pass.Fset, pass.Pkg, limit.Pos(), limit, info2) == nil {
   252  								tLimit := types.Default(info2.TypeOf(limit))
   253  								if types.AssignableTo(tLimit, tVar) {
   254  									beforeLimit, afterLimit = "", ""
   255  								}
   256  							}
   257  						}
   258  
   259  						pass.Report(analysis.Diagnostic{
   260  							Pos:     init.Pos(),
   261  							End:     loop.Post.End(),
   262  							Message: "for loop can be modernized using range over int",
   263  							SuggestedFixes: []analysis.SuggestedFix{{
   264  								Message: fmt.Sprintf("Replace for loop with range %s",
   265  									astutil.Format(pass.Fset, limit)),
   266  								TextEdits: append(edits, []analysis.TextEdit{
   267  									// for i := 0; i < limit; i++ {}
   268  									//     -----              ---
   269  									//          -------
   270  									// for i := range  limit      {}
   271  
   272  									// Delete init.
   273  									{
   274  										Pos:     init.Rhs[0].Pos(),
   275  										End:     limit.Pos(),
   276  										NewText: []byte("range "),
   277  									},
   278  									// Add "int(" before limit, if needed.
   279  									{
   280  										Pos:     limit.Pos(),
   281  										End:     limit.Pos(),
   282  										NewText: []byte(beforeLimit),
   283  									},
   284  									// Delete inc.
   285  									{
   286  										Pos: limit.End(),
   287  										End: loop.Post.End(),
   288  									},
   289  									// Add ")" after limit, if needed.
   290  									{
   291  										Pos:     limit.End(),
   292  										End:     limit.End(),
   293  										NewText: []byte(afterLimit),
   294  									},
   295  								}...),
   296  							}},
   297  						})
   298  					}
   299  				}
   300  			}
   301  		}
   302  	}
   303  	return nil, nil
   304  }
   305  
   306  // isScalarLvalue reports whether the specified identifier is
   307  // address-taken or appears on the left side of an assignment.
   308  //
   309  // This function is valid only for scalars (x = ...),
   310  // not for aggregates (x.a[i] = ...)
   311  func isScalarLvalue(info *types.Info, curId inspector.Cursor) bool {
   312  	// Unfortunately we can't simply use info.Types[e].Assignable()
   313  	// as it is always true for a variable even when that variable is
   314  	// used only as an r-value. So we must inspect enclosing syntax.
   315  
   316  	cur := curId
   317  
   318  	// Strip enclosing parens.
   319  	ek, _ := cur.ParentEdge()
   320  	for ek == edge.ParenExpr_X {
   321  		cur = cur.Parent()
   322  		ek, _ = cur.ParentEdge()
   323  	}
   324  
   325  	switch ek {
   326  	case edge.AssignStmt_Lhs:
   327  		assign := cur.Parent().Node().(*ast.AssignStmt)
   328  		if assign.Tok != token.DEFINE {
   329  			return true // i = j or i += j
   330  		}
   331  		id := curId.Node().(*ast.Ident)
   332  		if v, ok := info.Defs[id]; ok && v.Pos() != id.Pos() {
   333  			return true // reassignment of i (i, j := 1, 2)
   334  		}
   335  	case edge.RangeStmt_Key:
   336  		rng := cur.Parent().Node().(*ast.RangeStmt)
   337  		if rng.Tok == token.ASSIGN {
   338  			return true // "for k, v = range x" is like an AssignStmt to k, v
   339  		}
   340  	case edge.IncDecStmt_X:
   341  		return true // i++, i--
   342  	case edge.UnaryExpr_X:
   343  		if cur.Parent().Node().(*ast.UnaryExpr).Op == token.AND {
   344  			return true // &i
   345  		}
   346  	}
   347  	return false
   348  }
   349  
   350  // enclosingSignature returns the signature of the innermost
   351  // function enclosing the syntax node denoted by cur
   352  // or nil if the node is not within a function.
   353  //
   354  // TODO(adonovan): factor with gopls/internal/util/typesutil.EnclosingSignature.
   355  func enclosingSignature(cur inspector.Cursor, info *types.Info) *types.Signature {
   356  	if c, ok := enclosingFunc(cur); ok {
   357  		switch n := c.Node().(type) {
   358  		case *ast.FuncDecl:
   359  			if f, ok := info.Defs[n.Name]; ok {
   360  				return f.Type().(*types.Signature)
   361  			}
   362  		case *ast.FuncLit:
   363  			if f, ok := info.Types[n]; ok {
   364  				return f.Type.(*types.Signature)
   365  			}
   366  		}
   367  	}
   368  	return nil
   369  }
   370
View as plain text