validate map keys the same way as switch cases

README.md

@@ -1,6 +1,6 @@
 ## exhaustive [![Godoc][2]][1]
 
-Check exhaustiveness of enum switch statements in Go source code.
+Check exhaustiveness of enum switch statements and map literals in Go source code.
 
 ```
 go install github.com/nishanths/exhaustive/cmd/exhaustive@latest
@@ -39,7 +39,7 @@ const (
 )
 ```
 
-and the switch statement
+and the code
 
 ```go
 package calc
@@ -54,12 +54,21 @@ func f(t token.Token) {
 	default:
 	}
 }
+
+func g(t token.Token) string {
+	return map[token.Token]string{
+		token.Add:      "add",
+		token.Subtract: "subtract",
+		token.Multiply: "multiply",
+	}[t]
+}
 ```
 
 running exhaustive will print
 
 ```
 calc.go:6:2: missing cases in switch of type token.Token: Quotient, Remainder
+calc.go:15:9: missing map keys of type token.Token: Quotient, Remainder
 ```
 
 ## Contributing

comment.go

@@ -2,6 +2,7 @@ package exhaustive
 
 import (
 	"go/ast"
+	"go/token"
 	"regexp"
 	"strings"
 )
@@ -50,10 +51,28 @@ func isGeneratedFileComment(s string) bool {
 	return generatedCodeRe.MatchString(s)
 }
 
+type generatedCache map[*ast.File]bool
+
+func (c generatedCache) IsGenerated(file *ast.File) bool {
+	if _, ok := c[file]; !ok {
+		c[file] = isGeneratedFile(file)
+	}
+	return c[file]
+}
+
 // ignoreDirective is used to exclude checking of specific switch statements.
 const ignoreDirective = "//exhaustive:ignore"
 const enforceDirective = "//exhaustive:enforce"
 
