diff --git a/bit/buf_reader.go b/bit/buf_reader.go index 806fec0..20a40d1 100644 --- a/bit/buf_reader.go +++ b/bit/buf_reader.go @@ -4,14 +4,14 @@ import ( "io" ) -type bufReader struct { +type BufReader struct { src []byte // source of data n uint64 // bit buffer bits uint // number of valid bits in n err error // stored error } -func (r *bufReader) ReadBits(bits uint) (n uint64) { +func (r *BufReader) ReadBits(bits uint) (n uint64) { for bits > r.bits { if len(r.src) == 0 { r.err = io.EOF @@ -27,7 +27,7 @@ func (r *bufReader) ReadBits(bits uint) (n uint64) { return } -func (r *bufReader) ReadByte() byte { +func (r *BufReader) ReadByte() byte { if r.bits == 0 { if len(r.src) == 0 { r.err = io.EOF @@ -40,7 +40,7 @@ func (r *bufReader) ReadByte() byte { return byte(r.ReadBits(8)) } -func (r *bufReader) Read(buf []byte) int { +func (r *BufReader) Read(buf []byte) int { if r.bits == 0 { if len(r.src) < len(buf) { r.err = io.EOF @@ -61,7 +61,7 @@ func (r *bufReader) Read(buf []byte) int { return len(buf) } -func (r *bufReader) DiscardBytes(n int) { +func (r *BufReader) DiscardBytes(n int) { if r.bits == 0 { if len(r.src) < n { r.err = io.EOF @@ -85,4 +85,11 @@ func (r *bufReader) DiscardBytes(n int) { r.src = r.src[n:] } -func (r *bufReader) Err() error { return r.err } +func (r *BufReader) SetSource(b []byte) { + r.src = b + r.bits = 0 + r.err = nil + r.n = 0 +} + +func (r *BufReader) Err() error { return r.err } diff --git a/bit/reader.go b/bit/reader.go index 89c4366..d122d44 100644 --- a/bit/reader.go +++ b/bit/reader.go @@ -15,16 +15,16 @@ type Reader interface { } // NewReader creates a new bit.Reader for any arbitrary reader. -func NewReader(r io.Reader) Reader { +func NewReader(r io.Reader) *StreamReader { br, ok := r.(io.ByteReader) if !ok { br = bufio.NewReader(r) } - return &streamReader{src: br} + return &StreamReader{src: br} } // NewByteReader creates a bit.Reader for a static slice of bytes. It's just // using a bytes.Reader internally. -func NewBytesReader(b []byte) Reader { - return &bufReader{src: b} +func NewBytesReader(b []byte) *BufReader { + return &BufReader{src: b} } diff --git a/bit/stream_reader.go b/bit/stream_reader.go index 3351db6..9dc71b9 100644 --- a/bit/stream_reader.go +++ b/bit/stream_reader.go @@ -7,7 +7,7 @@ import ( // bit.Reader allows for bit-level reading of arbitrary source data. This is // based on the bit reader found in the standard library's bzip2 package. // https://golang.org/src/compress/bzip2/bit_reader.go -type streamReader struct { +type StreamReader struct { src io.ByteReader // source of data n uint64 // bit buffer bits uint // number of valid bits in n @@ -16,7 +16,7 @@ type streamReader struct { // ReadBits reads the given number of bits and returns them in the // least-significant part of a uint64. -func (r *streamReader) ReadBits(bits uint) (n uint64) { +func (r *StreamReader) ReadBits(bits uint) (n uint64) { if r.err != nil { return 0 } @@ -37,7 +37,7 @@ func (r *streamReader) ReadBits(bits uint) (n uint64) { } // ReadByte reads a single byte, regardless of alignment. -func (r *streamReader) ReadByte() byte { +func (r *StreamReader) ReadByte() byte { if r.bits == 0 { b, err := r.src.ReadByte() if err != nil { @@ -50,7 +50,7 @@ func (r *streamReader) ReadByte() byte { } // Read reads like an io.Reader, taking care of alignment internally. -func (r *streamReader) Read(buf []byte) int { +func (r *StreamReader) Read(buf []byte) int { for i := 0; i < len(buf); i++ { b := r.ReadByte() if r.err != nil { @@ -62,10 +62,10 @@ func (r *streamReader) Read(buf []byte) int { } // discards N byte of data on the reader or until EOF -func (r *streamReader) DiscardBytes(n int) { +func (r *StreamReader) DiscardBytes(n int) { for i := 0; i < n; i++ { r.ReadByte() } } -func (r *streamReader) Err() error { return r.err } +func (r *StreamReader) Err() error { return r.err } diff --git a/gen/main.go b/gen/main.go index 1402468..c90c50b 100644 --- a/gen/main.go +++ b/gen/main.go @@ -65,6 +65,8 @@ var ( "CDOTAUserMsg_UnitEvent": true, "CDOTAUserMsg_TE_ProjectileLoc": true, "CDOTAUserMsg_OverheadEvent": true, + "CDOTAUserMsg_LocationPing": true, + "CDOTAUserMsg_TE_DotaBloodImpact": true, } // EBaseUserMessages_UM_HandHapticPulse tpl = `package main diff --git a/generated.go b/generated.go index 98dafec..47f8330 100644 --- a/generated.go +++ b/generated.go @@ -690,6 +690,8 @@ func (m *messageFactory) BuildEntity(id entityType) (proto.Message, error) { func (m *messageFactory) Return(msg proto.Message) { switch msg.(type) { + case *dota.CDOTAUserMsg_LocationPing: + p_CDOTAUserMsg_LocationPing.Put(msg) case *dota.CDOTAUserMsg_OverheadEvent: p_CDOTAUserMsg_OverheadEvent.Put(msg) case *dota.CDOTAUserMsg_ParticleManager: @@ -698,6 +700,8 @@ func (m *messageFactory) Return(msg proto.Message) { p_CDOTAUserMsg_SpectatorPlayerClick.Put(msg) case *dota.CDOTAUserMsg_SpectatorPlayerUnitOrders: p_CDOTAUserMsg_SpectatorPlayerUnitOrders.Put(msg) + case *dota.CDOTAUserMsg_TE_DotaBloodImpact: + p_CDOTAUserMsg_TE_DotaBloodImpact.Put(msg) case *dota.CDOTAUserMsg_TE_Projectile: p_CDOTAUserMsg_TE_Projectile.Put(msg) case *dota.CDOTAUserMsg_TE_ProjectileLoc: @@ -726,10 +730,12 @@ func (m *messageFactory) Return(msg proto.Message) { } var ( + p_CDOTAUserMsg_LocationPing = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_LocationPing) }} p_CDOTAUserMsg_OverheadEvent = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_OverheadEvent) }} p_CDOTAUserMsg_ParticleManager = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_ParticleManager) }} p_CDOTAUserMsg_SpectatorPlayerClick = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_SpectatorPlayerClick) }} p_CDOTAUserMsg_SpectatorPlayerUnitOrders = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_SpectatorPlayerUnitOrders) }} + p_CDOTAUserMsg_TE_DotaBloodImpact = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_TE_DotaBloodImpact) }} p_CDOTAUserMsg_TE_Projectile = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_TE_Projectile) }} p_CDOTAUserMsg_TE_ProjectileLoc = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_TE_ProjectileLoc) }} p_CDOTAUserMsg_TE_UnitAnimation = &sync.Pool{New: func() interface{} { return new(dota.CDOTAUserMsg_TE_UnitAnimation) }} @@ -1291,7 +1297,7 @@ var messages = messageFactory{ EDotaUserMessages_DOTA_UM_GlobalLightColor: func() proto.Message { return new(dota.CDOTAUserMsg_GlobalLightColor) }, EDotaUserMessages_DOTA_UM_GlobalLightDirection: func() proto.Message { return new(dota.CDOTAUserMsg_GlobalLightDirection) }, EDotaUserMessages_DOTA_UM_InvalidCommand: func() proto.Message { return new(dota.CDOTAUserMsg_InvalidCommand) }, - EDotaUserMessages_DOTA_UM_LocationPing: func() proto.Message { return new(dota.CDOTAUserMsg_LocationPing) }, + EDotaUserMessages_DOTA_UM_LocationPing: func() proto.Message { return p_CDOTAUserMsg_LocationPing.Get().(*dota.CDOTAUserMsg_LocationPing) }, EDotaUserMessages_DOTA_UM_MapLine: func() proto.Message { return new(dota.CDOTAUserMsg_MapLine) }, EDotaUserMessages_DOTA_UM_MiniKillCamInfo: func() proto.Message { return new(dota.CDOTAUserMsg_MiniKillCamInfo) }, EDotaUserMessages_DOTA_UM_MinimapDebugPoint: func() proto.Message { return new(dota.CDOTAUserMsg_MinimapDebugPoint) }, @@ -1337,7 +1343,9 @@ var messages = messageFactory{ EDotaUserMessages_DOTA_UM_TE_ProjectileLoc: func() proto.Message { return p_CDOTAUserMsg_TE_ProjectileLoc.Get().(*dota.CDOTAUserMsg_TE_ProjectileLoc) }, - EDotaUserMessages_DOTA_UM_TE_DotaBloodImpact: func() proto.Message { return new(dota.CDOTAUserMsg_TE_DotaBloodImpact) }, + EDotaUserMessages_DOTA_UM_TE_DotaBloodImpact: func() proto.Message { + return p_CDOTAUserMsg_TE_DotaBloodImpact.Get().(*dota.CDOTAUserMsg_TE_DotaBloodImpact) + }, EDotaUserMessages_DOTA_UM_TE_UnitAnimation: func() proto.Message { return p_CDOTAUserMsg_TE_UnitAnimation.Get().(*dota.CDOTAUserMsg_TE_UnitAnimation) }, diff --git a/packet.go b/packet.go index 79c1d99..1db410d 100644 --- a/packet.go +++ b/packet.go @@ -2,11 +2,16 @@ package main import ( "fmt" + "sync" "github.com/golang/protobuf/proto" "github.com/golang/snappy" ) +var ( + p_decode_bufs = &sync.Pool{New: func() interface{} { return make([]byte, 1<<18) }} +) + // packet represents the top-level envelope in the dota replay format. All // data in the replay file is packed into packets of at most 65kb. type packet struct { @@ -31,7 +36,9 @@ func (p *packet) Open(m *messageFactory, pbuf *proto.Buffer) (proto.Message, err } if p.compressed { - buf, err := snappy.Decode(nil, p.body[:p.size]) + buf := p_decode_bufs.Get().([]byte) + defer p_decode_bufs.Put(buf) + buf, err := snappy.Decode(buf, p.body[:p.size]) if err != nil { return nil, wrap(err, "packet open failed snappy decode") } diff --git a/parser.go b/parser.go index ca59c13..ab5d686 100644 --- a/parser.go +++ b/parser.go @@ -23,6 +23,8 @@ type parser struct { pwl packetWhitelist packets *sync.Pool + + br *bit.BufReader } func newParser(r io.Reader) *parser { @@ -37,6 +39,7 @@ func newParser(r io.Reader) *parser { return &packet{body: make([]byte, 1<<16)} }, }, + br: new(bit.BufReader), } } @@ -77,19 +80,19 @@ func (p *parser) run(out chan maybe) { } func (p *parser) emitChildren(pkt *dota.CDemoPacket, c chan maybe) { - br := bit.NewBytesReader(pkt.GetData()) + p.br.SetSource(pkt.GetData()) for { - t := entityType(bit.ReadUBitVar(br)) - s := bit.ReadVarInt(br) + t := entityType(bit.ReadUBitVar(p.br)) + s := bit.ReadVarInt(p.br) if p.ewl[t] { - br.Read(p.scratch[:s]) + p.br.Read(p.scratch[:s]) } else { - br.DiscardBytes(int(s)) + p.br.DiscardBytes(int(s)) continue } - switch err := br.Err(); err { + switch err := p.br.Err(); err { case nil: break case io.EOF: