hyperstone/gen/main.go

package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"io"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"text/template"
)

var (
	messageTypes    = make(map[string]bool)
	enumTypes       = make(map[string]bool)
	entityTypes     = make(map[int]string)
	cmdTypes        = make(map[int]string)
	cmdEnumType     = "EDemoCommands"
	entityEnumTypes = map[string]bool{
		"NET_Messages":        true,
		"SVC_Messages":        true,
		"EBaseUserMessages":   true,
		"EBaseEntityMessages": true,
		"EBaseGameEvents":     true,
		"EDotaUserMessages":   true,
	}
	prefixes = map[string]string{
		"EDemoCommands_DEM_":         "CDemo",
		"NET_Messages_net_":          "CNETMsg_",
		"SVC_Messages_svc_":          "CSVCMsg_",
		"EBaseUserMessages_UM_":      "CUserMessage",
		"EBaseEntityMessages_EM_":    "CEntityMessage",
		"EBaseGameEvents_GE_":        "CMsg",
		"EDotaUserMessages_DOTA_UM_": "CDOTAUserMsg_",
	}
	specials = map[string]string{
		"EDemoCommands_DEM_SignonPacket":              "CDemoPacket",
		"EDotaUserMessages_DOTA_UM_StatsHeroDetails":  "CDOTAUserMsg_StatsHeroMinuteDetails",
		"EDotaUserMessages_DOTA_UM_CombatLogDataHLTV": "CMsgDOTACombatLogEntry",
		"EDotaUserMessages_DOTA_UM_TournamentDrop":    "CMsgGCToClientTournamentItemDrop",
		"EDotaUserMessages_DOTA_UM_MatchMetadata":     "CDOTAClientMsg_MatchMetadata",
	}
	skipped = map[string]bool{
		"EDemoCommands_DEM_IsCompressed":                  true,
		"EDemoCommands_DEM_Max":                           true,
		"EDotaUserMessages_DOTA_UM_MatchDetails":          true,
		"EBaseUserMessages_UM_ParticleManager":            true,
		"EBaseUserMessages_UM_CustomGameEvent":            true,
		"EDotaUserMessages_DOTA_UM_AddUnitToSelection":    true,
		"EBaseUserMessages_UM_HudError":                   true,
		"EDotaUserMessages_DOTA_UM_CombatLogData":         true,
		"EDotaUserMessages_DOTA_UM_CharacterSpeakConcept": true,
	}
	// EBaseUserMessages_UM_HandHapticPulse
	tpl = `package main
                                                                                         
                                                                                         
/*------------------------------------------------------------------------------

                          .aMMMb  .aMMMb  dMMMMb  dMMMMMP
                         dMP"VMP dMP"dMP dMP VMP dMP
                        dMP     dMP dMP dMP dMP dMMMP
                       dMP.aMP dMP.aMP dMP.aMP dMP
                       VMMMP"  VMMMP" dMMMMP" dMMMMMP

      .aMMMMP dMMMMMP dMMMMb  dMMMMMP dMMMMb  .aMMMb dMMMMMMP dMMMMMP dMMMMb
     dMP"    dMP     dMP dMP dMP     dMP.dMP dMP"dMP   dMP   dMP     dMP VMP
    dMP MMP"dMMMP   dMP dMP dMMMP   dMMMMK" dMMMMMP   dMP   dMMMP   dMP dMP
   dMP.dMP dMP     dMP dMP dMP     dMP"AMF dMP dMP   dMP   dMP     dMP.aMP
   VMMMP" dMMMMMP dMP dMP dMMMMMP dMP dMP dMP dMP   dMP   dMMMMMP dMMMMP"


This code was generated by a code-generation program. It was NOT written by
hand. Do not edit this file by hand! Your edits will be destroyed!

This file can be regenerated by running "go generate"

The generator program is defined in "gen/main.go"

------------------------------------------------------------------------------*/

import (
	"github.com/golang/protobuf/proto"
	"github.com/jordanorelli/hyperstone/dota"
)

type protoFactory map[int]func() proto.Message

func (p protoFactory) BuildMessage(id int) proto.Message {
	fn, ok := p[id]
	if !ok {
		return nil
	}
	return fn()
}

var cmdFactory = protoFactory{
{{- range $id, $name := .Commands }}
	{{$id}}: func() proto.Message { return new(dota.{{$name}}) },
{{- end }}
}

var entFactory = protoFactory{
{{- range $id, $name := .Entities }}
	{{$id}}: func() proto.Message { return new(dota.{{$name}}) },
{{- end }}
}
`
)

