tile38/pkg/collection/collection.go

package collection

import (
	"math"

	"github.com/tidwall/btree"
	"github.com/tidwall/tile38/pkg/ds"
	"github.com/tidwall/tile38/pkg/geojson"
	"github.com/tidwall/tile38/pkg/index"
)

const (
	idOrdered    = 0
	valueOrdered = 1
)

type itemT struct {
	id     string
	object geojson.Object
}

func (i *itemT) Less(item btree.Item, ctx interface{}) bool {
	switch ctx {
	default:
		return false
	case idOrdered:
		return i.id < item.(*itemT).id
	case valueOrdered:
		i1, i2 := i.object.String(), item.(*itemT).object.String()
		if i1 < i2 {
			return true
		}
		if i1 > i2 {
			return false
		}
		// the values match so we will compare the ids, which are always unique.
		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       ds.BTree     // items sorted by keys
	values      *btree.BTree // items sorted by value+key
	index       *index.Index // items geospatially indexed
	fieldMap    map[string]int
	fieldValues map[string][]float64
	weight      int
	points      int
	objects     int // geometry count
	nobjects    int // non-geometry count
}

var counter uint64

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

func (c *Collection) setFieldValues(id string, values []float64) {
	if c.fieldValues == nil {
		c.fieldValues = make(map[string][]float64)
	}
	c.fieldValues[id] = values
}

func (c *Collection) getFieldValues(id string) (values []float64) {
	if c.fieldValues == nil {
		return nil
	}
	return c.fieldValues[id]
}
func (c *Collection) deleteFieldValues(id string) {
	if c.fieldValues != nil {
		delete(c.fieldValues, id)
	}
}

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

// StringCount returns the number of string values.
func (c *Collection) StringCount() int {
	return c.nobjects
}

// 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
}

// Bounds returns the bounds of all the items in the collection.
func (c *Collection) Bounds() (minX, minY, maxX, maxY float64) {
	return c.index.Bounds()
}

// Set 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) Set(
	id string, obj geojson.Object, fields []string, values []float64,
) (
	oldObject geojson.Object, oldFields []float64, newFields []float64,
) {
	var oldItem *itemT
	newItem := &itemT{id: id, object: obj}
	// add the new item to main btree and remove the old one if needed
	oldItemPtr, _ := c.items.Set(id, newItem)
	if oldItemPtr != nil {
		// the old item was removed, now let's remove from the rtree
		// or strings tree.
		oldItem = oldItemPtr.(*itemT)
		if obj.IsGeometry() {
			// geometry
			c.index.Remove(oldItem)
			c.objects--
		} else {
			// string
			c.values.Delete(oldItem)
			c.nobjects--
		}
		// decrement the point count
		c.points -= oldItem.object.PositionCount()

		// decrement the weights
		c.weight -= len(c.getFieldValues(id)) * 8
		c.weight -= oldItem.object.Weight() + len(oldItem.id)

		// references
		oldObject = oldItem.object
		oldFields = c.getFieldValues(id)
		newFields = oldFields
	}
	// insert the new item into the rtree or strings tree.
	if obj.IsGeometry() {
		c.index.Insert(newItem)
		c.objects++
	} else {
		c.values.ReplaceOrInsert(newItem)
		c.nobjects++
	}
	// increment the point count
	c.points += obj.PositionCount()

	// add the new weights
	c.weight += len(newFields) * 8
	c.weight += obj.Weight() + len(id)
	if fields == nil {
		if len(values) > 0 {
			// directly set the field values, update weight
			c.weight -= len(newFields) * 8
			newFields = values
			c.setFieldValues(id, newFields)
			c.weight += len(newFields) * 8
		}
	} else {
		//if len(fields) == 0 {
		//	panic("if fields is empty, make it nil")
		//}
		// map field name to value
		for i, field := range fields {
			c.setField(newItem, field, values[i])
		}
		newFields = c.getFieldValues(id)
	}
	return oldObject, oldFields, newFields
}

// Delete removes an object and returns it.
// If the object does not exist then the 'ok' return value will be false.
func (c *Collection) Delete(id string) (
	obj geojson.Object, fields []float64, ok bool,
) {
	old, _ := c.items.Delete(id)
	if old == nil {
		return nil, nil, false
	}
	item := old.(*itemT)
	if item.object.IsGeometry() {
		c.index.Remove(item)
		c.objects--
	} else {
		c.values.Delete(item)
		c.nobjects--
	}
	fields = c.getFieldValues(id)
	c.deleteFieldValues(id)
	c.weight -= len(fields) * 8
	c.weight -= item.object.Weight() + len(item.id)
	c.points -= item.object.PositionCount()
	return item.object, 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,
) {
	val, _ := c.items.Get(id)
	if val == nil {
		return nil, nil, false
	}
	item := val.(*itemT)
	return item.object, c.getFieldValues(id), true
}

