Refactor to cache positive and negative suffix lookup results

betteralign.go

@@ -85,6 +85,7 @@ var (
 	testFiles         bool
 	generatedFiles    bool
 	reGeneratedBy     = regexp.MustCompile(`^//\s*Code generated by .* DO NOT EDIT\.$`)
+	testSuffixes      = []string{"_test.go"}
 	generatedSuffixes = []string{"_generated.go", "_gen.go", ".gen.go", ".pb.go", ".pb.gw.go"}
 	ErrStatFile       = errors.New("unable to stat the file")
 	ErrNotRegularFile = errors.New("not a regular file, skipping")
@@ -116,35 +117,28 @@ func run(pass *analysis.Pass) (interface{}, error) {
 	var dFile *dst.File
 
 	applyFixesFset := make(map[string][]byte)
-	generatedFset := make(map[string]struct{})
+	testFset := make(map[string]bool)
+	generatedFset := make(map[string]bool)
 
 	inspect.Preorder(nodeFilter, func(node ast.Node) {
 		fn := pass.Fset.File(node.Pos()).Name()
 
-		if !testFiles && strings.HasSuffix(fn, "_test.go") {
+		if !testFiles && hasSuffixes(testFset, fn, testSuffixes) {
 			return
 		}
 
-		if !generatedFiles {
-			if _, ok := generatedFset[fn]; ok {
-				return
-			}
-
-			for _, s := range generatedSuffixes {
-				if strings.HasSuffix(fn, s) {
-					generatedFset[fn] = struct{}{}
-					return
-				}
-			}
+		if !generatedFiles && hasSuffixes(generatedFset, fn, generatedSuffixes) {
+			return
 		}
 
 		if f, ok := node.(*ast.File); ok {
 			aFile = f
 			dFile, _ = dec.DecorateFile(aFile)
 
-			if isGeneratedFile(aFile) {
-				generatedFset[fn] = struct{}{}
+			if !generatedFiles && hasGeneratedComment(generatedFset, fn, aFile) {
+				return
 			}
+
 			return
 		}
 
@@ -195,7 +189,7 @@ func betteralign(pass *analysis.Pass, aNode *ast.StructType, typ *types.Struct,
 
 	dNode := dec.Dst.Nodes[aNode].(*dst.StructType)
 
-	if isIgnoreStruct(dNode.Fields) {
+	if hasIgnoreComment(dNode.Fields) {
 		return
 	}
 
@@ -467,14 +461,33 @@ func (s *gcSizes) ptrdata(T types.Type) int64 {
 	panic("impossible")
 }
 
-func isGeneratedFile(file *ast.File) bool {
+func hasSuffixes(fset map[string]bool, fn string, suffixes []string) bool {
+	if t, ok := fset[fn]; ok {
+		return t
+	} else {
+		for _, s := range suffixes {
+			if strings.HasSuffix(fn, s) {
+				fset[fn] = true
+
+				return true
+			}
+		}
+	}
+
+	fset[fn] = false
+
+	return false
+}
+
+func hasGeneratedComment(generatedFset map[string]bool, fn string, file *ast.File) bool {
 	for _, cg := range file.Comments {
 		if cg.Pos() > file.Package {
 			return false
 		}
 
 		for _, l := range cg.List {
 			if reGeneratedBy.MatchString(l.Text) {
+				generatedFset[fn] = true
 				return true
 			}
 		}
@@ -483,7 +496,7 @@ func isGeneratedFile(file *ast.File) bool {
 	return false
 }
 
-func isIgnoreStruct(node *dst.FieldList) bool {
+func hasIgnoreComment(node *dst.FieldList) bool {
 	for _, opening := range node.Decs.Opening.All() {
 		if strings.HasPrefix(opening, "//") && strings.Contains(opening, ignoreStruct) {
 			return true