func ensureNewline(t string) string {
	if strings.HasSuffix(t, "\n") {
		return t
	}
	return t + "\n"
}

func bail(status int, t string, args ...interface{}) {
	var out io.Writer
	if status == 0 {
		out = os.Stdout
	} else {
		out = os.Stderr
	}

	fmt.Fprintf(out, ensureNewline(t), args...)
	os.Exit(status)
}

// processes a single value specification
func processValueSpec(spec *ast.ValueSpec) {
	if spec.Type == nil {
		return
	}

	t, ok := spec.Type.(*ast.Ident)
	if !ok {
		return
	}

	var isEntityType bool
	switch {
	case t.Name == cmdEnumType:
		// it's a message type
	case entityEnumTypes[t.Name]:
		isEntityType = true
	default:
		return
	}

	for i, name := range spec.Names {
		if name.Name == "_" {
			continue
		}

		valExpr := spec.Values[i]
		litExpr, ok := valExpr.(*ast.BasicLit)
		if !ok {
			continue
		}

		if litExpr.Kind != token.INT {
			continue
		}

		n, err := strconv.Atoi(litExpr.Value)
		if err != nil {
			continue
		}

		if isEntityType {
			entityTypes[n] = name.Name
		} else {
			cmdTypes[n] = name.Name
		}
	}
}

// processes a single type specification from the Go source
// e.g.:
//       type Thing int32
//       type Message struct {
//           ...
//       }
func processTypeSpec(spec *ast.TypeSpec) {
	switch tt := spec.Type.(type) {
	case *ast.StructType:
		// the only structes that are defined in our generated proto code are
		// protobuf message types
		messageTypes[spec.Name.Name] = true
	case *ast.Ident:
		switch tt.Name {
		case "int32":
			// all protobuf enums generate int32s in Go
			enumTypes[spec.Name.Name] = true
		}
	}
}

// processes one specification from the Go source
func processSpec(spec ast.Spec) {
	switch t := spec.(type) {
	case *ast.ValueSpec:
		processValueSpec(t)
	case *ast.TypeSpec:
		processTypeSpec(t)
	}
}

// processes one protobuf-generated Go declaration
func processDeclaration(decl ast.Decl) {
	d, ok := decl.(*ast.GenDecl)
	if !ok {
		return
	}
	for _, spec := range d.Specs {
		processSpec(spec)
	}
}

// processes one protobuf-generated Go file
func processFile(name string, fi *ast.File) {
	for _, decl := range fi.Decls {
		processDeclaration(decl)
	}
}

// processes a package of protobuf-generated Go
func processPackage(name string, pkg *ast.Package) {
	for name, fi := range pkg.Files {
		processFile(name, fi)
	}
}

// given an enum name, finds the appropriate message type
func typeName(enumName string) string {
	for prefix, replacement := range prefixes {
		if strings.HasPrefix(enumName, prefix) {
			candidate := strings.Replace(enumName, prefix, replacement, 1)
			if messageTypes[candidate] {
				return candidate
			}
		}
	}
	return ""
}

func main() {
	flag.Parse()
	if flag.NArg() != 1 {
		bail(1, "gen should get exactly one argument: the directory to operate on")
	}
	path := flag.Arg(0)

	fs := token.NewFileSet()
	packages, err := parser.ParseDir(fs, path, nil, 0)
	if err != nil {
		bail(1, "go parser error: %v", err)
	}

	for name, pkg := range packages {
		processPackage(name, pkg)
	}

	type typeMap map[int]string

	type context struct {
		Commands typeMap
		Entities typeMap
	}

	ctx := context{make(typeMap), make(typeMap)}

	for id, name := range cmdTypes {
		realName := typeName(name)
		if realName == "" {
			continue
		}
		ctx.Commands[id] = realName
	}
	for id, name := range entityTypes {
		realName := typeName(name)
		if realName == "" {
			continue
		}
		ctx.Entities[id] = realName
	}

	t, err := template.New("out.go").Parse(tpl)
	if err != nil {
		bail(1, "bad template: %v", err)
	}

	buf := bytes.NewBuffer(nil)
	if err := t.Execute(buf, ctx); err != nil {
		bail(1, "template error: %v", err)
	}

	source, err := format.Source(buf.Bytes())
	if err != nil {
		bail(1, "fmt error: %v", err)
	}

	if err := ioutil.WriteFile("generated.go", source, 0644); err != nil {
		bail(1, "error writing source output: %v", err)
	}
}