+type commentsCache map[*ast.File]ast.CommentMap
+
+func (c commentsCache) GetComments(file *ast.File, set *token.FileSet) ast.CommentMap {
+	if _, ok := c[file]; !ok {
+		c[file] = ast.NewCommentMap(set, file, file.Comments)
+	}
+	return c[file]
+}
+
 func containsDirective(comments []*ast.CommentGroup, directive string) bool {
 	for _, c := range comments {
 		for _, cc := range c.List {

exhaustive.go

@@ -1,6 +1,6 @@
 /*
 Package exhaustive provides an analyzer that checks exhaustiveness of enum
-switch statements in Go source code.
+switch statements and map literals in Go source code.
 
 Definition of enum
 
@@ -52,6 +52,11 @@ The analyzer will produce a diagnostic about unhandled enum members if the
 required memebers are not listed in a switch statement's cases (this applies
 even if the switch statement has a 'default' case).
 
+Map literals
+
+All above relates to map literals as well, if key is an enum type.
+But empty map is ignored because it's an alternative for make(map...).
+
 Type aliases
 
 The analyzer handles type aliases for an enum type in the following manner.
@@ -115,6 +120,7 @@ All of these flags are optional.
     flag                            type    default value
 
     -explicit-exhaustive-switch     bool    false
+    -explicit-exhaustive-map        bool    false
     -check-generated                bool    false
     -default-signifies-exhaustive   bool    false
     -ignore-enum-members            string  (none)
@@ -126,6 +132,11 @@ switch statements explicitly marked with the comment text
 ("exhaustive:enforce"). Otherwise, it runs on every enum switch statement not
 marked with the comment text ("exhaustive:ignore").
 
+If the -explicit-exhaustive-map flag is enabled, the analyzer only runs on
+map literals explicitly marked with the comment text
+("exhaustive:enforce"). Otherwise, it runs on every enum map literal not
+marked with the comment text ("exhaustive:ignore").
+
 If the -check-generated flag is enabled, switch statements in generated Go
 source files are also checked. Otherwise, by default, switch statements in
 generated files are not checked. See https://golang.org/s/generatedcode for the
@@ -176,6 +187,7 @@ package exhaustive
 
 import (
 	"flag"
+	"go/ast"
 	"regexp"
 
 	"golang.org/x/tools/go/analysis"
@@ -215,6 +227,7 @@ func (v *regexpFlag) value() *regexp.Regexp { return v.r }
 
 func init() {
 	Analyzer.Flags.BoolVar(&fExplicitExhaustiveSwitch, ExplicitExhaustiveSwitchFlag, false, "only run exhaustive check on switches with \"//exhaustive:enforce\" comment")
+	Analyzer.Flags.BoolVar(&fExplicitExhaustiveMap, ExplicitExhaustiveMapFlag, false, "only run exhaustive check on map literals with \"//exhaustive:enforce\" comment")
 	Analyzer.Flags.BoolVar(&fCheckGenerated, CheckGeneratedFlag, false, "check switch statements in generated files")
 	Analyzer.Flags.BoolVar(&fDefaultSignifiesExhaustive, DefaultSignifiesExhaustiveFlag, false, "presence of \"default\" case in switch statements satisfies exhaustiveness, even if all enum members are not listed")
 	Analyzer.Flags.Var(&fIgnoreEnumMembers, IgnoreEnumMembersFlag, "enum members matching `regex` do not have to be listed in switch statements to satisfy exhaustiveness")
@@ -229,6 +242,7 @@ func init() {
 // driver programs.
 const (
 	ExplicitExhaustiveSwitchFlag   = "explicit-exhaustive-switch"
+	ExplicitExhaustiveMapFlag      = "explicit-exhaustive-map"
 	CheckGeneratedFlag             = "check-generated"
 	DefaultSignifiesExhaustiveFlag = "default-signifies-exhaustive"
 	IgnoreEnumMembersFlag          = "ignore-enum-members"
@@ -240,6 +254,7 @@ const (
 
 var (
 	fExplicitExhaustiveSwitch   bool
+	fExplicitExhaustiveMap      bool
 	fCheckGenerated             bool
 	fDefaultSignifiesExhaustive bool
 	fIgnoreEnumMembers          regexpFlag
@@ -250,6 +265,7 @@ var (
 // Useful in tests.
 func resetFlags() {
 	fExplicitExhaustiveSwitch = false
+	fExplicitExhaustiveMap = false
 	fCheckGenerated = false
 	fDefaultSignifiesExhaustive = false
 	fIgnoreEnumMembers = regexpFlag{}
@@ -258,7 +274,7 @@ func resetFlags() {
 
 var Analyzer = &analysis.Analyzer{
 	Name:      "exhaustive",
-	Doc:       "check exhaustiveness of enum switch statements",
+	Doc:       "check exhaustiveness of enum switch statements and map literals",
 	Run:       run,
 	Requires:  []*analysis.Analyzer{inspect.Analyzer},
 	FactTypes: []analysis.Fact{&enumMembersFact{}},
@@ -271,11 +287,47 @@ func run(pass *analysis.Pass) (interface{}, error) {
 		exportFact(pass, typ, members)
 	}
 
-	checkSwitchStatements(pass, inspect, config{
-		explicitExhaustiveSwitch:   fExplicitExhaustiveSwitch,
-		defaultSignifiesExhaustive: fDefaultSignifiesExhaustive,
-		checkGeneratedFiles:        fCheckGenerated,
-		ignoreEnumMembers:          fIgnoreEnumMembers.value(),
+	generated := make(generatedCache)
+	comments := make(commentsCache)
+
+	checkSwitch := switchStmtChecker(
+		pass,
+		switchConfig{
+			explicitExhaustiveSwitch:   fExplicitExhaustiveSwitch,
+			defaultSignifiesExhaustive: fDefaultSignifiesExhaustive,
+			checkGeneratedFiles:        fCheckGenerated,
+			ignoreEnumMembers:          fIgnoreEnumMembers.value(),
+		},
+		generated,
+		comments,
+	)
+
+	checkMap := mapChecker(
+		pass,
+		mapConfig{
+			explicitExhaustiveMap: fExplicitExhaustiveMap,
+			checkGeneratedFiles:   fCheckGenerated,
+			ignoreEnumMembers:     fIgnoreEnumMembers.value(),
+		},
+		generated,
+		comments,
+	)
+
+	filter := []ast.Node{
+		&ast.SwitchStmt{},
+		&ast.CompositeLit{},
+	}
+
+	inspect.WithStack(filter, func(n ast.Node, push bool, stack []ast.Node) bool {
+		var proceed bool
+		switch n.(type) {
+		case *ast.SwitchStmt:
+			proceed, _ = checkSwitch(n, push, stack)
+		case *ast.CompositeLit:
+			proceed, _ = checkMap(n, push, stack)
+		}
+		return proceed
 	})
+
 	return nil, nil
 }

exhaustive_test.go

@@ -112,7 +112,10 @@ func TestExhaustive(t *testing.T) {
 
 	// Switch statements without enforce directive comment should not be checked during explicitly exhaustive
 	// switch mode
-	run(t, "enforce-comment/...", func() { fExplicitExhaustiveSwitch = true })
+	run(t, "enforce-comment/...", func() {
+		fExplicitExhaustiveSwitch = true
+		fExplicitExhaustiveMap = true
+	})
 
 	// To satisfy exhaustiveness, it is sufficient for each unique constant
 	// value of the members to be listed, not each member by name.

map.go

@@ -0,0 +1,147 @@
+package exhaustive
+
+import (
+	"fmt"
+	"go/ast"
+	"go/types"
+	"golang.org/x/tools/go/analysis"
+	"regexp"
+	"strings"
+)
+
+// mapConfig is configuration for mapChecker.
+type mapConfig struct {
+	explicitExhaustiveMap bool
+	checkGeneratedFiles   bool
+	ignoreEnumMembers     *regexp.Regexp // can be nil
+}
+
+// mapChecker returns a node visitor that checks exhaustiveness
+// of enum keys in map literal for the supplied pass, and reports diagnostics if non-exhaustive.
+// It expects to only see *ast.CompositeLit nodes.
+func mapChecker(pass *analysis.Pass, cfg mapConfig, generated generatedCache, comments commentsCache) nodeVisitor {
+	return func(n ast.Node, push bool, stack []ast.Node) (bool, string) {
+		if !push {
+			// The proceed return value should not matter; it is ignored by
+			// inspector package for pop calls.
+			// Nevertheless, return true to be on the safe side for the future.
+			return true, resultNotPush
+		}
+
+		file := stack[0].(*ast.File)
+
+		if !cfg.checkGeneratedFiles && generated.IsGenerated(file) {
+			// Don't check this file.
+			// Return false because the children nodes of node `n` don't have to be checked.
+			return false, resultGeneratedFile
+		}
+
+		lit := n.(*ast.CompositeLit)
+
+		mapType, ok := pass.TypesInfo.Types[lit.Type].Type.(*types.Map)
+		if !ok {
+			namedType, ok2 := pass.TypesInfo.Types[lit.Type].Type.(*types.Named)
+			if !ok2 {
+				return true, resultNotMapLiteral
+			}
+
+			mapType, ok = namedType.Underlying().(*types.Map)
+			if !ok {
+				return true, resultNotMapLiteral
+			}
+		}
+
+		if len(lit.Elts) == 0 {
+			// because it may be used as an alternative for make(map[...]...)
+			return false, resultEmptyMapLiteral
+		}
+
+		keyType, ok := mapType.Key().(*types.Named)
+		if !ok {
+			return true, resultMapKeyIsNotNamedType
+		}
+
+		fileComments := comments.GetComments(file, pass.Fset)
+		var relatedComments []*ast.CommentGroup
+		for i := range stack {
+			// iterate over stack in the reverse order (from bottom to top)
+			node := stack[len(stack)-1-i]
+			switch node.(type) {
+			// need to check comments associated with following nodes,
+			// because logic of ast package doesn't allow to associate comment with *ast.CompositeLit
+			case *ast.CompositeLit, // stack[len(stack)-1]
+				*ast.ReturnStmt, // return ...
+				*ast.IndexExpr,  // map[enum]...{...}[key]
+				*ast.CallExpr,   // myfunc(map...)
+				*ast.UnaryExpr,  // &map...
+				*ast.AssignStmt, // variable assignment (without var keyword)
+				*ast.DeclStmt,   // var declaration, parent of *ast.GenDecl
+				*ast.GenDecl,    // var declaration, parent of *ast.ValueSpec
+				*ast.ValueSpec:  // var declaration
+				relatedComments = append(relatedComments, fileComments[node]...)
+				continue
+			}
+			// stop iteration on the first inappropriate node
+			break
+		}
+
+		if !cfg.explicitExhaustiveMap && containsIgnoreDirective(relatedComments) {
+			// Skip checking of this map literal due to ignore directive comment.
+			// Still return true because there may be nested map literals
+			// that are not to be ignored.
+			return true, resultMapIgnoreComment
+		}
+		if cfg.explicitExhaustiveMap && !containsEnforceDirective(relatedComments) {
+			// Skip checking of this map literal due to missing enforce directive comment.
+			return true, resultMapNoEnforceComment
+		}
+
+		keyPkg := keyType.Obj().Pkg()
+		if keyPkg == nil {
+			// The Go documentation says: nil for labels and objects in the Universe scope.
+			// This happens for the `error` type, for example.
+			return true, resultNilMapKeyTypePkg
+		}
+
+		enumTyp := enumType{keyType.Obj()}
+		members, ok := importFact(pass, enumTyp)
+		if !ok {
+			return true, resultMapKeyNotEnum
+		}
+
+		samePkg := keyPkg == pass.Pkg // do the map literal and the map key type (i.e. enum type) live in the same package?
+		checkUnexported := samePkg    // we want to include unexported members in the exhaustiveness check only if we're in the same package
+		checklist := makeChecklist(members, keyPkg, checkUnexported, cfg.ignoreEnumMembers)
+
+		for _, e := range lit.Elts {
+			expr, ok := e.(*ast.KeyValueExpr)
+			if !ok {
+				continue // is it possible for valid map literal?
+			}
+			analyzeCaseClauseExpr(expr.Key, pass.TypesInfo, checklist.found)
+		}
+
+		if len(checklist.remaining()) == 0 {
+			// All enum members accounted for.
+			// Nothing to report.
+			return true, resultEnumMembersAccounted
+		}
+
+		pass.Report(makeMapDiagnostic(lit, samePkg, enumTyp, members, checklist.remaining()))
+		return true, resultReportedDiagnostic
+	}
+}
+
+// Makes a "missing map keys" diagnostic.
+// samePkg should be true if the enum type and the map literal are defined in the same package.
+func makeMapDiagnostic(lit *ast.CompositeLit, samePkg bool, enumTyp enumType, allMembers enumMembers, missingMembers map[string]struct{}) analysis.Diagnostic {
+	message := fmt.Sprintf("missing map keys of type %s: %s",
+		diagnosticEnumTypeName(enumTyp.TypeName, samePkg),
+		strings.Join(diagnosticMissingMembers(missingMembers, allMembers), ", "))
+
+	return analysis.Diagnostic{
+		Pos:     lit.Pos(),
+		End:     lit.End(),
+		Message: message,
+	}
+}