Using go/ast for iota declarations

user_6398738 0
Viewed 406

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!

iota go abstract-syntax-tree
1 Answers

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