fast_blob_decoder.go

package gosmonaut

import (
	"errors"
	"io"

	"github.com/MorbZ/gosmonaut/OSMPBF"
	"google.golang.org/protobuf/encoding/protowire"
)

const (
	protoPBStringTable    = 1
	protoPBPrimitiveGroup = 2
	protoPBGranularity    = 17
	protoPBLatOffset      = 19
	protoPBLonOffset      = 20

	protoSTString = 1

	protoPGNode       = 1
	protoPGDenseNodes = 2
	protoPGWay        = 3
	protoPGRelation   = 4

	protoDNID       = 1
	protoDNLat      = 8
	protoDNLon      = 9
	protoDNKeysVals = 10

	protoEntityID   = 1
	protoEntityKeys = 2
	protoEntityVals = 3

	// protoNID   = 1
	// protoNKeys = 2
	// protoNVals = 3
	protoNLat = 8
	protoNLon = 9

	// protoWID   = 1
	// protoWKeys = 2
	// protoWVals = 3
	protoWRefs = 8

	// protoRID       = 1
	// protoRKeys     = 2
	// protoRVals     = 3
	protoRRolesSid = 8
	protoRMemIDs   = 9
	protoRTypes    = 10

	protoMTNode     = 0
	protoMTWay      = 1
	protoMTRelation = 2
)

// fastBlobDecoder is a custom implementation of the protobuf format. It is
// optimized for decoding of PBF files. Rather than unmarshalling it streams
// the entities and thus reduces GC overhead. The fast blob decoder lacks
// support of some protobuf features which include groups and unpacked varint
// arrays. It is supposed to fail when it encounters a feature it doesn't
// support.
type fastBlobDecoder struct {
	q                                 []entityParser
	types                             OSMTypeSet
	granularity, latOffset, lonOffset int64
}

func newFastBlobDecoder() *fastBlobDecoder {
	return &fastBlobDecoder{
		granularity: 100,
	}
}

func (dec *fastBlobDecoder) decode(blob *OSMPBF.Blob, t OSMType) ([]entityParser, OSMTypeSet, error) {
	data, err := getBlobData(blob, nil)
	if err != nil {
		return nil, 0, err
	}

	// Decode primitive block
	var pgBufs [][]byte
	var stBuf []byte
	p := newProtoBuffer(data)
	for {
		m, err := p.consumeProtoMessage()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, 0, err
		}

		switch m.wire {
		case protowire.BytesType:
			switch m.tag {
			case protoPBStringTable:
				stBuf = m.buf
			case protoPBPrimitiveGroup:
				pgBufs = append(pgBufs, m.buf)
			}
		case protowire.VarintType:
			switch m.tag {
			case protoPBGranularity:
				dec.granularity = int64(m.i)
			case protoPBLatOffset:
				dec.latOffset = int64(m.i)
			case protoPBLonOffset:
				dec.lonOffset = int64(m.i)
			}
		}
	}

	// Decode primitive groups
	dec.q = nil
	var st stringTable
	for _, buf := range pgBufs {
		if err := dec.decodePrimitiveGroup(buf, &st, t); err != nil {
			return nil, 0, err
		}
	}

	// Decode string table
	if len(dec.q) > 0 {
		if st, err = dec.decodeStringTable(stBuf); err != nil {
			return nil, 0, err
		}
	}

	return dec.q, dec.types, nil
}

func (dec *fastBlobDecoder) decodePrimitiveGroup(buf []byte, st *stringTable, t OSMType) error {
	p := newProtoBuffer(buf)

	// Read first tag to evaluate the entity type
	m, err := p.consumeProtoMessage()
	if err != nil {
		return err
	}
	if m.wire == protowire.BytesType {
		var parser entityParser
		switch m.tag {
		case protoPGDenseNodes:
			if t == NodeType {
				// We use the nested buffer here since there is only 1 dense
				// nodes message.
				if d, err := newFastDenseNodesParser(
					st, m.buf,
					dec.granularity, dec.latOffset, dec.lonOffset,
				); err == nil {
					parser = d
				} else {
					return err
				}
			}
			dec.types.Set(NodeType, true)
		case protoPGNode:
			if t == NodeType {
				parser = newFastNodeParser(
					st, buf,
					dec.granularity, dec.latOffset, dec.lonOffset,
				)
			}
			dec.types.Set(NodeType, true)
		case protoPGWay:
			if t == WayType {
				parser = newFastWayParser(st, buf)
			}
			dec.types.Set(WayType, true)
		case protoPGRelation:
			if t == RelationType {
				parser = newFastRelationParser(st, buf)
			}
			dec.types.Set(RelationType, true)
		}
		if parser != nil {
			dec.q = append(dec.q, parser)
		}
	}
	return nil
}

