1 // Copyright 2020 Google Inc. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package s2 16 17 import ( 18 "sort" 19 ) 20 21 const ( 22 // A special label indicating that the ContentsIterator done is true. 23 cellIndexDoneContents = -1 24 ) 25 26 // cellIndexNode represents a node in the CellIndex. Cells are organized in a 27 // tree such that the ancestors of a given node contain that node. 28 type cellIndexNode struct { 29 cellID CellID 30 label int32 31 parent int32 32 } 33 34 // newCellIndexNode returns a node with the appropriate default values. 35 func newCellIndexNode() cellIndexNode { 36 return cellIndexNode{ 37 cellID: 0, 38 label: cellIndexDoneContents, 39 parent: -1, 40 } 41 } 42 43 // A rangeNode represents a range of leaf CellIDs. The range starts at 44 // startID (a leaf cell) and ends at the startID field of the next 45 // rangeNode. contents points to the node of the CellIndex cellTree 46 // representing the cells that overlap this range. 47 type rangeNode struct { 48 startID CellID // First leaf cell contained by this range. 49 contents int32 // Contents of this node (an index within the cell tree). 50 } 51 52 // CellIndexIterator is an iterator that visits the entire set of indexed 53 // (CellID, label) pairs in an unspecified order. 54 type CellIndexIterator struct { 55 // TODO(roberts): Implement 56 } 57 58 // NewCellIndexIterator creates an iterator for the given CellIndex. 59 func NewCellIndexIterator(index *CellIndex) *CellIndexIterator { 60 return &CellIndexIterator{} 61 } 62 63 // CellIndexRangeIterator is an iterator that seeks and iterates over a set of 64 // non-overlapping leaf cell ranges that cover the entire sphere. The indexed 65 // (CellID, label) pairs that intersect the current leaf cell range can be 66 // visited using CellIndexContentsIterator (see below). 67 type CellIndexRangeIterator struct { 68 rangeNodes []rangeNode 69 pos int 70 nonEmpty bool 71 } 72 73 // NewCellIndexRangeIterator creates an iterator for the given CellIndex. 74 // The iterator is initially *unpositioned*; you must call a positioning method 75 // such as Begin() or Seek() before accessing its contents. 76 func NewCellIndexRangeIterator(index *CellIndex) *CellIndexRangeIterator { 77 return &CellIndexRangeIterator{ 78 rangeNodes: index.rangeNodes, 79 } 80 } 81 82 // NewCellIndexNonEmptyRangeIterator creates an iterator for the given CellIndex. 83 // The iterator is initially *unpositioned*; you must call a positioning method such as 84 // Begin() or Seek() before accessing its contents. 85 func NewCellIndexNonEmptyRangeIterator(index *CellIndex) *CellIndexRangeIterator { 86 return &CellIndexRangeIterator{ 87 rangeNodes: index.rangeNodes, 88 nonEmpty: true, 89 } 90 } 91 92 // StartID reports the CellID of the start of the current range of leaf CellIDs. 93 // 94 // If done is true, this returns the last possible CellID. This property means 95 // that most loops do not need to test done explicitly. 96 func (c *CellIndexRangeIterator) StartID() CellID { 97 return c.rangeNodes[c.pos].startID 98 } 99 100 // LimitID reports the non-inclusive end of the current range of leaf CellIDs. 101 // 102 // This assumes the iterator is not done. 103 func (c *CellIndexRangeIterator) LimitID() CellID { 104 return c.rangeNodes[c.pos+1].startID 105 } 106 107 // IsEmpty reports if no (CellID, label) pairs intersect this range. 108 // Also returns true if done() is true. 109 func (c *CellIndexRangeIterator) IsEmpty() bool { 110 return c.rangeNodes[c.pos].contents == cellIndexDoneContents 111 } 112 113 // Begin positions the iterator at the first range of leaf cells (if any). 114 func (c *CellIndexRangeIterator) Begin() { 115 c.pos = 0 116 for c.nonEmpty && c.IsEmpty() && !c.Done() { 117 c.pos++ 118 } 119 } 120 121 // Prev positions the iterator at the previous entry and reports whether it was not 122 // already positioned at the beginning. 123 func (c *CellIndexRangeIterator) Prev() bool { 124 if c.nonEmpty { 125 return c.nonEmptyPrev() 126 } 127 return c.prev() 128 } 129 130 // prev is used to position the iterator at the previous entry without checking 131 // if nonEmpty is true to prevent unwanted recursion. 132 func (c *CellIndexRangeIterator) prev() bool { 133 if c.pos == 0 { 134 return false 135 } 136 137 c.pos-- 138 return true 139 } 140 141 // Prev positions the iterator at the previous entry, and reports whether it was 142 // already positioned at the beginning. 143 func (c *CellIndexRangeIterator) nonEmptyPrev() bool { 144 for c.prev() { 145 if !c.IsEmpty() { 146 return true 147 } 148 } 149 150 // Return the iterator to its original position. 151 if c.IsEmpty() && !c.Done() { 152 c.Next() 153 } 154 return false 155 } 156 157 // Next advances the iterator to the next range of leaf cells. 158 // 159 // This assumes the iterator is not done. 160 func (c *CellIndexRangeIterator) Next() { 161 c.pos++ 162 for c.nonEmpty && c.IsEmpty() && !c.Done() { 163 c.pos++ 164 } 165 } 166 167 // Advance reports if advancing would leave it positioned on a valid range. If 168 // the value would not be valid, the positioning is not changed. 169 func (c *CellIndexRangeIterator) Advance(n int) bool { 170 // Note that the last element of rangeNodes is a sentinel value. 171 if n >= len(c.rangeNodes)-1-c.pos { 172 return false 173 } 174 c.pos += n 175 return true 176 } 177 178 // Finish positions the iterator so that done is true. 179 func (c *CellIndexRangeIterator) Finish() { 180 // Note that the last element of rangeNodes is a sentinel value. 181 c.pos = len(c.rangeNodes) - 1 182 } 183 184 // Done reports if the iterator is positioned beyond the last valid range. 185 func (c *CellIndexRangeIterator) Done() bool { 186 return c.pos >= len(c.rangeNodes)-1 187 } 188 189 // Seek positions the iterator at the first range with startID >= target. 190 // Such an entry always exists as long as "target" is a valid leaf cell. 191 // 192 // Note that it is valid to access startID even when done is true. 193 func (c *CellIndexRangeIterator) Seek(target CellID) { 194 c.pos = sort.Search(len(c.rangeNodes), func(i int) bool { 195 return c.rangeNodes[i].startID > target 196 }) - 1 197 198 // Ensure we don't go beyond the beginning. 199 if c.pos < 0 { 200 c.pos = 0 201 } 202 203 // Nonempty needs to find the next non-empty entry. 204 for c.nonEmpty && c.IsEmpty() && !c.Done() { 205 // c.Next() 206 c.pos++ 207 } 208 } 209 210 // CellIndexContentsIterator is an iterator that visits the (CellID, label) pairs 211 // that cover a set of leaf cell ranges (see CellIndexRangeIterator). Note that 212 // when multiple leaf cell ranges are visited, this iterator only guarantees that 213 // each result will be reported at least once, i.e. duplicate values may be 214 // suppressed. If you want duplicate values to be reported again, be sure to call 215 // Clear first. 216 // 217 // In particular, the implementation guarantees that when multiple leaf 218 // cell ranges are visited in monotonically increasing order, then each 219 // (CellID, label) pair is reported exactly once. 220 type CellIndexContentsIterator struct { 221 // The maximum index within the cellTree slice visited during the 222 // previous call to StartUnion. This is used to eliminate duplicate 223 // values when StartUnion is called multiple times. 224 nodeCutoff int32 225 226 // The maximum index within the cellTree visited during the 227 // current call to StartUnion. This is used to update nodeCutoff. 228 nextNodeCutoff int32 229 230 // The value of startID from the previous call to StartUnion. 231 // This is used to check whether these values are monotonically 232 // increasing. 233 prevStartID CellID 234 235 // The cell tree from CellIndex 236 cellTree []cellIndexNode 237 238 // A copy of the current node in the cell tree. 239 node cellIndexNode 240 } 241 242 // NewCellIndexContentsIterator returns a new contents iterator. 243 // 244 // Note that the iterator needs to be positioned using StartUnion before 245 // it can be safely used. 246 func NewCellIndexContentsIterator(index *CellIndex) *CellIndexContentsIterator { 247 it := &CellIndexContentsIterator{ 248 cellTree: index.cellTree, 249 prevStartID: 0, 250 nodeCutoff: -1, 251 nextNodeCutoff: -1, 252 node: cellIndexNode{label: cellIndexDoneContents}, 253 } 254 return it 255 } 256 257 // Clear clears all state with respect to which range(s) have been visited. 258 func (c *CellIndexContentsIterator) Clear() { 259 c.prevStartID = 0 260 c.nodeCutoff = -1 261 c.nextNodeCutoff = -1 262 c.node.label = cellIndexDoneContents 263 } 264 265 // CellID returns the current CellID. 266 func (c *CellIndexContentsIterator) CellID() CellID { 267 return c.node.cellID 268 } 269 270 // Label returns the current Label. 271 func (c *CellIndexContentsIterator) Label() int32 { 272 return c.node.label 273 } 274 275 // Next advances the iterator to the next (CellID, label) pair covered by the 276 // current leaf cell range. 277 // 278 // This requires the iterator to not be done. 279 func (c *CellIndexContentsIterator) Next() { 280 if c.node.parent <= c.nodeCutoff { 281 // We have already processed this node and its ancestors. 282 c.nodeCutoff = c.nextNodeCutoff 283 c.node.label = cellIndexDoneContents 284 } else { 285 c.node = c.cellTree[c.node.parent] 286 } 287 } 288 289 // Done reports if all (CellID, label) pairs have been visited. 290 func (c *CellIndexContentsIterator) Done() bool { 291 return c.node.label == cellIndexDoneContents 292 } 293 294 // StartUnion positions the ContentsIterator at the first (cell_id, label) pair 295 // that covers the given leaf cell range. Note that when multiple leaf cell 296 // ranges are visited using the same ContentsIterator, duplicate values 297 // may be suppressed. If you don't want this behavior, call Reset() first. 298 func (c *CellIndexContentsIterator) StartUnion(r *CellIndexRangeIterator) { 299 if r.StartID() < c.prevStartID { 300 c.nodeCutoff = -1 // Can't automatically eliminate duplicates. 301 } 302 c.prevStartID = r.StartID() 303 304 contents := r.rangeNodes[r.pos].contents 305 if contents <= c.nodeCutoff { 306 c.node.label = cellIndexDoneContents 307 } else { 308 c.node = c.cellTree[contents] 309 } 310 311 // When visiting ancestors, we can stop as soon as the node index is smaller 312 // than any previously visited node index. Because indexes are assigned 313 // using a preorder traversal, such nodes are guaranteed to have already 314 // been reported. 315 c.nextNodeCutoff = contents 316 } 317 318 // CellIndex stores a collection of (CellID, label) pairs. 319 // 320 // The CellIDs may be overlapping or contain duplicate values. For example, a 321 // CellIndex could store a collection of CellUnions, where each CellUnion 322 // gets its own non-negative int32 label. 323 // 324 // Similar to ShapeIndex and PointIndex which map each stored element to an 325 // identifier, CellIndex stores a label that is typically used to map the 326 // results of queries back to client's specific data. 327 // 328 // The zero value for a CellIndex is sufficient when constructing a CellIndex. 329 // 330 // To build a CellIndex where each Cell has a distinct label, call Add for each 331 // (CellID, label) pair, and then Build the index. For example: 332 // 333 // // contents is a mapping of an identifier in my system (restaurantID, 334 // // vehicleID, etc) to a CellID 335 // var contents = map[int32]CellID{...} 336 // 337 // for key, val := range contents { 338 // index.Add(val, key) 339 // } 340 // 341 // index.Build() 342 // 343 // There is also a helper method that adds all elements of CellUnion with the 344 // same label: 345 // 346 // index.AddCellUnion(cellUnion, label) 347 // 348 // Note that the index is not dynamic; the contents of the index cannot be 349 // changed once it has been built. Adding more after calling Build results in 350 // undefined behavior of the index. 351 // 352 // There are several options for retrieving data from the index. The simplest 353 // is to use a built-in method such as IntersectingLabels (which returns 354 // the labels of all cells that intersect a given target CellUnion): 355 // 356 // labels := index.IntersectingLabels(targetUnion); 357 // 358 // Alternatively, you can use a ClosestCellQuery which computes the cell(s) 359 // that are closest to a given target geometry. 360 // 361 // For example, here is how to find all cells that are closer than 362 // distanceLimit to a given target point: 363 // 364 // query := NewClosestCellQuery(cellIndex, opts) 365 // target := NewMinDistanceToPointTarget(targetPoint); 366 // for result := range query.FindCells(target) { 367 // // result.Distance() is the distance to the target. 368 // // result.CellID() is the indexed CellID. 369 // // result.Label() is the label associated with the CellID. 370 // DoSomething(targetPoint, result); 371 // } 372 // 373 // Internally, the index consists of a set of non-overlapping leaf cell ranges 374 // that subdivide the sphere and such that each range intersects a particular 375 // set of (cellID, label) pairs. 376 // 377 // Most clients should use either the methods such as VisitIntersectingCells 378 // and IntersectingLabels, or a helper such as ClosestCellQuery. 379 type CellIndex struct { 380 // A tree of (cellID, label) pairs such that if X is an ancestor of Y, then 381 // X.cellID contains Y.cellID. The contents of a given range of leaf 382 // cells can be represented by pointing to a node of this tree. 383 cellTree []cellIndexNode 384 385 // The last element of rangeNodes is a sentinel value, which is necessary 386 // in order to represent the range covered by the previous element. 387 rangeNodes []rangeNode 388 } 389 390 // Add adds the given CellID and Label to the index. 391 func (c *CellIndex) Add(id CellID, label int32) { 392 if label < 0 { 393 panic("labels must be non-negative") 394 } 395 c.cellTree = append(c.cellTree, cellIndexNode{cellID: id, label: label, parent: -1}) 396 } 397 398 // AddCellUnion adds all of the elements of the given CellUnion to the index with the same label. 399 func (c *CellIndex) AddCellUnion(cu CellUnion, label int32) { 400 if label < 0 { 401 panic("labels must be non-negative") 402 } 403 for _, cell := range cu { 404 c.Add(cell, label) 405 } 406 } 407 408 // Build builds the index for use. This method should only be called once. 409 func (c *CellIndex) Build() { 410 // To build the cell tree and leaf cell ranges, we maintain a stack of 411 // (CellID, label) pairs that contain the current leaf cell. This struct 412 // represents an instruction to push or pop a (cellID, label) pair. 413 // 414 // If label >= 0, the (cellID, label) pair is pushed on the stack. 415 // If CellID == SentinelCellID, a pair is popped from the stack. 416 // Otherwise the stack is unchanged but a rangeNode is still emitted. 417 418 // delta represents an entry in a stack of (CellID, label) pairs used in the 419 // construction of the CellIndex structure. 420 type delta struct { 421 startID CellID 422 cellID CellID 423 label int32 424 } 425 426 deltas := make([]delta, 0, 2*len(c.cellTree)+2) 427 428 // Create two deltas for each (cellID, label) pair: one to add the pair to 429 // the stack (at the start of its leaf cell range), and one to remove it from 430 // the stack (at the end of its leaf cell range). 431 for _, node := range c.cellTree { 432 deltas = append(deltas, delta{ 433 startID: node.cellID.RangeMin(), 434 cellID: node.cellID, 435 label: node.label, 436 }) 437 deltas = append(deltas, delta{ 438 startID: node.cellID.RangeMax().Next(), 439 cellID: SentinelCellID, 440 label: -1, 441 }) 442 } 443 444 // We also create two special deltas to ensure that a RangeNode is emitted at 445 // the beginning and end of the CellID range. 446 deltas = append(deltas, delta{ 447 startID: CellIDFromFace(0).ChildBeginAtLevel(MaxLevel), 448 cellID: CellID(0), 449 label: -1, 450 }) 451 deltas = append(deltas, delta{ 452 startID: CellIDFromFace(5).ChildEndAtLevel(MaxLevel), 453 cellID: CellID(0), 454 label: -1, 455 }) 456 457 sort.Slice(deltas, func(i, j int) bool { 458 // deltas are sorted first by startID, then in reverse order by cellID, 459 // and then by label. This is necessary to ensure that (1) larger cells 460 // are pushed on the stack before smaller cells, and (2) cells are popped 461 // off the stack before any new cells are added. 462 463 if si, sj := deltas[i].startID, deltas[j].startID; si != sj { 464 return si < sj 465 } 466 if si, sj := deltas[i].cellID, deltas[j].cellID; si != sj { 467 return si > sj 468 } 469 return deltas[i].label < deltas[j].label 470 }) 471 472 // Now walk through the deltas to build the leaf cell ranges and cell tree 473 // (which is essentially a permanent form of the "stack" described above). 474 c.cellTree = nil 475 c.rangeNodes = nil 476 contents := int32(-1) 477 for i := 0; i < len(deltas); { 478 startID := deltas[i].startID 479 // Process all the deltas associated with the current startID. 480 for ; i < len(deltas) && deltas[i].startID == startID; i++ { 481 if deltas[i].label >= 0 { 482 c.cellTree = append(c.cellTree, cellIndexNode{ 483 cellID: deltas[i].cellID, 484 label: deltas[i].label, 485 parent: contents}) 486 contents = int32(len(c.cellTree) - 1) 487 } else if deltas[i].cellID == SentinelCellID { 488 contents = c.cellTree[contents].parent 489 } 490 } 491 c.rangeNodes = append(c.rangeNodes, rangeNode{startID, contents}) 492 } 493 } 494 495 // TODO(roberts): Differences from C++ 496 // IntersectingLabels 497 // VisitIntersectingCells 498 // CellIndexIterator 499