cloud-services/pkg/utils/quadkey/quadkey.go

package quadkey

import (
	"math"
	"strings"
)

const (
	MAX_LATITUDE        = 85.05112877980659
	MAX_LONGITUDE       = 180.0
	zoom                = 32
	EARTH_RADIUS        = 6378137.0 //meters
	EARTH_CIRCUMFERENCE = 2.0 * math.Pi * EARTH_RADIUS
	METERS_TO_XY        = 1.0 / EARTH_CIRCUMFERENCE * ((1 << zoom) - 1) //conversion coefficient between meters to XY
	XY_TO_METERS        = EARTH_CIRCUMFERENCE / ((1 << zoom) - 1)       //conversion coefficient from XY to meters
)

// Returns integer (X, Y) in range 0 - (1<<32) where 1<<32 is mapped to MAX LATITUDE, MAX LONGITUDE
// XY (0,0) corresponds to LatLong (90, -180)
func LatLongToXY(lat float64, long float64) (uint64, uint64) {
	lat = math.Min(MAX_LATITUDE, math.Max(-MAX_LATITUDE, lat))
	long = math.Min(MAX_LONGITUDE, math.Max(-MAX_LONGITUDE, long))

	//fx and fy compute fractional x, y in range 0.0-1.0
	fx := long/360.0 + 0.5
	sinlat := math.Sin(lat * math.Pi / 180.0)
	fy := 0.5 - math.Log((1+sinlat)/(1-sinlat))/(4*math.Pi)

	scale := 1 << zoom
	x := math.Min(float64(scale-1), math.Max(0, math.Floor(fx*float64(scale))))
	y := math.Min(float64(scale-1), math.Max(0, math.Floor(fy*float64(scale))))
	return uint64(x), uint64(y)
}

// converts XY to quadkey, where XY is int in range 0 - (1<<32).
// taken from example code https://github.com/CartoDB/python-quadkey/
// this implementation only works with ZOOM=32
func XYToQuadkey(x uint64, y uint64) uint64 {
	B := []uint64{0x5555555555555555, 0x3333333333333333, 0x0F0F0F0F0F0F0F0F, 0x00FF00FF00FF00FF, 0x0000FFFF0000FFFF}
	S := []uint64{1, 2, 4, 8, 16}

	x = (x | (x << S[4])) & B[4]
	y = (y | (y << S[4])) & B[4]

	x = (x | (x << S[3])) & B[3]
	y = (y | (y << S[3])) & B[3]

	x = (x | (x << S[2])) & B[2]
	y = (y | (y << S[2])) & B[2]

	x = (x | (x << S[1])) & B[1]
	y = (y | (y << S[1])) & B[1]

	x = (x | (x << S[0])) & B[0]
	y = (y | (y << S[0])) & B[0]

	return x | (y << 1)
}

func LatLongToQuadKey(lat float64, long float64) uint64 {
	x, y := LatLongToXY(lat, long)
	return XYToQuadkey(x, y)
}

// converts quadkey back to XY, where XY is int in range 0 - (1<<32).
// taken from example code https://github.com/CartoDB/python-quadkey/
func QuadkeyToXY(quadkey uint64) (uint64, uint64) {
	B := []uint64{0x5555555555555555, 0x3333333333333333, 0x0F0F0F0F0F0F0F0F, 0x00FF00FF00FF00FF, 0x0000FFFF0000FFFF, 0x00000000FFFFFFFF}
	S := []uint64{0, 1, 2, 4, 8, 16}
	x := quadkey
	y := quadkey >> 1

	x = (x | (x >> S[0])) & B[0]
	y = (y | (y >> S[0])) & B[0]

	x = (x | (x >> S[1])) & B[1]
	y = (y | (y >> S[1])) & B[1]

	x = (x | (x >> S[2])) & B[2]
	y = (y | (y >> S[2])) & B[2]

	x = (x | (x >> S[3])) & B[3]
	y = (y | (y >> S[3])) & B[3]

	x = (x | (x >> S[4])) & B[4]
	y = (y | (y >> S[4])) & B[4]

	x = (x | (x >> S[5])) & B[5]
	y = (y | (y >> S[5])) & B[5]
	return x, y
}

// converts a uint64 quadkey to a string representation.
// the string can be visualized using https://labs.mapbox.com/what-the-tile/
func QuadkeyStr(quadkey uint64) string {
	sb := strings.Builder{}
	sb.Grow(zoom)

	for offset := 62; offset >= 0; offset -= 2 {
		digit := (quadkey >> offset) & 3
		switch digit {
		case 0:
			sb.WriteByte('0')
		case 1:
			sb.WriteByte('1')
		case 2:
			sb.WriteByte('2')
		case 3:
			sb.WriteByte('3')
		}
	}
	return sb.String()
}

