chore: use iterator

Author: ldez

sloglint.go

@@ -9,6 +9,7 @@ import (
 	"go/token"
 	"go/types"
 	"go/version"
+	"iter"
 	"slices"
 	"strconv"
 	"strings"
@@ -320,30 +321,33 @@ func visit(pass *analysis.Pass, opts *Options, node ast.Node, stack []ast.Node)
 	}
 
 	if opts.NoRawKeys {
-		forEachKey(pass.TypesInfo, keys, attrs, func(key ast.Expr) {
+		for key := range AllKeys(pass.TypesInfo, keys, attrs) {
 			if sel, ok := key.(*ast.SelectorExpr); ok {
 				key = sel.Sel // the key is defined in another package, e.g. pkg.ConstKey.
 			}
+
 			isConst := false
+
 			if ident, ok := key.(*ast.Ident); ok {
 				if obj := pass.TypesInfo.ObjectOf(ident); obj != nil {
 					if _, ok := obj.(*types.Const); ok {
 						isConst = true
 					}
 				}
 			}
+
 			if !isConst {
 				pass.Reportf(call.Pos(), "raw keys should not be used")
 			}
-		})
+		}
 	}
 
 	checkKeyNamingCase := func(caseFn func(string) string, caseName string) {
-		forEachKey(pass.TypesInfo, keys, attrs, func(key ast.Expr) {
+		for key := range AllKeys(pass.TypesInfo, keys, attrs) {
 			if name, ok := getKeyName(key); ok && name != caseFn(name) {
 				pass.Reportf(call.Pos(), "keys should be written in %s", caseName)
 			}
-		})
+		}
 	}
 
 	switch opts.KeyNamingCase {
@@ -358,11 +362,11 @@ func visit(pass *analysis.Pass, opts *Options, node ast.Node, stack []ast.Node)
 	}
 
 	if len(opts.ForbiddenKeys) > 0 {
-		forEachKey(pass.TypesInfo, keys, attrs, func(key ast.Expr) {
+		for key := range AllKeys(pass.TypesInfo, keys, attrs) {
 			if name, ok := getKeyName(key); ok && slices.Contains(opts.ForbiddenKeys, name) {
 				pass.Reportf(call.Pos(), "%q key is forbidden and should not be used", name)
 			}
-		})
+		}
 	}
 
 	if opts.ArgsOnSepLines && areArgsOnSameLine(pass.Fset, call, keys, attrs) {
@@ -444,36 +448,52 @@ func isValidMsgStyle(msg, style string) bool {
 	}
 }
 
-func forEachKey(info *types.Info, keys, attrs []ast.Expr, fn func(key ast.Expr)) {
-	for _, key := range keys {
-		fn(key)
-	}
-
-	for _, attr := range attrs {
-		switch attr := attr.(type) {
-		case *ast.CallExpr: // e.g. slog.Int()
-			callee := typeutil.StaticCallee(info, attr)
-			if callee == nil {
-				continue
-			}
-			if _, ok := attrFuncs[callee.FullName()]; !ok {
-				continue
+func AllKeys(info *types.Info, keys, attrs []ast.Expr) iter.Seq[ast.Expr] {
+	return func(yield func(key ast.Expr) bool) {
+		for _, key := range keys {
+			if !yield(key) {
+				return
 			}
-			fn(attr.Args[0])
+		}
+
+		for _, attr := range attrs {
+			switch attr := attr.(type) {
+			case *ast.CallExpr: // e.g. slog.Int()
+				callee := typeutil.StaticCallee(info, attr)
+				if callee == nil {
+					continue
+				}
+				if _, ok := attrFuncs[callee.FullName()]; !ok {
+					continue
+				}
 
-		case *ast.CompositeLit: // slog.Attr{}
-			switch len(attr.Elts) {
-			case 1: // slog.Attr{Key: ...} | slog.Attr{Value: ...}
-				if kv := attr.Elts[0].(*ast.KeyValueExpr); kv.Key.(*ast.Ident).Name == "Key" {
-					fn(kv.Value)
+				if !yield(attr.Args[0]) {
+					return
 				}
-			case 2: // slog.Attr{Key: ..., Value: ...} | slog.Attr{Value: ..., Key: ...} | slog.Attr{..., ...}
-				if kv, ok := attr.Elts[0].(*ast.KeyValueExpr); ok && kv.Key.(*ast.Ident).Name == "Key" {
-					fn(kv.Value)
-				} else if kv, ok := attr.Elts[1].(*ast.KeyValueExpr); ok && kv.Key.(*ast.Ident).Name == "Key" {
-					fn(kv.Value)
-				} else {
-					fn(attr.Elts[0])
+
+			case *ast.CompositeLit: // slog.Attr{}
+				switch len(attr.Elts) {
+				case 1: // slog.Attr{Key: ...} | slog.Attr{Value: ...}
+					if kv := attr.Elts[0].(*ast.KeyValueExpr); kv.Key.(*ast.Ident).Name == "Key" {
+						if !yield(kv.Value) {
+							return
+						}
+					}
+
+				case 2: // slog.Attr{Key: ..., Value: ...} | slog.Attr{Value: ..., Key: ...} | slog.Attr{..., ...}
+					if kv, ok := attr.Elts[0].(*ast.KeyValueExpr); ok && kv.Key.(*ast.Ident).Name == "Key" {
+						if !yield(kv.Value) {
+							return
+						}
+					} else if kv, ok := attr.Elts[1].(*ast.KeyValueExpr); ok && kv.Key.(*ast.Ident).Name == "Key" {
+						if !yield(kv.Value) {
+							return
+						}
+					} else {
+						if !yield(attr.Elts[0]) {
+							return
+						}
+					}
 				}
 			}
 		}