Create a LISP like representation from a Go AST

main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"log"
	"os"
)

var (
	flagFn = flag.String("fn", "", "input filename")
)

func writeIndent(buf *bytes.Buffer, i int) {
	for j := 0; j < i; j++ {
		buf.WriteString(" ")
	}
}

func main() {
	flag.Parse()

	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, *flagFn, nil, 0)
	if err != nil {
		log.Fatal(err)
	}

	var buf bytes.Buffer

	tree := buildTree(f)

	var visitTree func(SyntaxNode, int)
	visitTree = func(n SyntaxNode, i int) {
		switch n := n.(type) {
		case *SyntaxTree:
			{
				buf.WriteString(fmt.Sprintf("(%v", n.Label()))

				children := n.Children()
				if len(children) > 0 {
					if len(children) > 1 {
						buf.WriteString("\n")

						for _, child := range children {
							writeIndent(&buf, i+1)
							visitTree(child, i+1)
							buf.WriteString("\n")
						}

						writeIndent(&buf, i)
					} else {
						buf.WriteString(" ")

						for _, child := range children {
							visitTree(child, i+1)
						}
					}
				}

				buf.WriteString(")")
			}
		case *SyntaxToken:
			{
				buf.WriteString(fmt.Sprintf("(%v)", n.Label()))
			}
		}
	}

	visitTree(tree, 0)

	os.Stdout.WriteString(buf.String())
}

type NodeIterator interface {
	Len() int
	Get(int) ast.Node
}

type DeclNodeIterator struct {
	decls []ast.Decl
}

func (it *DeclNodeIterator) Len() int           { return len(it.decls) }
func (it *DeclNodeIterator) Get(i int) ast.Node { return it.decls[i] }

type SpecNodeIterator struct {
	specs []ast.Spec
}

func (it *SpecNodeIterator) Len() int           { return len(it.specs) }
func (it *SpecNodeIterator) Get(i int) ast.Node { return it.specs[i] }

type IdentNodeIterator struct {
	idents []*ast.Ident
}

func (it *IdentNodeIterator) Len() int           { return len(it.idents) }
func (it *IdentNodeIterator) Get(i int) ast.Node { return it.idents[i] }

type ExprNodeIterator struct {
	exprs []ast.Expr
}

func (it *ExprNodeIterator) Len() int           { return len(it.exprs) }
func (it *ExprNodeIterator) Get(i int) ast.Node { return it.exprs[i] }

type FieldNodeIterator struct {
	fields []*ast.Field
}

func (it *FieldNodeIterator) Len() int           { return len(it.fields) }
func (it *FieldNodeIterator) Get(i int) ast.Node { return it.fields[i] }

type StmtNodeIterator struct {
	stmts []ast.Stmt
}

func (it *StmtNodeIterator) Len() int           { return len(it.stmts) }
func (it *StmtNodeIterator) Get(i int) ast.Node { return it.stmts[i] }

