Using go/ast for iota declarations

Question

I've been working with go/ast to parse go source code and copy it into another file as part of a vendoring exercise. I've got most things handled - functions, types etc - but I'm struggling with const declarations which use iota. I'm iterating through the items in ast.File.Scope.Objects and copying over the source for objects with Scope.Outer == nil and their Decl == ast.ValueSpec, basically implying top level variables and constants.

In a block of type:

const (
    a = iota
    b
    c
    d
)

...each one of them registers as a separate object, which is fair enough. However, I'm struggling to assign them with values because the objects can also be out of order when I'm iterating through them. I'm able to see the value as ast.Object.Data for these, but it also seems off when it's set to 1 << iota and so on. Does anyone have any thoughts on how I can get a grouped const declaration with the correct iota values assigned?

Thank you!

Answer 1

I was working on this problem for the exhaustive analyzer, which, during its enum discovery phase, needs to find constant values.

---

Playground: https://play.golang.org/p/nZLmgE4rJZH

Consider the following ConstDecl. It comprises of 3 ConstSpec, and each ConstSpec has 2 names. (This example uses iota, but the approach below should work generally for any ConstDecl.)

package example

const (
    A, B = iota, iota * 100 // 0, 0
    _, D                    // 1, 100
    E, F                    // 2, 200
)

With go/ast and go/types (or x/tools/go/packages), we can obtain a *ast.GenDecl representing the above ConstDecl and a *types.Info for the package.

var decl *ast.GenDecl
var info *types.Info

The following will be true of decl.

decl.Tok == token.CONST

To obtain the constant values, we can do:

func printValuesConst(decl *ast.GenDecl, info *types.Info) {
    for _, s := range decl.Specs {
        v := s.(*ast.ValueSpec) // safe because decl.Tok == token.CONST
        for _, name := range v.Names {
            c := info.ObjectOf(name).(*types.Const)
            fmt.Println(name, c.Val().ExactString())
        }
    }
}

This will, as expected, print:

A 0
B 0
_ 1
D 100
E 2
F 200

Side note: var instead of const

Note that the code above works for const blocks; for a var block we will have to obtain the value using v.Values[i] (assuming a value exists at the index i).

Playground: https://play.golang.org/p/f4mYjXvsvHB

decl.Tok == token.VAR

func printValuesVar(decl *ast.GenDecl, info *types.Info) {
    for _, s := range decl.Specs {
        v := s.(*ast.ValueSpec) // safe because decl.Tok == token.VAR
        for i, name := range v.Names {
            if len(v.Values) <= i {
                fmt.Println(name, "(no AST value)")
                continue
            }
            tv := info.Types[v.Values[i]]
            if tv.Value == nil {
                fmt.Println(name, "(not constant value)")
                continue
            }
            fmt.Println(name, tv.Value.ExactString())
        }
    }
}

Answer 2

On the *ast.File type:

• Scope.Objects tells you the names and values of the const declarations, but does not have them in the correct order, so you can't use it to infer the type.

• Decls gives you the declarations in order, so you can use it to infer the type

You can iterate on both list. First on the Scope.Objects to get the const values, then on Decls to infer the type.

Samplecode on the playground: https://go.dev/play/p/yEz6G1Kqoe3

I'll provide a sample code below. To make it more concrete, there are two assumptions:

• We have a file *ast.File which we got from the parser package (for example, by calling parser.ParseDir(..)

• We have a list of enum type name we want to get the consts for.

type EnumInfo struct {
    Name     string
    TypeName string // base type
    Consts   []ConstValue
}

type ConstValue struct {
    Name  string
    Value any // int or string
}

and we have a enumTypesMap map[string]*EnumInfo that maps a type name to the EnumInfo object.

The purpose of the following snippet is to populate the Consts list of the EnumInfo struct. We assume we have the name and base typename via other means (e.g. reflection).

Modify to suit your own purposes

type EnumInfo struct {
    Name     string
    TypeName string // base type
    Consts   []ConstValue
}

type ConstValue struct {
    Name  string
    Value any // int or string
}

func PopulateEnumInfo(enumTypesMap map[string]*EnumInfo, file *ast.File) {
    // phase 1: iterate scope objects to get the values
    var nameValues = make(map[string]any)

    for _, object := range file.Scope.Objects {
        if object.Kind == ast.Con {
            nameValues[object.Name] = object.Data
        }
    }

    // phase 2: iterate decls to get the type and names in order
    for _, decl := range file.Decls {
        genDecl, ok := decl.(*ast.GenDecl)
        if !ok {
            continue
        }
        if genDecl.Tok != token.CONST {
            continue
        }
        var enumInfo *EnumInfo
        for _, spec := range genDecl.Specs {
            valSpec := spec.(*ast.ValueSpec)
            if typeIdent, ok := valSpec.Type.(*ast.Ident); ok {
                enumInfo = enumTypesMap[typeIdent.String()]
            }
            if enumInfo != nil {
                for _, nameIdent := range valSpec.Names {
                    name := nameIdent.String()
                    if name == "_" {
                        continue
                    }
                    value := nameValues[name]
                    enumInfo.Consts = append(enumInfo.Consts, ConstValue{
                        Name:  name,
                        Value: value,
                    })
                }
            }
        }
    }
}