gosmonaut.go

package gosmonaut

import (
	"errors"
	"fmt"
	"io"
	"runtime"
	"sync"
	"time"
)

/* Gosmonaut */

type osmPair struct {
	i OSMEntity
	e error
}

// Config defines the configuration for Gosmonaut.
type Config struct {
	// DebugMode prints duration and memory info and runs the garbage collector
	// after every processing step if enabled.
	DebugMode bool

	// PrintWarnings prints warnings to stdout if enabled. Possible warnings
	// include missing referenced entites or unsupported features.
	PrintWarnings bool

	// Set the number of processes that are used for decoding.
	// If not set the amount of available logical CPUs will be used.
	NumProcessors int

	// Decoder sets the PBF blob decoder. Defaults to `FastDecoder`.
	Decoder DecoderType
}

// Gosmonaut is responsible for decoding an OpenStreetMap pbf file.
// For creating an instance the NewGosmonaut() function must be used.
type Gosmonaut struct {
	dec    *decoder
	stream chan osmPair
	lock   sync.Mutex

	// Store header block
	header Header

	// Defined by caller
	types            OSMTypeSet
	funcEntityNeeded func(OSMType, OSMTags) bool

	// ID trackers
	nodeIDTracker, wayIDTracker idTracker

	// Entitiy caches
	nodeCache *binaryNodeEntityMap
	wayCache  *binaryWayEntityMap
	nodeTags  map[int64]OSMTags
	wayTags   map[int64]OSMTags

	// For debug mode
	debugMode     bool
	printWarnings bool
	timeStarted   time.Time
	timeLast      time.Time
}

// NewGosmonaut creates a new Gosmonaut instance and parses the meta information
// (Header) of the given file. Either zero or exactly one `Config` object must
// be passed.
func NewGosmonaut(file io.ReadSeeker, config ...Config) (*Gosmonaut, error) {
	// Check config
	var conf Config
	if len(config) > 1 {
		return nil, errors.New("only 1 Config object is allowed")
	} else if len(config) == 1 {
		conf = config[0]
	}

	// Get number of processes
	var nProcs int
	if conf.NumProcessors < 1 {
		nProcs = runtime.NumCPU()
	} else {
		nProcs = conf.NumProcessors
	}

	// Create decoder
	dec, header, err := newDecoder(file, nProcs, conf.Decoder)
	if err != nil {
		return nil, err
	}

	return &Gosmonaut{
		dec:           dec,
		header:        header,
		debugMode:     conf.DebugMode,
		printWarnings: conf.PrintWarnings,
	}, nil
}

// Header returns the meta information of the PBF file.
func (g *Gosmonaut) Header() Header {
	return g.header
}

// Typical PrimitiveBlock contains 8k OSM entities
const entitiesPerPrimitiveBlock = 8000

// Start starts the decoding process. The function call will block until the
// previous run has finished.
// Only types that are enabled in `types` will be sent to the caller.
// funcEntityNeeded will be called to determine if the caller needs a specific
// OSM entity.
// Found entities and encountered errors can be received by polling the Next()
// method.
func (g *Gosmonaut) Start(
	types OSMTypeSet,
	funcEntityNeeded func(OSMType, OSMTags) bool,
) {
	// Block until previous run finished
	g.lock.Lock()
	g.stream = make(chan osmPair, entitiesPerPrimitiveBlock)

	// Init vars
	g.funcEntityNeeded = funcEntityNeeded
	g.types = types

	go func() {
		// Decode
		g.decode()

		// Finish
		close(g.stream)
		g.lock.Unlock()
	}()
}

func (g *Gosmonaut) decode() {
	// Init debug vars
	{
		timeNow := time.Now()
		g.timeStarted = timeNow
		g.timeLast = timeNow
	}
	g.printDebugInfo("Decoding started")

	// Scan relation dependencies
	g.nodeIDTracker = newBitsetIDTracker()
	g.wayIDTracker = newBitsetIDTracker()

	if g.types.Get(RelationType) {
		if err := g.scanRelationDependencies(); err != nil {
			g.streamError(err)
			return
		}
		g.printDebugInfo(fmt.Sprintf("Scanned relation dependencies [length: %d]", g.wayIDTracker.len()))
	}

	// Scan way dependencies
	if g.types.Get(WayType) || g.wayIDTracker.len() != 0 {
		if err := g.scanWayDependencies(); err != nil {
			g.streamError(err)
			return
		}
		g.printDebugInfo(fmt.Sprintf("Scanned way dependencies [length: %d]", g.nodeIDTracker.len()))
	}

	g.nodeCache = newBinaryNodeEntityMap(g.nodeIDTracker.len())
	g.nodeTags = make(map[int64]OSMTags)
	g.printDebugInfo("Created node cache")

	// Scan nodes
	if g.types.Get(NodeType) || g.nodeIDTracker.len() != 0 {
		if err := g.scanNodes(); err != nil {
			g.streamError(err)
			return
		}
		g.printDebugInfo("Scanned nodes")
	}

	g.nodeIDTracker = nil
	g.printDebugInfo("Deleted node ID tracker")

	g.nodeCache.prepare()
	g.printDebugInfo("Prepared node cache")

	g.wayCache = newBinaryWayEntityMap(g.wayIDTracker.len())
	g.wayTags = make(map[int64]OSMTags)
	g.printDebugInfo("Created way cache")

	// Scan ways
	if g.types.Get(WayType) || g.wayIDTracker.len() != 0 {
		if err := g.scanWays(); err != nil {
			g.streamError(err)
			return
		}
		g.printDebugInfo("Scanned ways")
	}

	g.wayIDTracker = nil
	g.printDebugInfo("Deleted way ID tracker")

	g.wayCache.prepare()
	g.printDebugInfo("Prepared way cache")

	// Scan relations
	if g.types.Get(RelationType) {
		if err := g.scanRelations(); err != nil {
			g.streamError(err)
			return
		}
		g.printDebugInfo("Scanned relations")
	}

	g.wayCache = nil
	g.nodeCache = nil
	g.wayTags = nil
	g.nodeTags = nil
	g.printDebugInfo("Deleted entity caches")

	if g.debugMode {
		fmt.Println("Elapsed time:", time.Since(g.timeStarted))
	}
}