func buildTree(n ast.Node) SyntaxNode {
	switch n := n.(type) {
	case *ast.File:
		return &SyntaxTree{
			node:     n,
			children: append([]SyntaxNode{buildTree(n.Name)}, buildSubtreeFromList(&DeclNodeIterator{n.Decls})...),
		}
	case *ast.Ident:
		return &SyntaxToken{
			node: n,
			v:    n.Name,
		}
	case *ast.GenDecl:
		return &SyntaxTree{
			node:     n,
			children: buildSubtreeFromList(&SpecNodeIterator{n.Specs}),
		}
	case *ast.ImportSpec:
		if n.Name == nil {
			return &SyntaxTree{
				node:     n,
				children: []SyntaxNode{buildTree(n.Path)},
			}
		}
		return &SyntaxTree{
			node:     n,
			children: []SyntaxNode{buildTree(n.Name), buildTree(n.Path)},
		}
	case *ast.BasicLit:
		return &SyntaxToken{
			node: n,
			v:    n.Value,
		}
	case *ast.ValueSpec:
		var children []SyntaxNode

		children = append(children, buildSubtreeFromList(&IdentNodeIterator{n.Names})...)
		if n.Type != nil {
			children = append(children, buildTree(n.Type))
		}
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.Values})...)

		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.CallExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.Fun))
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.Args})...)
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.SelectorExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		children = append(children, buildTree(n.Sel))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.FuncDecl:
		var children []SyntaxNode
		if n.Recv != nil {
			children = append(children, buildTree(n.Recv))
		}
		children = append(children, buildTree(n.Name))
		children = append(children, buildTree(n.Type))
		if n.Body != nil {
			children = append(children, buildTree(n.Body))
		}
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.FuncType:
		var children []SyntaxNode
		if n.Params != nil {
			children = append(children, buildTree(n.Params))
		}
		if n.Results != nil {
			children = append(children, buildTree(n.Results))
		}
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.FieldList:
		return &SyntaxTree{
			node:     n,
			children: buildSubtreeFromList(&FieldNodeIterator{n.List}),
		}
	case *ast.Field:
		var children []SyntaxNode
		children = append(children, buildSubtreeFromList(&IdentNodeIterator{n.Names})...)
		children = append(children, buildTree(n.Type))
		if n.Tag != nil {
			children = append(children, buildTree(n.Tag))
		}
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.StarExpr:
		return &SyntaxTree{
			node:     n,
			children: []SyntaxNode{buildTree(n.X)},
		}
	case *ast.BlockStmt:
		return &SyntaxTree{
			node:     n,
			children: buildSubtreeFromList(&StmtNodeIterator{n.List}),
		}
	case *ast.ForStmt:
		var children []SyntaxNode
		if n.Init != nil {
			children = append(children, buildTree(n.Init))
		}
		if n.Cond != nil {
			children = append(children, buildTree(n.Cond))
		}
		if n.Post != nil {
			children = append(children, buildTree(n.Post))
		}
		children = append(children, buildTree(n.Body))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.AssignStmt:
		var children []SyntaxNode
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.Lhs})...)
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.Rhs})...)
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.BinaryExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		children = append(children, buildTree(n.Y))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.IncDecStmt:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.ExprStmt:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.IfStmt:
		var children []SyntaxNode
		if n.Init != nil {
			children = append(children, buildTree(n.Init))
		}
		children = append(children, buildTree(n.Cond))
		children = append(children, buildTree(n.Body))
		if n.Else != nil {
			children = append(children, buildTree(n.Else))
		}
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.DeclStmt:
		var children []SyntaxNode
		children = append(children, buildTree(n.Decl))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.FuncLit:
		var children []SyntaxNode
		children = append(children, buildTree(n.Type))
		children = append(children, buildTree(n.Body))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.TypeSpec:
		var children []SyntaxNode
		children = append(children, buildTree(n.Name))
		children = append(children, buildTree(n.Type))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.TypeSwitchStmt:
		var children []SyntaxNode
		if n.Init != nil {
			children = append(children, buildTree(n.Init))
		}
		children = append(children, buildTree(n.Assign))
		children = append(children, buildTree(n.Body))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.InterfaceType:
		var children []SyntaxNode
		children = append(children, buildTree(n.Methods))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.StructType:
		var children []SyntaxNode
		children = append(children, buildTree(n.Fields))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.ReturnStmt:
		var children []SyntaxNode
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.Results})...)
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.TypeAssertExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		if n.Type != nil {
			children = append(children, buildTree(n.Type))
		}
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.CaseClause:
		var children []SyntaxNode
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.List})...)
		children = append(children, buildSubtreeFromList(&StmtNodeIterator{n.Body})...)
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.ArrayType:
		var children []SyntaxNode
		if n.Len != nil {
			children = append(children, buildTree(n.Len))
		}
		children = append(children, buildTree(n.Elt))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.RangeStmt:
		var children []SyntaxNode
		if n.Key != nil {
			children = append(children, buildTree(n.Key))
		}
		if n.Value != nil {
			children = append(children, buildTree(n.Value))
		}
		children = append(children, buildTree(n.X))
		children = append(children, buildTree(n.Body))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.UnaryExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.IndexExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.X))
		children = append(children, buildTree(n.Index))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.CompositeLit:
		var children []SyntaxNode
		if n.Type != nil {
			children = append(children, buildTree(n.Type))
		}
		children = append(children, buildSubtreeFromList(&ExprNodeIterator{n.Elts})...)
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	case *ast.KeyValueExpr:
		var children []SyntaxNode
		children = append(children, buildTree(n.Key))
		children = append(children, buildTree(n.Value))
		return &SyntaxTree{
			node:     n,
			children: children,
		}
	default:
		fmt.Printf("error: don't know what to do with %T\n", n)
		return &SyntaxTree{node: n}
	}
}

func buildSubtreeFromList(it NodeIterator) []SyntaxNode {
	var subtree []SyntaxNode
	for i := 0; i < it.Len(); i++ {
		subtree = append(subtree, buildTree(it.Get(i)))
	}
	return subtree
}

///

type SyntaxNode interface {
	Label() string
	Children() []SyntaxNode
}

type SyntaxTree struct {
	node     ast.Node
	children []SyntaxNode
}

func (n *SyntaxTree) Label() string {
	return fmt.Sprintf("%T", n.node)
}

func (n *SyntaxTree) Children() []SyntaxNode {
	return n.children
}

type SyntaxToken struct {
	node ast.Node
	v    string
}

func (t *SyntaxToken) Label() string {
	return fmt.Sprintf("%T %v", t.node, t.v)
}

func (t *SyntaxToken) Children() []SyntaxNode {
	return nil
}