mtsatellite/common/coverage.go

package common

import "sync"

type zRange struct {
	y1     int16
	y2     int16
	xRange *Span
}

type Coverage3D struct {
	pool    *SpanPool
	zRanges map[int16]*zRange
	mu      sync.RWMutex
}

type Range struct {
	Z  int16
	Y1 int16
	Y2 int16
	X1 int16
	X2 int16
}

func NewCoverage3D() *Coverage3D {
	return &Coverage3D{
		pool:    NewSpanPool(),
		zRanges: map[int16]*zRange{}}
}

func (c3d *Coverage3D) Insert(c Coord) {
	c3d.mu.Lock()
	defer c3d.mu.Unlock()
	zr := c3d.zRanges[c.Z]
	if zr == nil {
		xr := c3d.pool.Alloc()
		xr.From = int32(c.X)
		xr.To = int32(c.X)
		xr.Next = nil
		c3d.zRanges[c.Z] = &zRange{
			y1:     c.Y,
			y2:     c.Y,
			xRange: xr}
		return
	}
	zr.xRange = c3d.pool.Insert(zr.xRange, int32(c.X), 0)
	if c.Y < zr.y1 {
		zr.y1 = c.Y
	}
	if c.Y > zr.y2 {
		zr.y2 = c.Y
	}
}

func (c3d *Coverage3D) Query(c1, c2 Coord) []Range {

	c1, c2 = MinCoord(c1, c2), MaxCoord(c1, c2)

	c3d.mu.RLock()
	defer c3d.mu.RUnlock()

	r := make([]Range, 0, 32)
	for z := c1.Z; z <= c2.Z; z++ {
		zr := c3d.zRanges[z]
		if zr == nil || c1.Y > zr.y2 || c2.Y < zr.y1 {
			continue
		}
		y1, y2 := max16(c1.Y, zr.y1), min16(c2.Y, zr.y2)
		for xr := zr.xRange; xr != nil && xr.From <= int32(c2.X); xr = xr.Next {
			if xr.To < int32(c1.X) {
				continue
			}
			r = append(r, Range{
				Z:  z,
				Y1: y1,
				Y2: y2,
				X1: max16(c1.X, int16(xr.From)),
				X2: min16(c2.X, int16(xr.To))})
		}
	}
	return r
}