// returns a list of quadkey buckets that can be used with a SQL query.
// this performs a dfs with depth limited by coarseness; so small coarseness can be slow.
// this is a much more accurate search for buckets in rect
//
//	radius is in meters
//	coarseness is in meters, and defines the depth of buckets
func QuadkeySearchBuckets(lat float64, long float64, radius float64, coarseness float64) map[uint64]QuadkeyBucket {
	maxDepth := coarseToDepth(coarseness)
	x, y := LatLongToXY(lat, long)
	radiusXY := uint64(radius * METERS_TO_XY)

	searchQuad := quad{
		x - radiusXY,
		x + radiusXY,
		y - radiusXY,
		y + radiusXY,
		0,
		0,
	}

	bucketsMap := map[uint64]QuadkeyBucket{}

	rectBucketSearch(&searchQuad, 0x0000000000000000, 0, maxDepth, bucketsMap)
	rectBucketSearch(&searchQuad, 0x4000000000000000, 0, maxDepth, bucketsMap)
	rectBucketSearch(&searchQuad, 0x8000000000000000, 0, maxDepth, bucketsMap)
	rectBucketSearch(&searchQuad, 0xC000000000000000, 0, maxDepth, bucketsMap)

	return bucketsMap
}

// returns a list of quadkey buckets that can be used with a SQL query.
// this does grid lookup for quadkeys, using coarseness as a stepsize.
// this is very inaccurate when coarseness is larger than radius
//
//		radius is in meters
//		coarseness is in meters, and defines the depth of buckets and stepsize for grid.
//	     coarseness is clamped to 2x radius
func QuadkeyGridBuckets(lat float64, long float64, radius float64, coarseness float64) map[uint64]QuadkeyBucket {
	depth := coarseToDepth(coarseness)
	x, y := LatLongToXY(lat, long)
	radiusXY := uint64(radius * METERS_TO_XY)
	stepXY := uint64(math.Min(coarseness, 2*radius) * METERS_TO_XY)

	var mask uint64 = 0xFFFFFFFFFFFFFFFF << (64 - depth*2 - 2)

	x -= radiusXY
	y -= radiusXY

	bucketsMap := map[uint64]QuadkeyBucket{}

	for i := uint64(0); i <= 2*radiusXY; i += stepXY {
		for j := uint64(0); j <= 2*radiusXY; j += stepXY {
			qkey := XYToQuadkey(x+i, y+j)
			qkey &= mask

			_, ok := bucketsMap[qkey]
			if !ok {
				bucketsMap[qkey] = NewQuadkeyBucket(qkey, depth)
			}
		}
	}

	return bucketsMap
}

// recursive dfs search down to desired maxdepth.
// quadkeys are appended to buckets map
func rectBucketSearch(searchQuad *quad, node uint64, depth int, maxDepth int, buckets map[uint64]QuadkeyBucket) {
	nodeQuad := QuadkeyToQuad(node, depth)

	if !nodeQuad.Intersects(searchQuad) {
		return
	}

	if depth < maxDepth {
		offset := 62 - (depth*2 + 2)
		rectBucketSearch(searchQuad, node, depth+1, maxDepth, buckets)             //child 0
		rectBucketSearch(searchQuad, node|(1<<offset), depth+1, maxDepth, buckets) //child 1
		rectBucketSearch(searchQuad, node|(2<<offset), depth+1, maxDepth, buckets) //child 2
		rectBucketSearch(searchQuad, node|(3<<offset), depth+1, maxDepth, buckets) //child 3
	} else {
		_, ok := buckets[node]
		if !ok {
			buckets[node] = NewQuadkeyBucket(node, depth)
		}
	}
}

// converts a coarseness (meters) to a quadtree depth (int)
// depth 0 at maximum coarseness is bit 63. depth 32 at max resolution is bit 0
func coarseToDepth(coarseness float64) int {
	size := EARTH_CIRCUMFERENCE / 2.0
	depth := 0
	for size > coarseness {
		depth++
		size /= 2.0
	}
	return depth - 1
}

// a quad is a rectangle used for intersection checks when performing
// quadtree searches
type quad struct {
	x0 uint64
	x1 uint64
	y0 uint64
	y1 uint64

	Quadkey uint64
	Depth   int
}

func (q *quad) Intersects(other *quad) bool {
	if q.x0 >= q.x1 || q.y0 >= q.y1 || other.x0 >= other.x1 || other.y0 >= other.y1 {
		return false
	}
	if q.x0 > other.x1 || other.x0 > q.x1 {
		return false
	}
	if q.y0 > other.y1 || other.y0 > q.y1 {
		return false
	}
	return true
}

// creates a quad from the given quadkey and depth
func QuadkeyToQuad(quadkey uint64, depth int) quad {
	var mask uint64 = 0xFFFFFFFFFFFFFFFF << (64 - depth*2 - 2)
	quadkey &= mask
	width := uint64(1<<(32-depth-1) - 1)

	x0, y0 := QuadkeyToXY(quadkey)

	q := quad{
		x0:      x0,
		y0:      y0,
		x1:      x0 + width,
		y1:      y0 + width,
		Quadkey: quadkey,
		Depth:   depth,
	}

	return q
}

// A QuadkeyBucket represents a contiguous range of quadkey indices that can be searched in a SQL db.
// an example query might use "WHERE quadkey BETWEEN {bucket.Start} AND {bucket.End}"
type QuadkeyBucket struct {
	Quadkey uint64 //the parent node for the bucket
	Depth   int    //parent node depth for the bucket
	Start   uint64 //Start Index of the bucket
	End     uint64 //End Index of the bucket
}

func NewQuadkeyBucket(quadkey uint64, depth int) QuadkeyBucket {
	var mask uint64 = 0xFFFFFFFFFFFFFFFF << (64 - depth*2 - 2)

	return QuadkeyBucket{
		quadkey,
		depth,
		quadkey & mask,
		quadkey | (^mask),
	}
}