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(typeMap)
	packetTypes     = make(typeMap)
	cmdEnumType     = "EDemoCommands"
	entityEnumTypes = map[string]bool{
		"NET_Messages":        true,
		"SVC_Messages":        true,
		"EBaseUserMessages":   true,
		"EBaseEntityMessages": true,
		"EBaseGameEvents":     true,
		"EDotaUserMessages":   true,
		"ETEProtobufIds":      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{
		"EDotaUserMessages_DOTA_UM_StatsHeroDetails":  "CDOTAUserMsg_StatsHeroMinuteDetails",
		"EDotaUserMessages_DOTA_UM_CombatLogDataHLTV": "CMsgDOTACombatLogEntry",
		"EDotaUserMessages_DOTA_UM_TournamentDrop":    "CMsgGCToClientTournamentItemDrop",
		"EDotaUserMessages_DOTA_UM_MatchMetadata":     "CDOTAClientMsg_MatchMetadata",
		"ETEProtobufIds_TE_EffectDispatchId":          "CMsgTEEffectDispatch",
		"EDemoCommands_DEM_SignonPacket":              "CDemoPacket",
	}
	pooled = map[string]bool{
		"CDemoPacket":                            true,
		"CMsgDOTACombatLogEntry":                 true,
		"CDOTAUserMsg_ParticleManager":           true,
		"CMsgSosStartSoundEvent":                 true,
		"CDOTAUserMsg_SpectatorPlayerUnitOrders": true,
		"CNETMsg_Tick":                           true,
		"CDOTAUserMsg_SpectatorPlayerClick":      true,
		"CDOTAUserMsg_TE_UnitAnimationEnd":       true,
		"CSVCMsg_PacketEntities":                 true,
		"CDOTAUserMsg_TE_Projectile":             true,
		"CDOTAUserMsg_TE_UnitAnimation":          true,
		"CSVCMsg_UpdateStringTable":              true,
		"CMsgSosSetSoundEventParams":             true,
		"CDOTAUserMsg_UnitEvent":                 true,
		"CDOTAUserMsg_TE_ProjectileLoc":          true,
		"CDOTAUserMsg_OverheadEvent":             true,
		"CDOTAUserMsg_LocationPing":              true,
		"CDOTAUserMsg_TE_DotaBloodImpact":        true,
		"CDOTAUserMsg_SharedCooldown":            true,
		"CDOTAUserMsg_ChatEvent":                 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 (
	"fmt"
	"sync"
	"github.com/golang/protobuf/proto"
	"github.com/jordanorelli/hyperstone/dota"
)

type packetType int32
type entityType int32

const (
	EDemoCommands_DEM_Error packetType = -1
{{- range $id, $spec := .Packets }}
	{{$spec.EnumName}} packetType = {{$id}}
{{- end }}
{{- range $id, $spec := .Entities }}
	{{$spec.EnumName}} entityType = {{$id}}
{{- end }}
)

func (t packetType) String() string {
	switch t {
{{- range $id, $spec := .Packets }}
	case {{$spec.EnumName}}:
		return "{{$spec.EnumName}}"
{{- end }}
	default:
		return fmt.Sprintf("UnknownPacketType_%d", t)
	}
}

func (t entityType) String() string {
	switch t {
{{- range $id, $spec := .Entities }}
	case {{$spec.EnumName}}:
		return "{{$spec.EnumName}}"
{{- end }}
	default:
		return fmt.Sprintf("UnknownEntityType_%d", t)
	}
}

type messageStatus int

const (
	m_Unknown messageStatus = iota
	m_Skipped
)

func (m messageStatus) Error() string {
	switch m {
	case m_Unknown:
		return "unknown message type"
	case m_Skipped:
		return "skipped message type"
	default:
		return "unknown message error"
	}
}

type packetFactory map[packetType]func() proto.Message
type entityFactory map[entityType]func() proto.Message

type messageFactory struct {
	packets packetFactory
	entities entityFactory
}

func (m *messageFactory) BuildPacket(id packetType) (proto.Message, error) {
	fn, ok := m.packets[id]
	if !ok {
		return nil, m_Unknown
	}
	return fn(), nil
}

func (m *messageFactory) BuildEntity(id entityType) (proto.Message, error) {
	fn, ok := m.entities[id]
	if !ok {
		return nil, m_Unknown
	}
	return fn(), nil
}

func (m *messageFactory) Return(msg proto.Message) {
	switch msg.(type) {
	{{- range $name, $ok := .Pooled }}
	case *dota.{{$name}}:
		p_{{$name}}.Put(msg)
	{{- end}}
	}
}

var (
	{{- range $name, $ok := .Pooled }}
	p_{{$name}} = &sync.Pool{New: func() interface{} { return new(dota.{{$name}}) }}
	{{- end }}
)

type packetWhitelist map[packetType]bool
type entityWhitelist map[entityType]bool

var allPackets = packetWhitelist{
	{{- range $id, $spec := .Packets }}
		{{$spec.EnumName}}: true,
	{{- end }}
}

var allEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{$spec.EnumName}}: true,
	{{- end }}
}

var netEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "NET_Messages" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}

var svcEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "SVC_Messages" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}

var baseUserEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "EBaseUserMessages" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}


//     ffffffffffffffff                                        kkkkkkkk
//    f::::::::::::::::f                                       k::::::k
//   f::::::::::::::::::f                                      k::::::k
//   f::::::fffffff:::::f                                      k::::::k
//   f:::::f       ffffffuuuuuu    uuuuuu      cccccccccccccccc k:::::k    kkkkkkk
//   f:::::f             u::::u    u::::u    cc:::::::::::::::c k:::::k   k:::::k
//  f:::::::ffffff       u::::u    u::::u   c:::::::::::::::::c k:::::k  k:::::k
//  f::::::::::::f       u::::u    u::::u  c:::::::cccccc:::::c k:::::k k:::::k
//  f::::::::::::f       u::::u    u::::u  c::::::c     ccccccc k::::::k:::::k
//  f:::::::ffffff       u::::u    u::::u  c:::::c              k:::::::::::k
//   f:::::f             u::::u    u::::u  c:::::c              k:::::::::::k
//   f:::::f             u:::::uuuu:::::u  c::::::c     ccccccc k::::::k:::::k
//  f:::::::f            u:::::::::::::::uuc:::::::cccccc:::::ck::::::k k:::::k
//  f:::::::f             u:::::::::::::::u c:::::::::::::::::ck::::::k  k:::::k
//  f:::::::f              uu::::::::uu:::u  cc:::::::::::::::ck::::::k   k:::::k
//  fffffffff                uuuuuuuu  uuuu    cccccccccccccccckkkkkkkk    kkkkkkk
var entityEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "EBaseEntityMessages" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}

var gameEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "EBaseGameEvents" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}

var dotaUserEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "EDotaUserMessages" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}

var eteEntities = entityWhitelist{
	{{- range $id, $spec := .Entities }}
		{{- if eq $spec.EnumType "ETEProtobufIds" }}
		{{$spec.EnumName}}: true,
		{{- end }}
	{{- end}}
}

var messages = messageFactory{
	packetFactory{
	{{- range $id, $spec := .Packets }}
		{{$spec.EnumName}}: {{$spec.CreateFn}},
	{{- end }}
	},
	entityFactory{
	{{- range $id, $spec := .Entities }}
		{{$spec.EnumName}}: {{$spec.CreateFn}},
	{{- end }}
	},
}
`
)

type messageSpec struct {
	EnumType  string
	EnumName  string
	EnumValue string
	TypeName  string
}

func (m messageSpec) CreateFn() string {
	if m.Pooled() {
		return fmt.Sprintf("func() proto.Message { return %s.Get().(*dota.%s) }", m.PoolName(), m.TypeName)
	}
	return m.New()
}

func (m messageSpec) New() string {
	return fmt.Sprintf("func() proto.Message { return new(dota.%s) }", m.TypeName)
}

func (m messageSpec) Pooled() bool {
	return pooled[m.TypeName]
}

func (m messageSpec) PoolName() string {
	return fmt.Sprintf("p_%s", m.TypeName)
}

func (m messageSpec) IsCmd() bool {
	return m.EnumType == cmdEnumType
}

type typeMap map[int]messageSpec

func (m typeMap) fillTypeNames() {
	for id, spec := range m {
		spec.TypeName = typeName(spec.EnumName)
		if spec.TypeName == "" {
			delete(m, id)
		} else {
			m[id] = spec
		}
	}
}

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] = messageSpec{EnumName: name.Name, EnumType: t.Name}
		} else {
			packetTypes[n] = messageSpec{EnumName: name.Name, EnumType: t.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 {
	if name, ok := specials[enumName]; ok {
		return name
	}
	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)
	}

	packetTypes.fillTypeNames()
	entityTypes.fillTypeNames()

	var ctx = struct {
		Packets  typeMap
		Entities typeMap
		Pooled   map[string]bool
	}{
		Packets:  packetTypes,
		Entities: entityTypes,
		Pooled:   pooled,
	}

	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 {
		fmt.Println(string(buf.Bytes()))
		bail(1, "fmt error: %v", err)
	}

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