persist/cmd/pgen/gen.go

package main

import (
	"log"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"golang.org/x/tools/go/ast/astutil"
	"os"
	"regexp"
)

type subspec struct {
	r *regexp.Regexp
	s string
}

var (
	firstone bool
	gfset *token.FileSet
	fspec, tspec, nspec subspec
	ntspecs []subspec
	nodes *ast.File
)

type gvisit struct{}

func subs(dst *string, specs []subspec) {
	for _,spec := range(specs) {
		sub(dst,spec)
	}
}

func sub(dst *string,spec subspec) {
	if spec.r.MatchString(*dst) {
		*dst = spec.r.ReplaceAllString(*dst,spec.s)
	}
}

func (v *gvisit) Visit(node ast.Node) ast.Visitor {
	switch n := node.(type) {
	case *ast.Ident:
		if n.Obj != nil && n.Obj.Kind == ast.Typ {
			subs(&n.Name,ntspecs)
			subs(&n.Obj.Name,ntspecs)
		}
	case *ast.TypeSwitchStmt:
		for _,s := range(n.Body.List) {
			cs, ok := s.(*ast.CaseClause); if !ok { continue }
			for _,c := range(cs.List) {
				ci, ok := c.(*ast.Ident); if !ok { continue }
				subs(&ci.Name,ntspecs)
			}
		}
	case *ast.TypeAssertExpr:
		nt, ok := n.Type.(*ast.Ident); if !ok { return v }
		subs(&nt.Name,ntspecs)
	case *ast.Ellipsis:
		nt, ok := n.Elt.(*ast.Ident); if !ok { return v }
		subs(&nt.Name,ntspecs)
	case *ast.Field:
		nt, ok := n.Type.(*ast.Ident); if !ok { return v }
		subs(&nt.Name,ntspecs)
	case *ast.StarExpr:
		nt, ok := n.X.(*ast.Ident); if !ok { return v }
		subs(&nt.Name,ntspecs)
	case *ast.TypeSpec:
		//ast.Print(gfset,n)
	case *ast.ValueSpec:
		nt, ok := n.Type.(*ast.Ident); if !ok { return v }
		subs(&nt.Name,ntspecs)
	case *ast.ArrayType:
		nt, ok := n.Elt.(*ast.Ident); if !ok { return v }
		subs(&nt.Name,ntspecs)
	case *ast.MapType:
		nt, ok := n.Key.(*ast.Ident); if ok {
			subs(&nt.Name,ntspecs)
		}
		nt, ok = n.Value.(*ast.Ident); if ok {
			subs(&nt.Name,ntspecs)
		}
	case *ast.FuncDecl:
		sub(&n.Name.Name,fspec)
		//ast.Print(gfset,n)
	}
	return v
}

func init() {
	firstone = true
	err := os.Remove("pgen.go")
	if err != nil && os.IsExist(err) {
		log.Fatal("Removing pgen.go:",err)
	}
}

func add(newName, typName, typ string) {
	fspec = subspec{regexp.MustCompile("New"), newName}
	nspec = subspec{regexp.MustCompile("_N"), typName}
	tspec = subspec{regexp.MustCompile("_T"), typ}
	ntspecs = []subspec{nspec,tspec}

	gfset = token.NewFileSet()
	f, err := parser.ParseFile(gfset, "", template, 0)
	if err != nil {
		log.Fatal("Parsing persist/template.T:",err)
	}
	var v = &gvisit{}
	ast.Walk(v,f)

	if firstone {
		nodes = f
	} else {
		for _,decl := range(f.Decls) {
			imp, ok := decl.(*ast.GenDecl); if ok {
				if imp.Tok == token.IMPORT {
					continue // skip imports
				}
			}
			nodes.Decls = append(nodes.Decls,decl)
		}
	}
	firstone = false
}

func newImportSpec(name string) *ast.ImportSpec {
	path := &ast.BasicLit{ Value: name }
	ret := &ast.ImportSpec{ Path: path }
	return ret
}

func gen(wantImps []string) {
	of, err := os.Create("pgen.go")
	if err != nil {
		log.Fatal("Cannot open pgen.go")
	}

	qexp := regexp.MustCompile(`"`)
	for _,v := range wantImps {
		v = qexp.ReplaceAllString(v,"")
		astutil.AddImport(gfset, nodes, v)
	}

	err = format.Node(of,gfset,nodes)
	if err != nil {
		log.Fatal("Generate error:",err)
	}
}