wg-quicker/vendor/github.com/nishanths/exhaustive/enum.go

147 lines
3.4 KiB
Go

package exhaustive
import (
"go/ast"
"go/token"
"go/types"
"golang.org/x/tools/go/analysis"
)
type enums map[string]*enumMembers // enum type name -> enum members
type enumMembers struct {
// Names in the order encountered in the AST.
OrderedNames []string
// Maps name -> (constant.Value).ExactString().
// If a name is missing in the map, it means that it does not have a
// corresponding constant.Value defined in the AST.
NameToValue map[string]string
// Maps (constant.Value).ExactString() -> names.
// Names that don't have a constant.Value defined in the AST (e.g., some
// iota constants) will not have a corresponding entry in this map.
ValueToNames map[string][]string
}
func (em *enumMembers) add(name string, constVal *string) {
em.OrderedNames = append(em.OrderedNames, name)
if constVal != nil {
if em.NameToValue == nil {
em.NameToValue = make(map[string]string)
}
em.NameToValue[name] = *constVal
if em.ValueToNames == nil {
em.ValueToNames = make(map[string][]string)
}
em.ValueToNames[*constVal] = append(em.ValueToNames[*constVal], name)
}
}
func (em *enumMembers) numMembers() int {
return len(em.OrderedNames)
}
func findEnums(pass *analysis.Pass) enums {
pkgEnums := make(enums)
// Gather enum types.
for _, f := range pass.Files {
for _, decl := range f.Decls {
gen, ok := decl.(*ast.GenDecl)
if !ok {
continue
}
if gen.Tok != token.TYPE {
continue
}
for _, s := range gen.Specs {
// Must be TypeSpec since we've filtered on token.TYPE.
t, ok := s.(*ast.TypeSpec)
obj := pass.TypesInfo.Defs[t.Name]
if obj == nil {
continue
}
named, ok := obj.Type().(*types.Named)
if !ok {
continue
}
basic, ok := named.Underlying().(*types.Basic)
if !ok {
continue
}
switch i := basic.Info(); {
case i&types.IsInteger != 0:
pkgEnums[named.Obj().Name()] = &enumMembers{}
case i&types.IsFloat != 0:
pkgEnums[named.Obj().Name()] = &enumMembers{}
case i&types.IsString != 0:
pkgEnums[named.Obj().Name()] = &enumMembers{}
}
}
}
}
// Gather enum members.
for _, f := range pass.Files {
for _, decl := range f.Decls {
gen, ok := decl.(*ast.GenDecl)
if !ok {
continue
}
if gen.Tok != token.CONST && gen.Tok != token.VAR {
continue
}
for _, s := range gen.Specs {
// Must be ValueSpec since we've filtered on token.CONST, token.VAR.
v := s.(*ast.ValueSpec)
for i, name := range v.Names {
obj := pass.TypesInfo.Defs[name]
if obj == nil {
continue
}
named, ok := obj.Type().(*types.Named)
if !ok {
continue
}
// Get the constant.Value representation, if any.
var constVal *string
if len(v.Values) > i {
value := v.Values[i]
if con, ok := pass.TypesInfo.Types[value]; ok && con.Value != nil {
str := con.Value.ExactString() // temp var to be able to take address
constVal = &str
}
}
em, ok := pkgEnums[named.Obj().Name()]
if !ok {
continue
}
em.add(obj.Name(), constVal)
pkgEnums[named.Obj().Name()] = em
}
}
}
}
// Delete member-less enum types.
// We can't call these enums, since we can't be sure without
// the existence of members. (The type may just be a named type,
// for instance.)
for k, v := range pkgEnums {
if v.numMembers() == 0 {
delete(pkgEnums, k)
}
}
return pkgEnums
}