// 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,
) {
	val, _ := c.items.Get(id)
	if val == nil {
		ok = false
		return
	}
	item := val.(*itemT)
	updated = c.setField(item, field, value)
	return item.object, c.getFieldValues(id), updated, true
}

// SetFields is similar to SetField, just setting multiple fields at once
func (c *Collection) SetFields(
	id string, inFields []string, inValues []float64,
) (
	obj geojson.Object, fields []float64, updatedCount int, ok bool,
) {
	val, _ := c.items.Get(id)
	if val == nil {
		ok = false
		return
	}
	item := val.(*itemT)
	for idx, field := range inFields {
		if c.setField(item, field, inValues[idx]) {
			updatedCount++
		}
	}
	return item.object, c.getFieldValues(id), updatedCount, 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
	}
	fields := c.getFieldValues(item.id)
	c.weight -= len(fields) * 8
	for idx >= len(fields) {
		fields = append(fields, 0)
	}
	c.weight += len(fields) * 8
	ovalue := fields[idx]
	fields[idx] = value
	c.setFieldValues(item.id, fields)
	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 ids.
func (c *Collection) Scan(desc bool,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	iter := func(key string, value interface{}) bool {
		iitm := value.(*itemT)
		keepon = iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id))
		return keepon
	}
	if desc {
		c.items.Reverse(iter)
	} else {
		c.items.Scan(iter)
	}
	return keepon
}

// ScanRange iterates though the collection starting with specified id.
func (c *Collection) ScanRange(start, end string, desc bool,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	iter := func(key string, value interface{}) bool {
		if !desc {
			if key >= end {
				return false
			}
		} else {
			if key <= end {
				return false
			}
		}
		iitm := value.(*itemT)
		keepon = iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id))
		return keepon
	}

	if desc {
		c.items.Descend(start, iter)
	} else {
		c.items.Ascend(start, iter)
	}
	return keepon
}

// SearchValues iterates though the collection values.
func (c *Collection) SearchValues(desc bool,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	iter := func(item btree.Item) bool {
		iitm := item.(*itemT)
		keepon = iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id))
		return keepon
	}
	if desc {
		c.values.Descend(iter)
	} else {
		c.values.Ascend(iter)
	}
	return keepon
}

// SearchValuesRange iterates though the collection values.
func (c *Collection) SearchValuesRange(start, end string, desc bool,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	iter := func(item btree.Item) bool {
		iitm := item.(*itemT)
		keepon = iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id))
		return keepon
	}
	if desc {
		c.values.DescendRange(&itemT{object: geojson.String(start)},
			&itemT{object: geojson.String(end)}, iter)
	} else {
		c.values.AscendRange(&itemT{object: geojson.String(start)},
			&itemT{object: geojson.String(end)}, iter)
	}
	return keepon
}

// ScanGreaterOrEqual iterates though the collection starting with specified id.
func (c *Collection) ScanGreaterOrEqual(id string, desc bool,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	iter := func(key string, value interface{}) bool {
		iitm := value.(*itemT)
		keepon = iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id))
		return keepon
	}
	if desc {
		c.items.Descend(id, iter)
	} else {
		c.items.Ascend(id, iter)
	}
	return keepon
}

func (c *Collection) geoSearch(
	bbox geojson.BBox,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	return c.index.Search(bbox.Min.X, bbox.Min.Y, bbox.Max.X, bbox.Max.Y,
		func(item interface{}) bool {
			iitm := item.(*itemT)
			if !iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id)) {
				return false
			}
			return true
		},
	)
}

