package index import ( "math" "github.com/tidwall/tile38/index/rtree" ) // Item represents an index item. type Item interface { Point() (x, y, z float64) Rect() (minX, minY, minZ, maxX, maxY, maxZ float64) } // FlexItem can represent a point or a rectangle type FlexItem struct { MinX, MinY, MinZ, MaxX, MaxY, MaxZ float64 } // Rect returns the rectangle func (item *FlexItem) Rect() (minX, minY, minZ, maxX, maxY, maxZ float64) { return item.MinX, item.MinY, item.MinZ, item.MaxX, item.MaxY, item.MaxZ } // Point returns the point func (item *FlexItem) Point() (x, y, z float64) { return item.MinX, item.MinY, item.MinZ } // Index is a geospatial index type Index struct { r *rtree.RTree nr map[*rtree.Rect]Item // normalized points nrr map[Item][]*rtree.Rect // normalized points mulm map[Item]bool // store items that contain multiple rects } // New create a new index func New() *Index { return &Index{ r: rtree.New(), mulm: make(map[Item]bool), nr: make(map[*rtree.Rect]Item), nrr: make(map[Item][]*rtree.Rect), } } // Insert inserts an item into the index func (ix *Index) Insert(item Item) { minX, minY, minZ, maxX, maxY, maxZ := item.Rect() if minX == maxX && minY == maxY { x, y, normd := normPoint(minY, minX) if normd { nitem := &rtree.Rect{MinX: x, MinY: y, MinZ: minZ, MaxX: x, MaxY: y, MaxZ: maxZ} ix.nr[nitem] = item ix.nrr[item] = []*rtree.Rect{nitem} ix.r.Insert(nitem) } else { ix.r.Insert(item) } } else { mins, maxs, normd := normRect(minY, minX, maxY, maxX) if normd { var nitems []*rtree.Rect for i := range mins { minX, minY, maxX, maxY := mins[i][0], mins[i][1], maxs[i][0], maxs[i][1] nitem := &rtree.Rect{MinX: minX, MinY: minY, MinZ: minZ, MaxX: maxX, MaxY: maxY, MaxZ: maxZ} ix.nr[nitem] = item nitems = append(nitems, nitem) ix.r.Insert(nitem) } ix.nrr[item] = nitems if len(mins) > 1 { ix.mulm[item] = true } } else { ix.r.Insert(item) } } return } // Remove removed an item from the index func (ix *Index) Remove(item Item) { if nitems, ok := ix.nrr[item]; ok { for _, nitem := range nitems { ix.r.Remove(nitem) delete(ix.nr, nitem) } delete(ix.nrr, item) } else { ix.r.Remove(item) } } // Count counts all items in the index. func (ix *Index) Count() int { count := 0 ix.Search(-90, -180, 90, 180, math.Inf(-1), math.Inf(+1), func(item Item) bool { count++ return true }) return count } // Bounds returns the minimum bounding rectangle of all items in the index. func (ix *Index) Bounds() (MinX, MinY, MinZ, MaxX, MaxY, MaxZ float64) { return ix.r.Bounds() } // RemoveAll removes all items from the index. func (ix *Index) RemoveAll() { ix.r.RemoveAll() } func (ix *Index) getRTreeItem(item rtree.Item) Item { switch item := item.(type) { case Item: return item case *rtree.Rect: return ix.nr[item] } return nil } func (ix *Index) NearestNeighbors(lat, lon float64, iterator func(item Item) bool) bool { x, y, _ := normPoint(lat, lon) return ix.r.NearestNeighbors(x, y, 0, func(item rtree.Item, dist float64) bool { iitm := ix.getRTreeItem(item) if item == nil { return true } return iterator(iitm) }) } // Search returns all items that intersect the bounding box. func (ix *Index) Search(swLat, swLon, neLat, neLon, minZ, maxZ float64, iterator func(item Item) bool) bool { var keepon = true var idm = make(map[Item]bool) mins, maxs, _ := normRect(swLat, swLon, neLat, neLon) // Points if len(mins) == 1 { // There is only one rectangle. // It's possible that a r rect may span multiple entries. Check mulm map for spanning rects. if keepon { ix.r.Search(mins[0][0], mins[0][1], minZ, maxs[0][0], maxs[0][1], maxZ, func(item rtree.Item) bool { iitm := ix.getRTreeItem(item) if iitm != nil { if ix.mulm[iitm] { if !idm[iitm] { idm[iitm] = true keepon = iterator(iitm) } } else { keepon = iterator(iitm) } } return keepon }) } } else { // There are multiple rectangles. Duplicates might occur. for i := range mins { if keepon { ix.r.Search(mins[i][0], mins[i][1], minZ, maxs[i][0], maxs[i][1], maxZ, func(item rtree.Item) bool { iitm := ix.getRTreeItem(item) if iitm != nil { if ix.mulm[iitm] { if !idm[iitm] { idm[iitm] = true keepon = iterator(iitm) } } else { keepon = iterator(iitm) } } return keepon }) } } } return keepon }