chore: Improve sortables detection (#1151)

Co-authored-by: chavacava <salvador.cavadini@gmail.com>

lint/package.go

@@ -1,6 +1,7 @@
 package lint
 
 import (
+	"errors"
 	"go/ast"
 	"go/importer"
 	"go/token"
@@ -9,6 +10,7 @@ import (
 
 	goversion "github.com/hashicorp/go-version"
 
+	"github.com/mgechev/revive/internal/astutils"
 	"github.com/mgechev/revive/internal/typeparams"
 )
 
@@ -31,7 +33,6 @@ type Package struct {
 var (
 	trueValue  = 1
 	falseValue = 2
-	notSet     = 3
 
 	go121 = goversion.Must(goversion.NewVersion("1.21"))
 	go122 = goversion.Must(goversion.NewVersion("1.22"))
@@ -111,6 +112,11 @@ func (p *Package) TypeCheck() error {
 		astFiles = append(astFiles, f.AST)
 	}
 
+	if anyFile == nil {
+		// this is unlikely to happen, but technically guarantees anyFile to not be nil
+		return errors.New("no ast.File found")
+	}
+
 	typesPkg, err := check(config, anyFile.AST.Name.Name, p.fset, astFiles, info)
 
 	// Remember the typechecking info, even if config.Check failed,
@@ -135,7 +141,7 @@ func check(config *types.Config, n string, fset *token.FileSet, astFiles []*ast.
 	return config.Check(n, fset, astFiles, info)
 }
 
-// TypeOf returns the type of an expression.
+// TypeOf returns the type of expression.
 func (p *Package) TypeOf(expr ast.Expr) types.Type {
 	if p.typesInfo == nil {
 		return nil
@@ -143,39 +149,32 @@ func (p *Package) TypeOf(expr ast.Expr) types.Type {
 	return p.typesInfo.TypeOf(expr)
 }
 
-type walker struct {
-	nmap map[string]int
-	has  map[string]int
-}
+type sortableMethodsFlags int
 
-func (w *walker) Visit(n ast.Node) ast.Visitor {
-	fn, ok := n.(*ast.FuncDecl)
-	if !ok || fn.Recv == nil || len(fn.Recv.List) == 0 {
-		return w
-	}
-	// TODO(dsymonds): We could check the signature to be more precise.
-	recv := typeparams.ReceiverType(fn)
-	if i, ok := w.nmap[fn.Name.Name]; ok {
-		w.has[recv] |= i
-	}
-	return w
-}
+// flags for sortable interface methods.
+const (
+	bfLen sortableMethodsFlags = 1 << iota
+	bfLess
+	bfSwap
+)
 
 func (p *Package) scanSortable() {
-	p.sortable = map[string]bool{}
-
-	// bitfield for which methods exist on each type.
-	const (
-		bfLen = 1 << iota
-		bfLess
-		bfSwap
-	)
-	nmap := map[string]int{"Len": bfLen, "Less": bfLess, "Swap": bfSwap}
-	has := map[string]int{}
+	sortableFlags := map[string]sortableMethodsFlags{}
 	for _, f := range p.files {
-		ast.Walk(&walker{nmap, has}, f.AST)
+		for _, decl := range f.AST.Decls {
+			fn, ok := decl.(*ast.FuncDecl)
+			isAMethodDeclaration := ok && fn.Recv != nil && len(fn.Recv.List) != 0
+			if !isAMethodDeclaration {
+				continue
+			}
+
+			recvType := typeparams.ReceiverType(fn)
+			sortableFlags[recvType] |= getSortableMethodFlagForFunction(fn)
+		}
 	}
-	for typ, ms := range has {
+
+	p.sortable = make(map[string]bool, len(sortableFlags))
+	for typ, ms := range sortableFlags {
 		if ms == bfLen|bfLess|bfSwap {
 			p.sortable[typ] = true
 		}
@@ -204,3 +203,16 @@ func (p *Package) IsAtLeastGo121() bool {
 func (p *Package) IsAtLeastGo122() bool {
 	return p.goVersion.GreaterThanOrEqual(go122)
 }
+
+func getSortableMethodFlagForFunction(fn *ast.FuncDecl) sortableMethodsFlags {
+	switch {
+	case astutils.FuncSignatureIs(fn, "Len", []string{}, []string{"int"}):
+		return bfLen
+	case astutils.FuncSignatureIs(fn, "Less", []string{"int", "int"}, []string{"bool"}):
+		return bfLess
+	case astutils.FuncSignatureIs(fn, "Swap", []string{"int", "int"}, []string{}):
+		return bfSwap
+	default:
+		return 0
+	}
+}