tile38/controller/collection/collection.go

package collection

import (
	"math"

	"github.com/google/btree"
	"github.com/tidwall/tile38/geojson"
	"github.com/tidwall/tile38/index"
)

type itemT struct {
	ID     string
	Object geojson.Object
	Fields []float64
}

func (i *itemT) Less(item btree.Item) bool {
	return i.ID < item.(*itemT).ID
}

func (i *itemT) Rect() (minX, minY, maxX, maxY float64) {
	bbox := i.Object.CalculatedBBox()
	return bbox.Min.X, bbox.Min.Y, bbox.Max.X, bbox.Max.Y
}

func (i *itemT) Point() (x, y float64) {
	x, y, _, _ = i.Rect()
	return
}

// Collection represents a collection of geojson objects.
type Collection struct {
	items    *btree.BTree
	index    *index.Index
	fieldMap map[string]int
	weight   int
	points   int
	objects  int
}

var counter uint64

// New creates an empty collection
func New() *Collection {
	col := &Collection{
		index:    index.New(),
		items:    btree.New(16),
		fieldMap: make(map[string]int),
	}
	return col
}

// Count returns the number of objects in collection.
func (c *Collection) Count() int {
	return c.objects
}

// PointCount returns the number of points (lat/lon coordinates) in collection.
func (c *Collection) PointCount() int {
	return c.points
}

// TotalWeight calculates the in-memory cost of the collection in bytes.
func (c *Collection) TotalWeight() int {
	return c.weight + c.overheadWeight()
}

func (c *Collection) overheadWeight() int {
	// the field map.
	mapweight := 0
	for field := range c.fieldMap {
		mapweight += len(field) + 8 // key + value
	}
	mapweight = int((float64(mapweight) * 1.05) + 28.0) // about an 8% pad plus golang 28 byte map overhead.
	// the btree. each object takes up 64bits for the interface head for each item.
	btreeweight := (c.objects * 8)
	// plus roughly one pointer for every item
	btreeweight += (c.objects * 8)
	// also the btree header weight
	btreeweight += 24
	return mapweight + btreeweight
}

// ReplaceOrInsert adds or replaces an object in the collection and returns the fields array.
// If an item with the same id is already in the collection then the new item will adopt the old item's fields.
// The fields argument is optional.
// The return values are the old object, the old fields, and the new fields
func (c *Collection) ReplaceOrInsert(id string, obj geojson.Object, fields []string, values []float64) (oldObject geojson.Object, oldFields []float64, newFields []float64) {
	oldItem, ok := c.remove(id)
	nitem := c.insert(id, obj)
	if ok {
		oldObject = oldItem.Object
		oldFields = oldItem.Fields
		nitem.Fields = oldFields
		c.weight += len(nitem.Fields) * 8
	}
	for i, field := range fields {
		c.setField(nitem, field, values[i])
	}
	return oldObject, oldFields, nitem.Fields
}

func (c *Collection) remove(id string) (item *itemT, ok bool) {
	i := c.items.Delete(&itemT{ID: id})
	if i == nil {
		return nil, false
	}
	item = i.(*itemT)
	c.index.Remove(item)
	c.weight -= len(item.Fields) * 8
	c.weight -= item.Object.Weight() + len(item.ID)
	c.points -= item.Object.PositionCount()
	c.objects--
	return item, true
}

func (c *Collection) insert(id string, obj geojson.Object) (item *itemT) {
	item = &itemT{ID: id, Object: obj}
	c.index.Insert(item)
	c.items.ReplaceOrInsert(item)
	c.weight += obj.Weight() + len(id)
	c.points += obj.PositionCount()
	c.objects++
	return item
}

// Remove removes an object and returns it.
// If the object does not exist then the 'ok' return value will be false.
func (c *Collection) Remove(id string) (obj geojson.Object, fields []float64, ok bool) {
	item, ok := c.remove(id)
	if !ok {
		return nil, nil, false
	}
	return item.Object, item.Fields, true
}

