stringscutprefix.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  	"strings"
    12  
    13  	"golang.org/x/tools/go/analysis"
    14  	"golang.org/x/tools/go/analysis/passes/inspect"
    15  	"golang.org/x/tools/go/types/typeutil"
    16  	"golang.org/x/tools/internal/analysis/analyzerutil"
    17  	typeindexanalyzer "golang.org/x/tools/internal/analysis/typeindex"
    18  	"golang.org/x/tools/internal/astutil"
    19  	"golang.org/x/tools/internal/moreiters"
    20  	"golang.org/x/tools/internal/refactor"
    21  	"golang.org/x/tools/internal/typesinternal"
    22  	"golang.org/x/tools/internal/typesinternal/typeindex"
    23  	"golang.org/x/tools/internal/versions"
    24  )
    25  
    26  var StringsCutPrefixAnalyzer = &analysis.Analyzer{
    27  	Name: "stringscutprefix",
    28  	Doc:  analyzerutil.MustExtractDoc(doc, "stringscutprefix"),
    29  	Requires: []*analysis.Analyzer{
    30  		inspect.Analyzer,
    31  		typeindexanalyzer.Analyzer,
    32  	},
    33  	Run: stringscutprefix,
    34  	URL: "https://pkg.go.dev/golang.org/x/tools/go/analysis/passes/modernize#stringscutprefix",
    35  }
    36  
    37  // stringscutprefix offers a fix to replace an if statement which
    38  // calls to the 2 patterns below with strings.CutPrefix or strings.CutSuffix.
    39  //
    40  // Patterns:
    41  //
    42  //  1. if strings.HasPrefix(s, pre) { use(strings.TrimPrefix(s, pre) }
    43  //     =>
    44  //     if after, ok := strings.CutPrefix(s, pre); ok { use(after) }
    45  //
    46  //  2. if after := strings.TrimPrefix(s, pre); after != s { use(after) }
    47  //     =>
    48  //     if after, ok := strings.CutPrefix(s, pre); ok { use(after) }
    49  //
    50  // Similar patterns apply for CutSuffix.
    51  //
    52  // The use must occur within the first statement of the block, and the offered fix
    53  // only replaces the first occurrence of strings.TrimPrefix/TrimSuffix.
    54  //
    55  // Variants:
    56  // - bytes.HasPrefix/HasSuffix usage as pattern 1.
    57  func stringscutprefix(pass *analysis.Pass) (any, error) {
    58  	var (
    59  		index = pass.ResultOf[typeindexanalyzer.Analyzer].(*typeindex.Index)
    60  		info  = pass.TypesInfo
    61  
    62  		stringsTrimPrefix = index.Object("strings", "TrimPrefix")
    63  		bytesTrimPrefix   = index.Object("bytes", "TrimPrefix")
    64  		stringsTrimSuffix = index.Object("strings", "TrimSuffix")
    65  		bytesTrimSuffix   = index.Object("bytes", "TrimSuffix")
    66  	)
    67  	if !index.Used(stringsTrimPrefix, bytesTrimPrefix, stringsTrimSuffix, bytesTrimSuffix) {
    68  		return nil, nil
    69  	}
    70  
    71  	for curFile := range filesUsingGoVersion(pass, versions.Go1_20) {
    72  		for curIfStmt := range curFile.Preorder((*ast.IfStmt)(nil)) {
    73  			ifStmt := curIfStmt.Node().(*ast.IfStmt)
    74  
    75  			// pattern1
    76  			if call, ok := ifStmt.Cond.(*ast.CallExpr); ok && ifStmt.Init == nil && len(ifStmt.Body.List) > 0 {
    77  
    78  				obj := typeutil.Callee(info, call)
    79  				if !typesinternal.IsFunctionNamed(obj, "strings", "HasPrefix", "HasSuffix") &&
    80  					!typesinternal.IsFunctionNamed(obj, "bytes", "HasPrefix", "HasSuffix") {
    81  					continue
    82  				}
    83  				isPrefix := strings.HasSuffix(obj.Name(), "Prefix")
    84  
    85  				// Replace the first occurrence of strings.TrimPrefix(s, pre) in the first statement only,
    86  				// but not later statements in case s or pre are modified by intervening logic (ditto Suffix).
    87  				firstStmt := curIfStmt.Child(ifStmt.Body).Child(ifStmt.Body.List[0])
    88  				for curCall := range firstStmt.Preorder((*ast.CallExpr)(nil)) {
    89  					call1 := curCall.Node().(*ast.CallExpr)
    90  					obj1 := typeutil.Callee(info, call1)
    91  					// bytesTrimPrefix or stringsTrimPrefix might be nil if the file doesn't import it,
    92  					// so we need to ensure the obj1 is not nil otherwise the call1 is not TrimPrefix and cause a panic (ditto Suffix).
    93  					if obj1 == nil ||
    94  						obj1 != stringsTrimPrefix && obj1 != bytesTrimPrefix &&
    95  							obj1 != stringsTrimSuffix && obj1 != bytesTrimSuffix {
    96  						continue
    97  					}
    98  
    99  					isPrefix1 := strings.HasSuffix(obj1.Name(), "Prefix")
   100  					var cutFuncName, varName, message, fixMessage string
   101  					if isPrefix && isPrefix1 {
   102  						cutFuncName = "CutPrefix"
   103  						varName = "after"
   104  						message = "HasPrefix + TrimPrefix can be simplified to CutPrefix"
   105  						fixMessage = "Replace HasPrefix/TrimPrefix with CutPrefix"
   106  					} else if !isPrefix && !isPrefix1 {
   107  						cutFuncName = "CutSuffix"
   108  						varName = "before"
   109  						message = "HasSuffix + TrimSuffix can be simplified to CutSuffix"
   110  						fixMessage = "Replace HasSuffix/TrimSuffix with CutSuffix"
   111  					} else {
   112  						continue
   113  					}
   114  
   115  					// Have: if strings.HasPrefix(s0, pre0) { ...strings.TrimPrefix(s, pre)... } (ditto Suffix)
   116  					var (
   117  						s0   = call.Args[0]
   118  						pre0 = call.Args[1]
   119  						s    = call1.Args[0]
   120  						pre  = call1.Args[1]
   121  					)
   122  
   123  					// check whether the obj1 uses the exact the same argument with strings.HasPrefix
   124  					// shadow variables won't be valid because we only access the first statement (ditto Suffix).
   125  					if astutil.EqualSyntax(s0, s) && astutil.EqualSyntax(pre0, pre) {
   126  						after := refactor.FreshName(info.Scopes[ifStmt], ifStmt.Pos(), varName)
   127  						prefix, importEdits := refactor.AddImport(
   128  							info,
   129  							curFile.Node().(*ast.File),
   130  							obj1.Pkg().Name(),
   131  							obj1.Pkg().Path(),
   132  							cutFuncName,
   133  							call.Pos(),
   134  						)
   135  						okVarName := refactor.FreshName(info.Scopes[ifStmt], ifStmt.Pos(), "ok")
   136  						pass.Report(analysis.Diagnostic{
   137  							// highlight at HasPrefix call (ditto Suffix).
   138  							Pos:     call.Pos(),
   139  							End:     call.End(),
   140  							Message: message,
   141  							SuggestedFixes: []analysis.SuggestedFix{{
   142  								Message: fixMessage,
   143  								// if              strings.HasPrefix(s, pre)     { use(strings.TrimPrefix(s, pre)) }
   144  								//    ------------ -----------------        -----      --------------------------
   145  								// if after, ok := strings.CutPrefix(s, pre); ok { use(after)                      }
   146  								// (ditto Suffix)
   147  								TextEdits: append(importEdits, []analysis.TextEdit{
   148  									{
   149  										Pos:     call.Fun.Pos(),
   150  										End:     call.Fun.Pos(),
   151  										NewText: fmt.Appendf(nil, "%s, %s :=", after, okVarName),
   152  									},
   153  									{
   154  										Pos:     call.Fun.Pos(),
   155  										End:     call.Fun.End(),
   156  										NewText: fmt.Appendf(nil, "%s%s", prefix, cutFuncName),
   157  									},
   158  									{
   159  										Pos:     call.End(),
   160  										End:     call.End(),
   161  										NewText: fmt.Appendf(nil, "; %s ", okVarName),
   162  									},
   163  									{
   164  										Pos:     call1.Pos(),
   165  										End:     call1.End(),
   166  										NewText: []byte(after),
   167  									},
   168  								}...),
   169  							}}},
   170  						)
   171  						break
   172  					}
   173  				}
   174  			}
   175  
   176  			// pattern2
   177  			if bin, ok := ifStmt.Cond.(*ast.BinaryExpr); ok &&
   178  				bin.Op == token.NEQ &&
   179  				ifStmt.Init != nil &&
   180  				isSimpleAssign(ifStmt.Init) {
   181  				assign := ifStmt.Init.(*ast.AssignStmt)
   182  				if call, ok := assign.Rhs[0].(*ast.CallExpr); ok && assign.Tok == token.DEFINE {
   183  					lhs := assign.Lhs[0]
   184  					obj := typeutil.Callee(info, call)
   185  
   186  					if obj == nil ||
   187  						obj != stringsTrimPrefix && obj != bytesTrimPrefix && obj != stringsTrimSuffix && obj != bytesTrimSuffix {
   188  						continue
   189  					}
   190  
   191  					isPrefix1 := strings.HasSuffix(obj.Name(), "Prefix")
   192  					var cutFuncName, message, fixMessage string
   193  					if isPrefix1 {
   194  						cutFuncName = "CutPrefix"
   195  						message = "TrimPrefix can be simplified to CutPrefix"
   196  						fixMessage = "Replace TrimPrefix with CutPrefix"
   197  					} else {
   198  						cutFuncName = "CutSuffix"
   199  						message = "TrimSuffix can be simplified to CutSuffix"
   200  						fixMessage = "Replace TrimSuffix with CutSuffix"
   201  					}
   202  
   203  					if astutil.EqualSyntax(lhs, bin.X) && astutil.EqualSyntax(call.Args[0], bin.Y) ||
   204  						(astutil.EqualSyntax(lhs, bin.Y) && astutil.EqualSyntax(call.Args[0], bin.X)) {
   205  						okVarName := freshName(info, index, info.Scopes[ifStmt], ifStmt.Pos(), curIfStmt, curIfStmt, token.NoPos, "ok")
   206  						// Have one of:
   207  						//   if rest := TrimPrefix(s, prefix); rest != s { (ditto Suffix)
   208  						//   if rest := TrimPrefix(s, prefix); s != rest { (ditto Suffix)
   209  
   210  						// We use AddImport not to add an import (since it exists already)
   211  						// but to compute the correct prefix in the dot-import case.
   212  						prefix, importEdits := refactor.AddImport(
   213  							info,
   214  							curFile.Node().(*ast.File),
   215  							obj.Pkg().Name(),
   216  							obj.Pkg().Path(),
   217  							cutFuncName,
   218  							call.Pos(),
   219  						)
   220  
   221  						// if x     := strings.TrimPrefix(s, pre); x != s ...
   222  						//     ----            ----------          ------
   223  						// if x, ok := strings.CutPrefix (s, pre); ok     ...
   224  						// (ditto Suffix)
   225  						edits := append(importEdits, []analysis.TextEdit{
   226  							{
   227  								Pos:     assign.Lhs[0].End(),
   228  								End:     assign.Lhs[0].End(),
   229  								NewText: fmt.Appendf(nil, ", %s", okVarName),
   230  							},
   231  							{
   232  								Pos:     call.Fun.Pos(),
   233  								End:     call.Fun.End(),
   234  								NewText: fmt.Appendf(nil, "%s%s", prefix, cutFuncName),
   235  							},
   236  							{
   237  								Pos:     ifStmt.Cond.Pos(),
   238  								End:     ifStmt.Cond.End(),
   239  								NewText: []byte(okVarName),
   240  							},
   241  						}...)
   242  
   243  						// Replace the "after" variable with "_" if is unused inside the if statement.
   244  						if id, ok := lhs.(*ast.Ident); ok {
   245  							if obj := info.ObjectOf(id); obj != nil && moreiters.Len(index.Uses(obj)) < 2 {
   246  								edits = append(edits, analysis.TextEdit{
   247  									Pos:     assign.Lhs[0].Pos(),
   248  									End:     assign.Lhs[0].End(),
   249  									NewText: []byte("_"),
   250  								})
   251  							}
   252  						}
   253  
   254  						pass.Report(analysis.Diagnostic{
   255  							// highlight from the init and the condition end.
   256  							Pos:     ifStmt.Init.Pos(),
   257  							End:     ifStmt.Cond.End(),
   258  							Message: message,
   259  							SuggestedFixes: []analysis.SuggestedFix{{
   260  								Message:   fixMessage,
   261  								TextEdits: edits,
   262  							}},
   263  						})
   264  					}
   265  				}
   266  			}
   267  		}
   268  	}
   269  	return nil, nil
   270  }
   271
View as plain text