func (g *Gosmonaut) streamError(err error) {
	g.stream <- osmPair{nil, err}
}

func (g *Gosmonaut) streamEntity(i OSMEntity) {
	g.stream <- osmPair{i, nil}
}

// Next returns the next decoded entity (x)or an error.
// If the error is io.EOF the file has successfully been decoded.
// If the error is not EOF decoding has been stopped due to another error.
func (g *Gosmonaut) Next() (OSMEntity, error) {
	p, ok := <-g.stream
	if !ok {
		return nil, io.EOF
	}
	return p.i, p.e
}

func (g *Gosmonaut) entityNeeded(t OSMType, tags OSMTags) bool {
	if !g.types.Get(t) {
		return false
	}
	return g.funcEntityNeeded(t, tags)
}

func (g *Gosmonaut) scanRelationDependencies() error {
	return g.scan(RelationType, func(id int64, tags OSMTags, v entityParser) error {
		// Needed by cache?
		if !g.entityNeeded(RelationType, tags) {
			return nil
		}

		// Get parser
		d, ok := v.(relationParser)
		if !ok {
			return fmt.Errorf("got invalid relation parser (%T)", v)
		}

		// Add members to ID caches
		ids, err := d.ids()
		if err != nil {
			return err
		}
		types, err := d.types()
		if err != nil {
			return err
		}
		if len(ids) != len(types) {
			return errors.New("length of relation ids and types differs")
		}
		for i, id := range ids {
			switch types[i] {
			case WayType:
				g.wayIDTracker.set(id)
			case NodeType:
				g.nodeIDTracker.set(id)
				// We don't support sub-relations yet
			}
		}
		return nil
	})
}

func (g *Gosmonaut) scanWayDependencies() error {
	return g.scan(WayType, func(id int64, tags OSMTags, v entityParser) error {
		// Needed by cache or stream?
		if !g.wayIDTracker.get(id) && !g.entityNeeded(WayType, tags) {
			return nil
		}

		// Get parser
		d, ok := v.(wayParser)
		if !ok {
			return fmt.Errorf("got invalid way parser (%T)", v)
		}

		// Add nodes to ID cache
		refs, err := d.refs()
		if err != nil {
			return err
		}
		for _, id := range refs {
			g.nodeIDTracker.set(id)
		}
		return nil
	})
}

func (g *Gosmonaut) scanNodes() error {
	return g.scan(NodeType, func(id int64, tags OSMTags, v entityParser) error {
		// Needed by cache or stream?
		rc, rs := g.nodeIDTracker.get(id), g.entityNeeded(NodeType, tags)
		if !rc && !rs {
			return nil
		}

		// Get parser
		d, ok := v.(nodeParser)
		if !ok {
			return fmt.Errorf("got invalid node parser (%T)", v)
		}

		// Get properties
		lat, lon, err := d.coords()
		if err != nil {
			return err
		}

		// Add to cache
		if rc {
			g.nodeCache.add(rawNode{
				id:  id,
				lat: newCoordinate(lat),
				lon: newCoordinate(lon),
			})

			// Add tags
			if tags.Len() != 0 {
				g.nodeTags[id] = tags
			}
		}

		// Send to output stream
		if rs {
			g.streamEntity(Node{
				ID:   id,
				Lat:  lat,
				Lon:  lon,
				Tags: tags,
			})
		}
		return nil
	})
}