func (c *Collection) get(id string) (obj geojson.Object, fields []float64, ok bool) {
	i := c.items.Get(&itemT{ID: id})
	if i == nil {
		return nil, nil, false
	}
	item := i.(*itemT)
	return item.Object, item.Fields, true
}

// Get returns an object.
// If the object does not exist then the 'ok' return value will be false.
func (c *Collection) Get(id string) (obj geojson.Object, fields []float64, ok bool) {
	return c.get(id)
}

// SetField set a field value for an object and returns that object.
// If the object does not exist then the 'ok' return value will be false.
func (c *Collection) SetField(id, field string, value float64) (obj geojson.Object, fields []float64, updated bool, ok bool) {
	i := c.items.Get(&itemT{ID: id})
	if i == nil {
		ok = false
		return
	}
	item := i.(*itemT)
	updated = c.setField(item, field, value)
	return item.Object, item.Fields, updated, true
}

func (c *Collection) setField(item *itemT, field string, value float64) (updated bool) {
	idx, ok := c.fieldMap[field]
	if !ok {
		idx = len(c.fieldMap)
		c.fieldMap[field] = idx
	}
	c.weight -= len(item.Fields) * 8
	for idx >= len(item.Fields) {
		item.Fields = append(item.Fields, math.NaN())
	}
	c.weight += len(item.Fields) * 8
	ovalue := item.Fields[idx]
	if math.IsNaN(ovalue) {
		ovalue = 0
	}
	item.Fields[idx] = value
	return ovalue != value
}

// FieldMap return a maps of the field names.
func (c *Collection) FieldMap() map[string]int {
	return c.fieldMap
}

// FieldArr return an array representation of the field names.
func (c *Collection) FieldArr() []string {
	arr := make([]string, len(c.fieldMap))
	for field, i := range c.fieldMap {
		arr[i] = field
	}
	return arr
}

// Scan iterates though the collection. A cursor can be used for paging.
func (c *Collection) Scan(cursor uint64, iterator func(id string, obj geojson.Object, fields []float64) bool) (ncursor uint64) {
	var i uint64
	var active = true
	c.items.Ascend(func(item btree.Item) bool {
		if i >= cursor {
			iitm := item.(*itemT)
			active = iterator(iitm.ID, iitm.Object, iitm.Fields)
		}
		i++
		return active
	})
	return i
}

// ScanGreaterOrEqual iterates though the collection starting with specified id. A cursor can be used for paging.
func (c *Collection) ScanGreaterOrEqual(id string, cursor uint64, iterator func(id string, obj geojson.Object, fields []float64) bool) (ncursor uint64) {
	var i uint64
	var active = true
	c.items.AscendGreaterOrEqual(&itemT{ID: id}, func(item btree.Item) bool {
		if i >= cursor {
			iitm := item.(*itemT)
			active = iterator(iitm.ID, iitm.Object, iitm.Fields)
		}
		i++
		return active
	})
	return i
}

func (c *Collection) search(cursor uint64, bbox geojson.BBox, iterator func(id string, obj geojson.Object, fields []float64) bool) (ncursor uint64) {
	return c.index.Search(cursor, bbox.Min.Y, bbox.Min.X, bbox.Max.Y, bbox.Max.X, func(item index.Item) bool {
		var iitm *itemT
		iitm, ok := item.(*itemT)
		if !ok {
			return true // just ignore
		}
		if !iterator(iitm.ID, iitm.Object, iitm.Fields) {
			return false
		}
		return true
	})
}

