Repository: mgechev/revive
Pull request #1160 (2024-12-04 06:43:32)

refactor(rule/context-as-argument): replace AST walker by iteration over declarations
Simplifies context-as-argument rule by replacing AST walker by an iteration over global declarations



Patch:
diff --git a/rule/argument_limit.go b/rule/argument_limit.go
index ce831e5..b44e1fd 100644
--- a/rule/argument_limit.go
+++ b/rule/argument_limit.go
@@ -35,50 +35,33 @@ func (r *ArgumentsLimitRule) Apply(file *lint.File, arguments lint.Arguments) []
 	r.configureOnce.Do(func() { r.configure(arguments) })
 
 	var failures []lint.Failure
-	onFailure := func(failure lint.Failure) {
-		failures = append(failures, failure)
-	}
-
-	walker := lintArgsNum{
-		max:       r.max,
-		onFailure: onFailure,
-	}
-
-	ast.Walk(walker, file.AST)
-
-	return failures
-}
-
-// Name returns the rule name.
-func (*ArgumentsLimitRule) Name() string {
-	return "argument-limit"
-}
 
-type lintArgsNum struct {
-	max       int
-	onFailure func(lint.Failure)
-}
+	for _, decl := range file.AST.Decls {
+		funcDecl, ok := decl.(*ast.FuncDecl)
+		if !ok {
+			continue
+		}
 
-func (w lintArgsNum) Visit(n ast.Node) ast.Visitor {
-	node, ok := n.(*ast.FuncDecl)
-	if !ok {
-		return w
-	}
+		numParams := 0
+		for _, l := range funcDecl.Type.Params.List {
+			numParams += len(l.Names)
+		}
 
-	num := 0
-	for _, l := range node.Type.Params.List {
-		for range l.Names {
-			num++
+		if numParams <= r.max {
+			continue
 		}
-	}
 
-	if num > w.max {
-		w.onFailure(lint.Failure{
+		failures = append(failures, lint.Failure{
 			Confidence: 1,
-			Failure:    fmt.Sprintf("maximum number of arguments per function exceeded; max %d but got %d", w.max, num),
-			Node:       node.Type,
+			Failure:    fmt.Sprintf("maximum number of arguments per function exceeded; max %d but got %d", r.max, numParams),
+			Node:       funcDecl.Type,
 		})
 	}
 
-	return nil // skip visiting the body of the function
+	return failures
+}
+
+// Name returns the rule name.
+func (*ArgumentsLimitRule) Name() string {
+	return "argument-limit"
 }
diff --git a/rule/confusing_results.go b/rule/confusing_results.go
index 1b79ada..c1804d1 100644
--- a/rule/confusing_results.go
+++ b/rule/confusing_results.go
@@ -13,54 +13,43 @@ type ConfusingResultsRule struct{}
 func (*ConfusingResultsRule) Apply(file *lint.File, _ lint.Arguments) []lint.Failure {
 	var failures []lint.Failure
 
-	fileAst := file.AST
-	walker := lintConfusingResults{
-		onFailure: func(failure lint.Failure) {
-			failures = append(failures, failure)
-		},
-	}
+	for _, decl := range file.AST.Decls {
+		funcDecl, ok := decl.(*ast.FuncDecl)
 
-	ast.Walk(walker, fileAst)
+		isFunctionWithMoreThanOneResult := ok && funcDecl.Type.Results != nil && len(funcDecl.Type.Results.List) > 1
+		if !isFunctionWithMoreThanOneResult {
+			continue
+		}
 
-	return failures
-}
+		resultsAreNamed := len(funcDecl.Type.Results.List[0].Names) > 0
+		if resultsAreNamed {
+			continue
+		}
 
-// Name returns the rule name.
-func (*ConfusingResultsRule) Name() string {
-	return "confusing-results"
-}
+		lastType := ""
+		for _, result := range funcDecl.Type.Results.List {
 
-type lintConfusingResults struct {
-	onFailure func(lint.Failure)
-}
+			resultTypeName := gofmt(result.Type)
 
-func (w lintConfusingResults) Visit(n ast.Node) ast.Visitor {
-	fn, ok := n.(*ast.FuncDecl)
-	if !ok || fn.Type.Results == nil || len(fn.Type.Results.List) < 2 {
-		return w
-	}
-	lastType := ""
-	for _, result := range fn.Type.Results.List {
-		if len(result.Names) > 0 {
-			return w
-		}
+			if resultTypeName == lastType {
+				failures = append(failures, lint.Failure{
+					Node:       result,
+					Confidence: 1,
+					Category:   "naming",
+					Failure:    "unnamed results of the same type may be confusing, consider using named results",
+				})
 
-		t, ok := result.Type.(*ast.Ident)
-		if !ok {
-			return w
-		}
+				break
+			}
 
-		if t.Name == lastType {
-			w.onFailure(lint.Failure{
-				Node:       n,
-				Confidence: 1,
-				Category:   "naming",
-				Failure:    "unnamed results of the same type may be confusing, consider using named results",
-			})
-			break
+			lastType = resultTypeName
 		}
-		lastType = t.Name
 	}
 
-	return w
+	return failures
+}
+
+// Name returns the rule name.
+func (*ConfusingResultsRule) Name() string {
+	return "confusing-results"
 }
diff --git a/rule/context_as_argument.go b/rule/context_as_argument.go
index 06dc0f2..f2f9aee 100644
--- a/rule/context_as_argument.go
+++ b/rule/context_as_argument.go
@@ -11,7 +11,7 @@ import (
 
 // ContextAsArgumentRule suggests that `context.Context` should be the first argument of a function.
 type ContextAsArgumentRule struct {
-	allowTypesLUT map[string]struct{}
+	allowTypes map[string]struct{}
 
 	configureOnce sync.Once
 }
@@ -21,20 +21,37 @@ func (r *ContextAsArgumentRule) Apply(file *lint.File, args lint.Arguments) []li
 	r.configureOnce.Do(func() { r.configure(args) })
 
 	var failures []lint.Failure
-	walker := lintContextArguments{
-		allowTypesLUT: r.allowTypesLUT,
-		onFailure: func(failure lint.Failure) {
-			failures = append(failures, failure)
-		},
-	}
+	for _, decl := range file.AST.Decls {
+		fn, ok := decl.(*ast.FuncDecl)
+		if !ok || len(fn.Type.Params.List) <= 1 {
+			continue // not a function or a function with less than 2 parameters
+		}
 
-	ast.Walk(walker, file.AST)
+		fnArgs := fn.Type.Params.List
+
+		// A context.Context should be the first parameter of a function.
+		// Flag any that show up after the first.
+		isCtxStillAllowed := true
+		for _, arg := range fnArgs {
+			argIsCtx := isPkgDot(arg.Type, "context", "Context")
+			if argIsCtx && !isCtxStillAllowed {
+				failures = append(failures, lint.Failure{
+					Node:       arg,
+					Category:   "arg-order",
+					Failure:    "context.Context should be the first parameter of a function",
+					Confidence: 0.9,
+				})
+
+				break // only flag one
+			}
 
-	return failures
-}
+			typeName := gofmt(arg.Type)
+			// a parameter of type context.Context is still allowed if the current arg type is in the allow types LookUpTable
+			_, isCtxStillAllowed = r.allowTypes[typeName]
+		}
+	}
 
-func (r *ContextAsArgumentRule) configure(arguments lint.Arguments) {
-	r.allowTypesLUT = getAllowTypesFromArguments(arguments)
+	return failures
 }
 
 // Name returns the rule name.
@@ -42,43 +59,11 @@ func (*ContextAsArgumentRule) Name() string {
 	return "context-as-argument"
 }
 
-type lintContextArguments struct {
-	allowTypesLUT map[string]struct{}
-	onFailure     func(lint.Failure)
-}
-
-func (w lintContextArguments) Visit(n ast.Node) ast.Visitor {
-	fn, ok := n.(*ast.FuncDecl)
-	if !ok || len(fn.Type.Params.List) <= 1 {
-		return w
-	}
-
-	fnArgs := fn.Type.Params.List
-
-	// A context.Context should be the first parameter of a function.
-	// Flag any that show up after the first.
-	isCtxStillAllowed := true
-	for _, arg := range fnArgs {
-		argIsCtx := isPkgDot(arg.Type, "context", "Context")
-		if argIsCtx && !isCtxStillAllowed {
-			w.onFailure(lint.Failure{
-				Node:       arg,
-				Category:   "arg-order",
-				Failure:    "context.Context should be the first parameter of a function",
-				Confidence: 0.9,
-			})
-			break // only flag one
-		}
-
-		typeName := gofmt(arg.Type)
-		// a parameter of type context.Context is still allowed if the current arg type is in the LUT
-		_, isCtxStillAllowed = w.allowTypesLUT[typeName]
-	}
-
-	return nil // avoid visiting the function body
+func (r *ContextAsArgumentRule) configure(arguments lint.Arguments) {
+	r.allowTypes = r.getAllowTypesFromArguments(arguments)
 }
 
-func getAllowTypesFromArguments(args lint.Arguments) map[string]struct{} {
+func (r *ContextAsArgumentRule) getAllowTypesFromArguments(args lint.Arguments) map[string]struct{} {
 	allowTypesBefore := []string{}
 	if len(args) >= 1 {
 		argKV, ok := args[0].(map[string]any)
diff --git a/rule/cyclomatic.go b/rule/cyclomatic.go
index 60f8441..7f42a1c 100644
--- a/rule/cyclomatic.go
+++ b/rule/cyclomatic.go
@@ -38,53 +38,30 @@ func (r *CyclomaticRule) Apply(file *lint.File, arguments lint.Arguments) []lint
 	r.configureOnce.Do(func() { r.configure(arguments) })
 
 	var failures []lint.Failure
-	fileAst := file.AST
-
-	walker := lintCyclomatic{
-		file:       file,
-		complexity: r.maxComplexity,
-		onFailure: func(failure lint.Failure) {
-			failures = append(failures, failure)
-		},
-	}
-
-	ast.Walk(walker, fileAst)
-
-	return failures
-}
-
-// Name returns the rule name.
-func (*CyclomaticRule) Name() string {
-	return "cyclomatic"
-}
-
-type lintCyclomatic struct {
-	file       *lint.File
-	complexity int
-	onFailure  func(lint.Failure)
-}
-
-func (w lintCyclomatic) Visit(_ ast.Node) ast.Visitor {
-	f := w.file
-	for _, decl := range f.AST.Decls {
+	for _, decl := range file.AST.Decls {
 		fn, ok := decl.(*ast.FuncDecl)
 		if !ok {
 			continue
 		}
 
 		c := complexity(fn)
-		if c > w.complexity {
-			w.onFailure(lint.Failure{
+		if c > r.maxComplexity {
+			failures = append(failures, lint.Failure{
 				Confidence: 1,
 				Category:   "maintenance",
 				Failure: fmt.Sprintf("function %s has cyclomatic complexity %d (> max enabled %d)",
-					funcName(fn), c, w.complexity),
+					funcName(fn), c, r.maxComplexity),
 				Node: fn,
 			})
 		}
 	}
 
-	return nil
+	return failures
+}
+
+// Name returns the rule name.
+func (*CyclomaticRule) Name() string {
+	return "cyclomatic"
 }
 
 // funcName returns the name representation of a function or method:
diff --git a/rule/error_return.go b/rule/error_return.go
index f26d807..c5d3f46 100644
--- a/rule/error_return.go
+++ b/rule/error_return.go
@@ -13,16 +13,34 @@ type ErrorReturnRule struct{}
 func (*ErrorReturnRule) Apply(file *lint.File, _ lint.Arguments) []lint.Failure {
 	var failures []lint.Failure
 
-	fileAst := file.AST
-	walker := lintErrorReturn{
-		file:    file,
-		fileAst: fileAst,
-		onFailure: func(failure lint.Failure) {
-			failures = append(failures, failure)
-		},
-	}
+	for _, decl := range file.AST.Decls {
+		funcDecl, ok := decl.(*ast.FuncDecl)
+		isFunctionWithMoreThanOneResult := ok && funcDecl.Type.Results != nil && len(funcDecl.Type.Results.List) > 1
+		if !isFunctionWithMoreThanOneResult {
+			continue
+		}
+
+		funcResults := funcDecl.Type.Results.List
+		isLastResultError := isIdent(funcResults[len(funcResults)-1].Type, "error")
+		if isLastResultError {
+			continue
+		}
 
-	ast.Walk(walker, fileAst)
+		// An error return parameter should be the last parameter.
+		// Flag any error parameters found before the last.
+		for _, r := range funcResults[:len(funcResults)-1] {
+			if isIdent(r.Type, "error") {
+				failures = append(failures, lint.Failure{
+					Category:   "style",
+					Confidence: 0.9,
+					Node:       funcDecl,
+					Failure:    "error should be the last type when returning multiple items",
+				})
+
+				break // only flag one
+			}
+		}
+	}
 
 	return failures
 }
@@ -31,37 +49,3 @@ func (*ErrorReturnRule) Apply(file *lint.File, _ lint.Arguments) []lint.Failure
 func (*ErrorReturnRule) Name() string {
 	return "error-return"
 }
-
-type lintErrorReturn struct {
-	file      *lint.File
-	fileAst   *ast.File
-	onFailure func(lint.Failure)
-}
-
-func (w lintErrorReturn) Visit(n ast.Node) ast.Visitor {
-	fn, ok := n.(*ast.FuncDecl)
-	if !ok || fn.Type.Results == nil {
-		return w
-	}
-	ret := fn.Type.Results.List
-	if len(ret) <= 1 {
-		return w
-	}
-	if isIdent(ret[len(ret)-1].Type, "error") {
-		return nil
-	}
-	// An error return parameter should be the last parameter.
-	// Flag any error parameters found before the last.
-	for _, r := range ret[:len(ret)-1] {
-		if isIdent(r.Type, "error") {
-			w.onFailure(lint.Failure{
-				Category:   "arg-order",
-				Confidence: 0.9,
-				Node:       fn,
-				Failure:    "error should be the last type when returning multiple items",
-			})
-			break // only flag one
-		}
-	}
-	return w
-}


Test patch:
diff --git a/test/argument_limit_test.go b/test/argument_limit_test.go
index cbe6be2..b89b0ec 100644
--- a/test/argument_limit_test.go
+++ b/test/argument_limit_test.go
@@ -7,12 +7,19 @@ import (
 	"github.com/mgechev/revive/rule"
 )
 
-func TestArgumentLimitDefault(t *testing.T) {
+func TestArgumentsLimitDefault(t *testing.T) {
 	testRule(t, "argument_limit_default", &rule.ArgumentsLimitRule{}, &lint.RuleConfig{})
 }
 
-func TestArgumentLimit(t *testing.T) {
+func TestArgumentsLimit(t *testing.T) {
 	testRule(t, "argument_limit", &rule.ArgumentsLimitRule{}, &lint.RuleConfig{
 		Arguments: []any{int64(3)},
 	})
 }
+
+func BenchmarkArgumentsLimit(b *testing.B) {
+	var t *testing.T
+	for i := 0; i <= b.N; i++ {
+		testRule(t, "argument_limit_default", &rule.ArgumentsLimitRule{}, &lint.RuleConfig{})
+	}
+}
diff --git a/testdata/confusing_results.go b/testdata/confusing_results.go
index fc47abe..6f45414 100644
--- a/testdata/confusing_results.go
+++ b/testdata/confusing_results.go
@@ -15,6 +15,14 @@ func GetTaz(a string, b int) string {
 
 }
 
-func (t *t) GetTaz(a int, b int)  {
+func (t *t) GetTaz(a int, b int) {
 
 }
+
+func namedResults() (a string, b string) {
+	return "nil", "nil"
+}
+
+func pointerResults() (*string, *string) { // MATCH /unnamed results of the same type may be confusing, consider using named results/
+	return nil, nil
+}