func (dec *fastBlobDecoder) decodeStringTable(buf []byte) (stringTable, error) {
	p := newProtoBuffer(buf)
	var st []string
	for {
		m, err := p.consumeProtoMessage()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		if m.wire == protowire.BytesType && m.tag == protoSTString {
			s := string(m.buf)
			st = append(st, s)
		}
	}
	return stringTable(st), nil
}

/* Dense Node Parser */
type fastDenseNodesParser struct {
	st                                *stringTable
	idBuf, latBuf, lonBuf, keyValsBuf *protoBuffer
	id, lat, lon                      int64
	granularity, latOffset, lonOffset int64
}

func newFastDenseNodesParser(
	st *stringTable,
	buf []byte,
	granularity, latOffset, lonOffset int64,
) (*fastDenseNodesParser, error) {
	d := &fastDenseNodesParser{
		st:          st,
		granularity: granularity,
		latOffset:   latOffset,
		lonOffset:   lonOffset,
	}

	// Collect packed field buffers
	var idBuf, latBuf, lonBuf, keyValsBuf []byte
	p := newProtoBuffer(buf)
	for {
		m, err := p.consumeProtoMessage()
		if err == io.EOF {
			break
		} else if err != nil {
			return nil, err
		}

		if m.wire == protowire.BytesType {
			var dest *[]byte
			switch m.tag {
			case protoDNID:
				dest = &idBuf
			case protoDNLat:
				dest = &latBuf
			case protoDNLon:
				dest = &lonBuf
			case protoDNKeysVals:
				dest = &keyValsBuf
			}
			if dest != nil {
				*dest = m.buf
			}
		}
	}

	// Create proto buffers
	d.idBuf = newProtoBuffer(idBuf)
	d.latBuf = newProtoBuffer(latBuf)
	d.lonBuf = newProtoBuffer(lonBuf)
	d.keyValsBuf = newProtoBuffer(keyValsBuf)

	return d, nil
}

func (d *fastDenseNodesParser) next() (id int64, tags OSMTags, err error) {
	// ID
	did, err := d.idBuf.consumeVarint()
	if err != nil {
		err = io.EOF
		return
	}
	d.id = protowire.DecodeZigZag(did) + d.id
	id = d.id

	// Lat
	dlat, err := d.latBuf.consumeVarint()
	if err != nil {
		return
	}
	d.lat += protowire.DecodeZigZag(dlat)

	// Lon
	dlon, err := d.lonBuf.consumeVarint()
	if err != nil {
		return
	}
	d.lon += protowire.DecodeZigZag(dlon)

	// Tags
	tags, err = d.nextTags()
	return
}

func (d *fastDenseNodesParser) coords() (lat, lon float64, err error) {
	lat = decodeCoord(d.latOffset, d.granularity, d.lat)
	lon = decodeCoord(d.lonOffset, d.granularity, d.lon)
	return
}

func (d *fastDenseNodesParser) nextTags() (OSMTags, error) {
	var tags OSMTags
	for {
		k, err := d.keyValsBuf.consumeVarint()
		if err != nil || k == 0 {
			break
		}

		v, err := d.keyValsBuf.consumeVarint()
		if err != nil {
			return nil, err
		}

		key, val, err := d.st.extractTag(int(k), int(v))
		if err != nil {
			return nil, err
		}
		tags.Set(key, val)
	}
	return tags, nil
}

/* Node Parser */
type fastNodeParser struct {
	st                                *stringTable
	p                                 *protoBuffer
	granularity, latOffset, lonOffset int64
	lat, lon                          uint64
}

func newFastNodeParser(
	st *stringTable,
	buf []byte,
	granularity, latOffset, lonOffset int64,
) *fastNodeParser {
	return &fastNodeParser{
		st:          st,
		p:           newProtoBuffer(buf),
		granularity: granularity,
		latOffset:   latOffset,
		lonOffset:   lonOffset,
	}
}