// Nearby returns all object that are nearby a point.
func (c *Collection) Nearby(cursor uint64, sparse uint8, lat, lon, meters float64, iterator func(id string, obj geojson.Object, fields []float64) bool) (ncursor uint64) {
	center := geojson.Position{X: lon, Y: lat, Z: 0}
	bbox := geojson.BBoxesFromCenter(lat, lon, meters)
	bboxes := bbox.Sparse(sparse)
	if sparse > 0 {
		for _, bbox := range bboxes {
			c.search(cursor, bbox, func(id string, obj geojson.Object, fields []float64) bool {
				if obj.Nearby(center, meters) {
					if iterator(id, obj, fields) {
						return false
					}
				}
				return true
			})
		}
		return 0
	}
	return c.search(cursor, bbox, func(id string, obj geojson.Object, fields []float64) bool {
		if obj.Nearby(center, meters) {
			return iterator(id, obj, fields)
		}
		return true
	})
}

// Within returns all object that are fully contained within an object or bounding box. Set obj to nil in order to use the bounding box.
func (c *Collection) Within(cursor uint64, sparse uint8, obj geojson.Object, minLat, minLon, maxLat, maxLon float64, iterator func(id string, obj geojson.Object, fields []float64) bool) (ncursor uint64) {
	var bbox geojson.BBox
	if obj != nil {
		bbox = obj.CalculatedBBox()
	} else {
		bbox = geojson.BBox{Min: geojson.Position{X: minLon, Y: minLat, Z: 0}, Max: geojson.Position{X: maxLon, Y: maxLat, Z: 0}}
	}
	bboxes := bbox.Sparse(sparse)
	if sparse > 0 {
		for _, bbox := range bboxes {
			if obj != nil {
				c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
					if o.Within(obj) {
						if iterator(id, o, fields) {
							return false
						}
					}
					return true
				})
			}
			c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
				if o.WithinBBox(bbox) {
					if iterator(id, o, fields) {
						return false
					}
				}
				return true
			})
		}
		return 0
	}
	if obj != nil {
		return c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
			if o.Within(obj) {
				return iterator(id, o, fields)
			}
			return true
		})
	}
	return c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
		if o.WithinBBox(bbox) {
			return iterator(id, o, fields)
		}
		return true
	})
}

// Intersects returns all object that are intersect an object or bounding box. Set obj to nil in order to use the bounding box.
func (c *Collection) Intersects(cursor uint64, sparse uint8, obj geojson.Object, minLat, minLon, maxLat, maxLon float64, iterator func(id string, obj geojson.Object, fields []float64) bool) (ncursor uint64) {
	var bbox geojson.BBox
	if obj != nil {
		bbox = obj.CalculatedBBox()
	} else {
		bbox = geojson.BBox{Min: geojson.Position{X: minLon, Y: minLat, Z: 0}, Max: geojson.Position{X: maxLon, Y: maxLat, Z: 0}}
	}
	var bboxes []geojson.BBox
	if sparse > 0 {
		split := 1 << sparse
		xpart := (bbox.Max.X - bbox.Min.X) / float64(split)
		ypart := (bbox.Max.Y - bbox.Min.Y) / float64(split)
		for y := bbox.Min.Y; y < bbox.Max.Y; y += ypart {
			for x := bbox.Min.X; x < bbox.Max.X; x += xpart {
				bboxes = append(bboxes, geojson.BBox{
					Min: geojson.Position{X: x, Y: y, Z: 0},
					Max: geojson.Position{X: x + xpart, Y: y + ypart, Z: 0},
				})
			}
		}
		for _, bbox := range bboxes {
			if obj != nil {
				c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
					if o.Intersects(obj) {
						if iterator(id, o, fields) {
							return false
						}
					}
					return true
				})
			}
			c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
				if o.IntersectsBBox(bbox) {
					if iterator(id, o, fields) {
						return false
					}
				}
				return true
			})
		}
		return 0
	}
	if obj != nil {
		return c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
			if o.Intersects(obj) {
				return iterator(id, o, fields)
			}
			return true
		})
	}
	return c.search(cursor, bbox, func(id string, o geojson.Object, fields []float64) bool {
		if o.IntersectsBBox(bbox) {
			return iterator(id, o, fields)
		}
		return true
	})
}