/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.lucene.spatial.prefix.tree;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import com.google.common.geometry.S2CellId;
import org.apache.lucene.util.BytesRef;
import org.locationtech.spatial4j.shape.Shape;
import org.locationtech.spatial4j.shape.SpatialRelation;

This class represents a S2 pixel in the RPT.
@lucene.internal
/** * This class represents a S2 pixel in the RPT. * * @lucene.internal */
class S2PrefixTreeCell implements CellCanPrune { //Faces of S2 Geometry private static S2CellId[] FACES = new S2CellId[6]; static { FACES[0] = S2CellId.fromFacePosLevel(0, 0, 0); FACES[1] = S2CellId.fromFacePosLevel(1, 0, 0); FACES[2] = S2CellId.fromFacePosLevel(2, 0, 0); FACES[3] = S2CellId.fromFacePosLevel(3, 0, 0); FACES[4] = S2CellId.fromFacePosLevel(4, 0, 0); FACES[5] = S2CellId.fromFacePosLevel(5, 0, 0); } /*Special character to define a cell leaf*/ private static final byte LEAF = '+'; /*Tokens are used to serialize cells*/ private static final byte[] TOKENS = {'.', '/', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z'}; /*Map containing mapping between tokens and integer values*/ private static final Map<Byte, Integer> PIXELS; static { PIXELS = new HashMap<>(TOKENS.length); for (int i = 0; i < TOKENS.length; i++) { PIXELS.put(TOKENS[i], i); } } S2CellId cellId; int level; //cache level S2PrefixTree tree; SpatialRelation shapeRel = null; boolean isLeaf; Shape shape = null; S2PrefixTreeCell(S2PrefixTree tree, S2CellId cellId) { this.cellId = cellId; this.tree = tree; setLevel(); if (getLevel() == tree.getMaxLevels()) { setLeaf(); } } void readCell(S2PrefixTree tree, BytesRef ref) { isLeaf = false; shape = null; shapeRel = null; this.tree = tree; cellId = getS2CellIdFromBytesRef(ref); setLevel(); if (isLeaf(ref) || getLevel() == tree.getMaxLevels()) { setLeaf(); } } @Override public SpatialRelation getShapeRel() { return shapeRel; } @Override public void setShapeRel(SpatialRelation rel) { shapeRel = rel; } @Override public boolean isLeaf() { return isLeaf; } @Override public void setLeaf() { isLeaf = true; } @Override public BytesRef getTokenBytesWithLeaf(BytesRef result) { result = getTokenBytesNoLeaf(result); //max levels do not have leaf if (isLeaf() && !(getLevel() == tree.getMaxLevels())) { //Add leaf byte result.bytes[result.offset + result.length] = LEAF; result.length++; } return result; } @Override public BytesRef getTokenBytesNoLeaf(BytesRef result) { if (result == null) { result = new BytesRef(); } getBytesRefFromS2CellId(cellId, result); return result; } @Override public int getLevel() { return this.level; }
Cache level of cell.
/** * Cache level of cell. */
private void setLevel() { if (this.cellId == null) { this.level = 0; } else { assert cellId.level() % tree.arity == 0; this.level = (this.cellId.level() / tree.arity) + 1; } } @Override public CellIterator getNextLevelCells(Shape shapeFilter) { S2CellId[] children; if (cellId == null) { // this is the world cell children = FACES; } else { int nChildren = (int) Math.pow(4, tree.arity); children = new S2CellId[nChildren]; children[0] = cellId.childBegin(cellId.level() + tree.arity); for (int i = 1; i < nChildren; i++) { children[i] = children[i - 1].next(); } } List<Cell> cells = new ArrayList<>(children.length); for (S2CellId pixel : children) { cells.add(new S2PrefixTreeCell(tree, pixel)); } return new FilterCellIterator(cells.iterator(), shapeFilter); } @Override public Shape getShape() { if (shape == null) { if (cellId == null) { //World cell shape = tree.getSpatialContext().getWorldBounds(); } else { shape = tree.s2ShapeFactory.getS2CellShape(cellId); } } return shape; } @Override public boolean isPrefixOf(Cell c) { if (cellId == null) { return true; } S2PrefixTreeCell cell = (S2PrefixTreeCell) c; return cellId.contains(cell.cellId); } @Override public int compareToNoLeaf(Cell fromCell) { if (cellId == null) { return 1; } S2PrefixTreeCell cell = (S2PrefixTreeCell) fromCell; return cellId.compareTo(cell.cellId); }
Check if a cell is a leaf.
Params:
  • ref – The Byteref of the leaf
Returns:true if it is a leaf, e.g last byte is the special Character.
/** * Check if a cell is a leaf. * * @param ref The Byteref of the leaf * @return true if it is a leaf, e.g last byte is the special Character. */
private boolean isLeaf(BytesRef ref) { return (ref.bytes[ref.offset + ref.length - 1] == LEAF); }
Get the S2CellId from the BytesRef representation.
Params:
  • ref – The bytes.
Returns:the corresponding S2 cell.
/** * Get the {@link S2CellId} from the {@link BytesRef} representation. * * @param ref The bytes. * @return the corresponding S2 cell. */
private S2CellId getS2CellIdFromBytesRef(BytesRef ref) { int length = ref.length; if (isLeaf(ref)) { length--; } if (length == 0) { return null; //world cell } int face = PIXELS.get(ref.bytes[ref.offset]); S2CellId cellId = FACES[face]; long id = cellId.id(); for (int i = ref.offset + 1; i < ref.offset + length; i++) { int thisLevel = i - ref.offset; int pos = PIXELS.get(ref.bytes[i]); // first child at level id = id - (id & -id) + (1L << (2 * (S2CellId.MAX_LEVEL - thisLevel * tree.arity))); // next until pos id = id + pos * ((id & -id) << 1); } return new S2CellId(id); }
Codify a S2CellId into its BytesRef representation.
Params:
  • cellId – The S2 Cell id to codify.
  • bref – The byteref representation.
/** * Codify a {@link S2CellId} into its {@link BytesRef} representation. * * @param cellId The S2 Cell id to codify. * @param bref The byteref representation. */
private void getBytesRefFromS2CellId(S2CellId cellId, BytesRef bref) { if (cellId == null) {//world cell bref.length = 0; return; } int length = getLevel() + 1; byte[] b = bref.bytes.length >= length ? bref.bytes : new byte[length]; b[0] = TOKENS[cellId.face()]; for (int i = 1; i < getLevel(); i++) { int offset = 0; int level = tree.arity * i; for (int j = 1; j < tree.arity; j++) { offset = 4 * offset + cellId.childPosition(level - tree.arity + j); } b[i] = TOKENS[4 * offset + cellId.childPosition(level)]; } bref.bytes = b; bref.length = getLevel(); bref.offset = 0; } @Override public int getSubCellsSize() { if (cellId == null) { //root node return 6; } return (int) Math.pow(4, tree.arity); } @Override public int hashCode() { if (cellId == null) { return super.hashCode(); } return this.cellId.hashCode(); } @Override public boolean equals(Object o) { S2PrefixTreeCell cell = (S2PrefixTreeCell) o; return Objects.equals(cellId, cell.cellId); } @Override public String toString() { if (cellId == null) { return "0"; } return cellId.toString(); } }