// Nearby returns all object that are nearby a point.
func (c *Collection) Nearby(
	sparse uint8, lat, lon, meters, minZ, maxZ float64,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	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 {
			bbox.Min.Z, bbox.Max.Z = minZ, maxZ
			keepon = c.geoSearch(bbox,
				func(id string, obj geojson.Object, fields []float64) bool {
					if obj.Nearby(center, meters) {
						if iterator(id, obj, fields) {
							return false
						}
					}
					return true
				},
			)
			if !keepon {
				break
			}
		}
		return keepon
	}
	bbox.Min.Z, bbox.Max.Z = minZ, maxZ
	return c.geoSearch(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(
	sparse uint8, obj geojson.Object,
	minLat, minLon, maxLat, maxLon, lat, lon, meters, minZ, maxZ float64,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	var keepon = true
	var bbox geojson.BBox
	center := geojson.Position{X: lon, Y: lat, Z: 0}

	if obj != nil {
		bbox = obj.CalculatedBBox()
		if minZ == math.Inf(-1) && maxZ == math.Inf(+1) {
			if bbox.Min.Z == 0 && bbox.Max.Z == 0 {
				bbox.Min.Z = minZ
				bbox.Max.Z = maxZ
			}
		}
	} else if meters != -1 {
		bbox = geojson.BBoxesFromCenter(lat, lon, meters)
	} else {
		bbox = geojson.BBox{
			Min: geojson.Position{X: minLon, Y: minLat, Z: minZ},
			Max: geojson.Position{X: maxLon, Y: maxLat, Z: maxZ},
		}
	}
	bboxes := bbox.Sparse(sparse)
	if sparse > 0 {
		for _, bbox := range bboxes {
			if obj != nil {
				keepon = c.geoSearch(bbox,
					func(id string, o geojson.Object, fields []float64) bool {
						if o.Within(obj) {
							if iterator(id, o, fields) {
								return false
							}
						}
						return true
					},
				)
			} else if meters != -1 {
				keepon = c.geoSearch(bbox,
					func(id string, o geojson.Object, fields []float64) bool {
						if o.WithinCircle(center, meters) {
							if iterator(id, o, fields) {
								return false
							}
						}
						return true
					},
				)
			}
			if keepon {
				keepon = c.geoSearch(bbox,
					func(id string, o geojson.Object, fields []float64) bool {
						if o.WithinBBox(bbox) {
							if iterator(id, o, fields) {
								return false
							}
						}
						return true
					},
				)
			}
			if !keepon {
				break
			}
		}
		return keepon
	}
	if obj != nil {
		return c.geoSearch(bbox,
			func(id string, o geojson.Object, fields []float64) bool {
				if o.Within(obj) {
					return iterator(id, o, fields)
				}
				return true
			},
		)
	} else if meters != -1 {
		return c.geoSearch(bbox,
			func(id string, o geojson.Object, fields []float64) bool {
				if o.WithinCircle(center, meters) {
					return iterator(id, o, fields)
				}
				return true
			},
		)
	}
	return c.geoSearch(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(
	sparse uint8, obj geojson.Object,
	minLat, minLon, maxLat, maxLon, lat, lon, meters, minZ, maxZ float64,
	doClip bool,
	iterator func(
		id string, obj geojson.Object, fields []float64, clipBox geojson.BBox,
	) bool,
) bool {

	var keepon = true
	var clipbox, bbox geojson.BBox
	center := geojson.Position{X: lon, Y: lat, Z: 0}
	if obj != nil {
		bbox = obj.CalculatedBBox()
		if minZ == math.Inf(-1) && maxZ == math.Inf(+1) {
			if bbox.Min.Z == 0 && bbox.Max.Z == 0 {
				bbox.Min.Z = minZ
				bbox.Max.Z = maxZ
			}
		}
	} else if meters != -1 {
		bbox = geojson.BBoxesFromCenter(lat, lon, meters)
	} else {
		bbox = geojson.BBox{
			Min: geojson.Position{X: minLon, Y: minLat, Z: minZ},
			Max: geojson.Position{X: maxLon, Y: maxLat, Z: maxZ},
		}
		if doClip {
			clipbox = bbox
		}
	}
	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: minZ},
					Max: geojson.Position{X: x + xpart, Y: y + ypart, Z: maxZ},
				})
			}
		}
		for _, bbox := range bboxes {
			if obj != nil {
				keepon = c.geoSearch(bbox,
					func(id string, o geojson.Object, fields []float64) bool {
						if o.Intersects(obj) {
							if iterator(id, o, fields, clipbox) {
								return false
							}
						}
						return true
					},
				)
			} else if meters != -1 {
				keepon = c.geoSearch(bbox,
					func(id string, o geojson.Object, fields []float64) bool {
						if o.IntersectsCircle(center, meters) {
							if iterator(id, o, fields, clipbox) {
								return false
							}
						}
						return true
					},
				)
			}
			if keepon {
				keepon = c.geoSearch(bbox,
					func(id string, o geojson.Object, fields []float64) bool {
						if o.IntersectsBBox(bbox) {
							if iterator(id, o, fields, clipbox) {
								return false
							}
						}
						return true
					},
				)
			}
			if !keepon {
				break
			}
		}
		return keepon
	}
	if obj != nil {
		return c.geoSearch(bbox,
			func(id string, o geojson.Object, fields []float64) bool {
				if o.Intersects(obj) {
					return iterator(id, o, fields, clipbox)
				}
				return true
			},
		)
	} else if meters != -1 {
		return c.geoSearch(bbox,
			func(id string, o geojson.Object, fields []float64) bool {
				if o.IntersectsCircle(center, meters) {
					return iterator(id, o, fields, clipbox)
				}
				return true
			},
		)
	}
	return c.geoSearch(bbox,
		func(id string, o geojson.Object, fields []float64) bool {
			if o.IntersectsBBox(bbox) {
				return iterator(id, o, fields, clipbox)
			}
			return true
		},
	)
}

// NearestNeighbors returns the nearest neighbors
func (c *Collection) NearestNeighbors(
	lat, lon float64,
	iterator func(id string, obj geojson.Object, fields []float64) bool,
) bool {
	return c.index.KNN(lon, lat, func(item interface{}) bool {
		var iitm *itemT
		iitm, ok := item.(*itemT)
		if !ok {
			return true // just ignore
		}
		if !iterator(iitm.id, iitm.object, c.getFieldValues(iitm.id)) {
			return false
		}
		return true
	})
}