func (d *fastNodeParser) next() (id int64, tags OSMTags, err error) {
	var hasLat, hasLon bool
	if id, tags, err = extractEntityBuffered(*d.st, d.p, protoPGNode, func(m protoMessage) {
		if m.wire == protowire.VarintType {
			switch m.tag {
			case protoNLat:
				d.lat = m.i
				hasLat = true
			case protoNLon:
				d.lon = m.i
				hasLon = true
			}
		}
	}); err != nil {
		return
	}

	if !hasLat || !hasLon {
		err = errors.New("Node does not have coordinates")
		return
	}

	id = protowire.DecodeZigZag(uint64(id))
	return
}

func (d *fastNodeParser) coords() (lat, lon float64, err error) {
	lat = decodeCoord(d.latOffset, d.granularity, protowire.DecodeZigZag(d.lat))
	lon = decodeCoord(d.lonOffset, d.granularity, protowire.DecodeZigZag(d.lon))
	return
}

/* Way Parser */
type fastWayParser struct {
	st     *stringTable
	p      *protoBuffer
	refBuf []byte
}

func newFastWayParser(st *stringTable, buf []byte) *fastWayParser {
	return &fastWayParser{
		st: st,
		p:  newProtoBuffer(buf),
	}
}

func (d *fastWayParser) next() (int64, OSMTags, error) {
	return extractEntityBuffered(*d.st, d.p, protoPGWay, func(m protoMessage) {
		if m.wire == protowire.BytesType && m.tag == protoWRefs {
			d.refBuf = m.buf
		}
	})
}

func (d *fastWayParser) refs() ([]int64, error) {
	return extractRefsBuffered(d.refBuf), nil
}

/* Relation Parser */
type fastRelationParser struct {
	st                               *stringTable
	p                                *protoBuffer
	rolesSidBuf, memIDsBuf, typesBuf []byte
}

func newFastRelationParser(st *stringTable, buf []byte) *fastRelationParser {
	return &fastRelationParser{
		st: st,
		p:  newProtoBuffer(buf),
	}
}

func (d *fastRelationParser) next() (int64, OSMTags, error) {
	return extractEntityBuffered(*d.st, d.p, protoPGRelation, func(m protoMessage) {
		if m.wire == protowire.BytesType {
			switch m.tag {
			case protoRRolesSid:
				d.rolesSidBuf = m.buf
			case protoRMemIDs:
				d.memIDsBuf = m.buf
			case protoRTypes:
				d.typesBuf = m.buf
			}
		}
	})
}

func (d *fastRelationParser) ids() ([]int64, error) {
	return extractRefsBuffered(d.memIDsBuf), nil
}

func (d *fastRelationParser) roles() ([]string, error) {
	p := newProtoBuffer(d.rolesSidBuf)
	roles := make([]string, 0, p.countVarints())
	for {
		if sid, err := p.consumeVarint(); err == io.ErrUnexpectedEOF {
			break
		} else if err != nil {
			return nil, err
		} else {
			if s, err := d.st.get(int(sid)); err == nil {
				roles = append(roles, s)
			} else {
				return nil, err
			}
		}
	}
	return roles, nil
}

func (d *fastRelationParser) types() ([]OSMType, error) {
	p := newProtoBuffer(d.typesBuf)
	types := make([]OSMType, 0, p.countVarints())
	for {
		if memberType, err := p.consumeVarint(); err == io.ErrUnexpectedEOF {
			break
		} else if err != nil {
			return nil, err
		} else {
			var t OSMType
			switch memberType {
			case protoMTNode:
				t = NodeType
			case protoMTWay:
				t = WayType
			case protoMTRelation:
				t = RelationType
			default:
				return nil, errors.New("unknown member type")
			}
			types = append(types, t)
		}
	}
	return types, nil
}

/* Helpers */

func extractRefsBuffered(buf []byte) []int64 {
	p := newProtoBuffer(buf)
	ids := make([]int64, 0, p.countVarints())
	var id int64
	for {
		delta, err := p.consumeVarint()
		if err != nil {
			break
		}
		id += protowire.DecodeZigZag(delta)
		ids = append(ids, id)
	}
	return ids
}

