Using go/ast for iota declarations - go

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())
}
}
}