func (p *Proximity) Search(p0, p1 s2.LatLng) []s2.CellID {
boundRect := s2.RectFromLatLng(p0)
boundRect = boundRect.AddPoint(p1)
startToken := []byte(s2.CellIDFromLatLng(boundRect.Hi()).ToToken())
endToken := []byte(s2.CellIDFromLatLng(boundRect.Lo()).ToToken())
results := []s2.CellID{}
p.DB.View(func(tx *bolt.Tx) error {
cursor := tx.Bucket([]byte("proximity")).Cursor()
for k, _ := cursor.Seek(startToken); k != nil && bytes.Compare(k, endToken) <= 0; k, _ = cursor.Next() {
cell := s2.CellIDFromToken(string(k))
results = append(results, cell)
}
return nil
})
return results
}
func (p *Proximity) Match(p0, p1 s2.LatLng) bool {
boundRect := s2.RectFromLatLng(p0)
boundRect = boundRect.AddPoint(p1)
startToken := []byte(s2.CellIDFromLatLng(boundRect.Hi()).ToToken())
endToken := []byte(s2.CellIDFromLatLng(boundRect.Lo()).ToToken())
found := false
p.DB.View(func(tx *bolt.Tx) error {
cursor := tx.Bucket([]byte("proximity")).Cursor()
for k, _ := cursor.Seek(startToken); k != nil && bytes.Compare(k, endToken) <= 0; k, _ = cursor.Next() {
found = true
break
}
return nil
})
return found
}
func (p *Proximity) AddLatlng(point s2.LatLng) error {
cell := s2.CellIDFromLatLng(point)
token := []byte(cell.ToToken())
return p.DB.Update(func(tx *bolt.Tx) error {
bucket := tx.Bucket([]byte("proximity"))
err := bucket.Put(token, []byte{})
return err
})
}
func TestProximitySearch(t *testing.T) {
points := []s2.LatLng{
s2.LatLngFromDegrees(40.724153, -73.992610), // nw
s2.LatLngFromDegrees(40.720533, -73.994144), // sw
s2.LatLngFromDegrees(40.722190, -73.986226), // ne
// Should not be included
s2.LatLngFromDegrees(40.717009, -73.983234), // SE
}
expectedCellIDs := []s2.CellID{
s2.CellIDFromLatLng(s2.LatLngFromDegrees(40.724153, -73.992610)), // nw
s2.CellIDFromLatLng(s2.LatLngFromDegrees(40.720533, -73.994144)), // sw
s2.CellIDFromLatLng(s2.LatLngFromDegrees(40.722190, -73.986226)), // ne
}
p, err := newProximity()
require.NoError(t, err)
p.AddLatLngs(points)
p0 := s2.LatLngFromDegrees(40.719657, -73.996632) // SW
p1 := s2.LatLngFromDegrees(40.723353, -73.984014) // NE
gotCellIDs := p.Search(p0, p1)
expected := make([]int, 0, len(expectedCellIDs))
for _, cell := range expectedCellIDs {
expected = append(expected, int(cell))
}
got := make([]int, 0, len(gotCellIDs))
for _, cell := range gotCellIDs {
got = append(got, int(cell))
}
sort.Ints(expected)
sort.Ints(got)
assert.Equal(t, expected, got, "Should contain the first three points")
}
func (p *Proximity) AddLatLngs(points []s2.LatLng) error {
var cell s2.CellID
var token []byte
return p.DB.Batch(func(tx *bolt.Tx) error {
bucket := tx.Bucket([]byte("proximity"))
for _, ll := range points {
cell = s2.CellIDFromLatLng(ll)
token = []byte(cell.ToToken())
err := bucket.Put(token, []byte{})
if err != nil {
return err
}
}
return nil
})
}
func TestAddLatLng(t *testing.T) {
points := []s2.LatLng{
s2.LatLngFromDegrees(40.724153, -73.992610), // nw
s2.LatLngFromDegrees(40.720533, -73.994144), // sw
s2.LatLngFromDegrees(40.722190, -73.986226), // ne
}
p, err := newProximity()
require.NoError(t, err)
for _, ll := range points {
assert.NoError(t, p.AddLatlng(ll))
key := []byte(s2.CellIDFromLatLng(ll).ToToken())
p.DB.View(func(tx *bolt.Tx) error {
b := tx.Bucket([]byte("proximity"))
v := b.Get(key)
assert.NoError(t, err)
assert.Equal(t, []byte{}, v)
return nil
})
}
}