func extractEntityBuffered(
	st stringTable,
	p *protoBuffer,
	tag protowire.Number,
	msgs func(protoMessage),
) (id int64, tags OSMTags, err error) {
	m, err := p.consumeProtoMessage()
	if err != nil {
		return
	}
	if m.wire != protowire.BytesType || m.tag != tag {
		err = errors.New("primitive group contains different types")
		return
	}

	var keyBuf, valBuf []byte
	var hasID bool
	p = newProtoBuffer(m.buf)
	for {
		m, err := p.consumeProtoMessage()
		if err == io.EOF {
			break
		} else if err != nil {
			return 0, nil, err
		}

		switch m.wire {
		case protowire.VarintType:
			switch m.tag {
			case protoEntityID:
				hasID = true
				id = int64(m.i)
			default:
				msgs(m)
			}
		case protowire.BytesType:
			switch m.tag {
			case protoEntityKeys:
				keyBuf = m.buf
			case protoEntityVals:
				valBuf = m.buf
			default:
				msgs(m)
			}
		}
	}

	if !hasID {
		err = errors.New("entity does not have an ID")
		return
	}

	// Extract tags
	kp := newProtoBuffer(keyBuf)
	vp := newProtoBuffer(valBuf)
	tags = NewOSMTags(kp.countVarints())
	for {
		var k, v uint64
		k, err = kp.consumeVarint()
		if err == io.ErrUnexpectedEOF {
			err = nil
			break
		} else if err != nil {
			return
		}

		v, err = vp.consumeVarint()
		if err != nil {
			return
		}

		var key, val string
		if key, val, err = st.extractTag(int(k), int(v)); err == nil {
			tags.Set(key, val)
		} else {
			return
		}
	}
	return
}

/* Proto Buffer */

type protoBuffer []byte

func newProtoBuffer(b []byte) *protoBuffer {
	p := protoBuffer(b)
	return &p
}

func (p *protoBuffer) empty() bool {
	return len(*p) == 0
}

func (p *protoBuffer) countVarints() int {
	n := 0
	for _, b := range *p {
		if b&0x80 == 0 {
			n++
		}
	}
	return n
}

func (p *protoBuffer) consume(i int) {
	*p = (*p)[i:]
}

func (p *protoBuffer) consumeTag() (protowire.Number, protowire.Type, error) {
	num, typ, n := protowire.ConsumeTag(*p)
	if n < 0 {
		return 0, 0, protowire.ParseError(n)
	}
	p.consume(n)
	return num, typ, nil
}

func (p *protoBuffer) consumeBytes() ([]byte, error) {
	buf, n := protowire.ConsumeBytes(*p)
	if n < 0 {
		return nil, protowire.ParseError(n)
	}
	p.consume(n)
	return buf, nil
}

func (p *protoBuffer) consumeVarint() (uint64, error) {
	i, n := protowire.ConsumeVarint(*p)
	if n < 0 {
		return 0, protowire.ParseError(n)
	}
	p.consume(n)
	return i, nil
}

func (p *protoBuffer) consumeFieldValue(num protowire.Number, typ protowire.Type) error {
	n := protowire.ConsumeFieldValue(num, typ, *p)
	if n < 0 {
		return protowire.ParseError(n)
	}
	p.consume(n)
	return nil
}

// Depending on the wire type either buf or i contain data
type protoMessage struct {
	tag  protowire.Number
	wire protowire.Type
	buf  []byte
	i    uint64
}

func (p *protoBuffer) consumeProtoMessage() (m protoMessage, err error) {
	for {
		m.tag, m.wire, err = p.consumeTag()
		if err != nil {
			if err == io.ErrUnexpectedEOF {
				err = io.EOF
			}
			return
		}

		switch m.wire {
		// Unneeded fields
		case protowire.Fixed32Type:
			fallthrough
		case protowire.Fixed64Type:
			if err = p.consumeFieldValue(m.tag, m.wire); err != nil {
				return
			}

		// Groups are deprecated and usually not used in PBF files
		case protowire.StartGroupType:
			fallthrough
		case protowire.EndGroupType:
			err = errors.New("protobuf groups are not supported")
			return

		// Valid fields
		case protowire.VarintType:
			m.i, err = p.consumeVarint()
			return
		case protowire.BytesType:
			m.buf, err = p.consumeBytes()
			return
		}
	}
}