func (g *Gosmonaut) scanWays() error {
	return g.scan(WayType, func(id int64, tags OSMTags, v entityParser) error {
		// Needed by cache or stream?
		rc, rs := g.wayIDTracker.get(id), g.entityNeeded(WayType, tags)
		if !rc && !rs {
			return nil
		}

		// Get parser
		d, ok := v.(wayParser)
		if !ok {
			return fmt.Errorf("got invalid way parser (%T)", v)
		}

		// Get properties
		refs, err := d.refs()
		if err != nil {
			return err
		}

		// Add to cache
		if rc {
			g.wayCache.add(rawWay{
				id:   id,
				refs: refs,
			})

			// Add tags
			if tags.Len() != 0 {
				g.wayTags[id] = tags
			}
		}

		// Send to output stream
		if rs {
			nodes, err := g.uncacheNodes(refs)
			if err != nil {
				return err
			}
			g.streamEntity(Way{
				ID:    id,
				Tags:  tags,
				Nodes: nodes,
			})
		}
		return nil
	})
}

func (g *Gosmonaut) scanRelations() error {
	return g.scan(RelationType, func(id int64, tags OSMTags, v entityParser) error {
		// Needed by stream?
		if !g.entityNeeded(RelationType, tags) {
			return nil
		}

		// Get parser
		d, ok := v.(relationParser)
		if !ok {
			return fmt.Errorf("got invalid relation parser (%T)", v)
		}

		// Get properties
		ids, err := d.ids()
		if err != nil {
			return err
		}
		types, err := d.types()
		if err != nil {
			return err
		}
		roles, err := d.roles()
		if err != nil {
			return err
		}
		n := len(ids)
		if len(types) != n || len(roles) != n {
			return errors.New("length of relation ids, roles and types differs")
		}

		// Build relation
		members := make([]Member, 0, n)
		for i, mid := range ids {
			var e OSMEntity
			switch types[i] {
			case WayType:
				if rw, ok := g.wayCache.get(mid); ok {
					nodes, err := g.uncacheNodes(rw.refs)
					if err != nil {
						return err
					}
					e = Way{
						ID:    rw.id,
						Tags:  g.wayTags[rw.id],
						Nodes: nodes,
					}
				} else {
					g.printWarning(fmt.Sprintf("Way #%d in not in file for relation #%d", mid, id))
					continue
				}
			case NodeType:
				if n, err := g.uncacheNode(mid); err == nil {
					e = n
				} else {
					g.printWarning(err.Error())
					continue
				}
			default:
				// We don't support sub-relations yet
				g.printWarning(fmt.Sprintf("Skipping sub-relation #%d in relation #%d (not supported)", mid, id))
				continue
			}
			members = append(members, Member{roles[i], e})
		}

		// Send to output stream
		g.streamEntity(Relation{
			ID:      id,
			Tags:    tags,
			Members: members,
		})
		return nil
	})
}

func (g *Gosmonaut) scan(t OSMType,
	receiver func(int64, OSMTags, entityParser) error,
) error {
	if err := g.dec.Start(t); err != nil {
		return err
	}

	// Decode file
	for {
		parsers, err := g.dec.nextPair()
		if err == io.EOF {
			return nil
		} else if err != nil {
			return err
		}

		// Iterate parsers
		for _, v := range parsers {
			// Iterate entities
			for {
				id, tags, err := v.next()
				if err == io.EOF {
					break
				} else if err != nil {
					return err
				}

				// Send to receiver
				if err := receiver(id, tags, v); err != nil {
					return err
				}
			}
		}
	}
}

/* Cache */

func (g *Gosmonaut) uncacheNode(id int64) (Node, error) {
	rn, ok := g.nodeCache.get(id)
	if !ok {
		return Node{}, fmt.Errorf("Node #%d in not in file", id)
	}
	return Node{
		ID:   rn.id,
		Lat:  rn.lat.float(),
		Lon:  rn.lon.float(),
		Tags: g.nodeTags[id],
	}, nil
}

func (g *Gosmonaut) uncacheNodes(refs []int64) ([]Node, error) {
	nodes := make([]Node, len(refs))
	for i, id := range refs {
		if n, err := g.uncacheNode(id); err == nil {
			nodes[i] = n
		} else {
			return nil, err
		}
	}
	return nodes, nil
}

/* Debug Mode */

func (g *Gosmonaut) printWarning(warning string) {
	if g.printWarnings {
		fmt.Println("Warning:", warning)
	}
}

func (g *Gosmonaut) printDebugInfo(state string) {
	if !g.debugMode {
		return
	}
	elapsed := time.Since(g.timeLast).Seconds()

	// Run garbage collector
	runtime.GC()

	// Print memory stats
	bToMb := func(b uint64) uint64 {
		return b / 1024 / 1024
	}
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	fmt.Printf("Alloc = %v MiB\tTotalAlloc = %v MiB\tSys = %v MiB\tNumGC = %v\n",
		bToMb(m.Alloc),
		bToMb(m.TotalAlloc),
		bToMb(m.Sys),
		m.NumGC,
	)

	// Print elapsed
	fmt.Printf("%.4fs - %v\n", elapsed, state)
	g.timeLast = time.Now()
}