Generate additional types for our entity system

main.go

package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/build"
	"go/parser"
	"go/token"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

type ComponentSet struct {
	FileName    string
	PackageName string
	Components  []Component
}

type Component struct {
	ComponentName       string
	ComponentBufferName string
}

type Generator struct {
	buf  bytes.Buffer
	vars map[string]string
}

func (g *Generator) w(s string) {
	g.buf.WriteString(s)
}

func (g *Generator) wf(format string, a ...interface{}) {
	g.buf.WriteString(fmt.Sprintf(format, a...))
}

var (
	replPattern = regexp.MustCompile("<([a-zA-Z0-9]+)>")
)

func (g *Generator) InstantiateTemplate(template string) {
	g.buf.WriteString(replPattern.ReplaceAllStringFunc(template, func(k string) string {
		if g.vars != nil {
			k = k[1 : len(k)-1] // drop angle brackets
			if v, ok := g.vars[k]; ok {
				return v
			}
		}
		return k
	}))
}

func (g *Generator) SetTemplateVar(k, v string) {
	if g.vars == nil {
		g.vars = make(map[string]string)
	}
	g.vars[k] = v
}

func main() {
	var css []ComponentSet

	pkg, err := build.Default.ImportDir(".", 0)
	if err != nil {
		log.Fatalf("cannot process directory %s: %s", ".", err)
	}
	fs := token.NewFileSet()
	for _, fn := range pkg.GoFiles {
		if strings.HasSuffix(fn, "_es.go") {
			continue
		}

		var cs ComponentSet

		cs.FileName = fn

		parsedFile, err := parser.ParseFile(fs, fn, nil, 0)
		if err != nil {
			log.Fatalf("parsing package: %s: %s", fn, err)
		}

		cs.PackageName = parsedFile.Name.Name

		ast.Inspect(parsedFile, func(n ast.Node) bool {
			switch x := n.(type) {
			case *ast.TypeSpec:
				typeName := x.Name.Name
				if strings.HasSuffix(typeName, "Component") {
					cs.Components = append(cs.Components, Component{
						ComponentName:       typeName,
						ComponentBufferName: strings.ToLower(typeName[:1]) + typeName[1:] + "Buffer", // package private
					})
				}
			}
			return true
		})

		if len(cs.Components) > 0 {
			css = append(css, cs)
		}
	}

	// generate the code
	var packageName string

	for _, cs := range css {
		var g Generator

		if len(packageName) > 0 && packageName != cs.PackageName {
			panic("different packages")
		}

		packageName = cs.PackageName

		g.wf("// Code generated by \"escgen %v\"; DO NOT EDIT\n", strings.Join(os.Args[1:], " "))
		g.w("\n")
		g.wf("package %s\n", cs.PackageName)
		g.w("\n")
		g.w("import \"github.com/leidegre/gendjinn-go/game/engine/es\"\n")

		for _, c := range cs.Components {
			g.SetTemplateVar("BufferTypeName", c.ComponentBufferName)

			g.w("\n")
			g.wf("func (*%v) ComponentType() es.ComponentType { return %vType }\n", c.ComponentName, c.ComponentName)

			g.w("\n")
			g.wf("type %v struct {	components []%v; ptrs []*%v; }\n", c.ComponentBufferName, c.ComponentName, c.ComponentName)
			g.w("\n")

			g.wf("func (arr *%v) Init(cap int) {\n", c.ComponentBufferName)
			g.wf("components, ptrs := make([]%v, cap), make([]*%v, cap)\n", c.ComponentName, c.ComponentName)
			g.w("for i := 0; i < cap; i++ { ptrs[i] = &components[i] }\n")
			g.w("arr.components, arr.ptrs = components, ptrs\n")
			g.w("}\n")

			g.wf("func (arr *%v) Len() int { return len(arr.components) }\n", c.ComponentBufferName)
			g.wf("func (arr *%v) Get(i int) es.Component { return &arr.components[i] }\n", c.ComponentBufferName)

			g.InstantiateTemplate(`
func (buf *<BufferTypeName>) Build(i, j int) uint16 {
	n := i
	for ; i < j; i++ {
		c := &buf.components[i]
		if c.TestFlag(es.ComponentFlagAllocated) {
			buf.ptrs[n] = c
			n++
		}
	}
	maxAllocComponentIndex := buf.ptrs[n-1].ComponentIndex()
	return maxAllocComponentIndex
}
`)

			g.wf("func (arr *%v) GetMany(n int) interface{} { return arr.ptrs[:n] }\n", c.ComponentBufferName)
		}

		bn := strings.TrimSuffix(filepath.Base(cs.FileName), filepath.Ext(cs.FileName))

		ioutil.WriteFile(filepath.Join(filepath.Dir(cs.FileName), fmt.Sprintf("%v_es.go", bn)), g.buf.Bytes(), 0644)
	}

	// generate the common stuff
	{
		var g Generator

		g.wf("// Code generated by \"escgen %v\"; DO NOT EDIT\n", strings.Join(os.Args[1:], " "))
		g.w("\n")
		g.wf("package %s\n", packageName)
		g.w("\n")
		g.w("import \"github.com/leidegre/gendjinn-go/game/engine/es\"\n")

		g.w("\n")
		g.w("const (\n")
		for _, cs := range css {
			for _, c := range cs.Components {
				g.wf("%vType es.ComponentType = iota + 1\n", c.ComponentName)
			}
		}
		g.w(")\n")

		g.w("\n")
		g.w("func init() {\n")
		for _, cs := range css {
			for _, c := range cs.Components {
				g.wf("es.RegisterComponent(%vType, &%v{}, 64)\n", c.ComponentName, c.ComponentBufferName)
			}
		}
		g.w("}\n")

		ioutil.WriteFile("init_es.go", g.buf.Bytes(), 0644)
	}
}