533 lines
18 KiB
Java
533 lines
18 KiB
Java
package com.hypixel.hytale.component.spatial;
|
|
|
|
import com.hypixel.hytale.math.vector.Vector3d;
|
|
import it.unimi.dsi.fastutil.objects.ObjectArrayList;
|
|
import it.unimi.dsi.fastutil.objects.ObjectListIterator;
|
|
import java.util.Comparator;
|
|
import java.util.List;
|
|
import java.util.function.Predicate;
|
|
import javax.annotation.Nonnull;
|
|
import javax.annotation.Nullable;
|
|
|
|
public class KDTree<T> implements SpatialStructure<T> {
|
|
@Nonnull
|
|
private final List<KDTree.Node<T>> nodePool = new ObjectArrayList();
|
|
private int nodePoolIndex = 0;
|
|
@Nonnull
|
|
private final List<List<T>> dataListPool = new ObjectArrayList();
|
|
private int dataListPoolIndex = 0;
|
|
private int size;
|
|
@Nonnull
|
|
private final Predicate<T> collectionFilter;
|
|
@Nullable
|
|
private KDTree.Node<T> root;
|
|
|
|
public KDTree(@Nonnull Predicate<T> collectionFilter) {
|
|
this.collectionFilter = collectionFilter;
|
|
}
|
|
|
|
@Override
|
|
public int size() {
|
|
return this.size;
|
|
}
|
|
|
|
@Override
|
|
public void rebuild(@Nonnull SpatialData<T> spatialData) {
|
|
this.root = null;
|
|
this.size = 0;
|
|
int spatialDataSize = spatialData.size();
|
|
if (spatialDataSize != 0) {
|
|
for (int i = 0; i < this.dataListPoolIndex; i++) {
|
|
this.dataListPool.get(i).clear();
|
|
}
|
|
|
|
this.nodePoolIndex = 0;
|
|
this.dataListPoolIndex = 0;
|
|
spatialData.sortMorton();
|
|
int mid = spatialDataSize / 2;
|
|
int sortedIndex = spatialData.getSortedIndex(mid);
|
|
Vector3d vector = spatialData.getVector(sortedIndex);
|
|
T data = spatialData.getData(sortedIndex);
|
|
List<T> list = this.getPooledDataList();
|
|
list.add(data);
|
|
|
|
int left;
|
|
for (left = mid - 1; left >= 0; left--) {
|
|
int leftSortedIndex = spatialData.getSortedIndex(left);
|
|
Vector3d leftVector = spatialData.getVector(leftSortedIndex);
|
|
if (!leftVector.equals(vector)) {
|
|
break;
|
|
}
|
|
|
|
T leftData = spatialData.getData(leftSortedIndex);
|
|
list.add(leftData);
|
|
}
|
|
|
|
int right;
|
|
for (right = mid + 1; right < spatialDataSize; right++) {
|
|
int rightSortedIndex = spatialData.getSortedIndex(right);
|
|
Vector3d rightVector = spatialData.getVector(rightSortedIndex);
|
|
if (!rightVector.equals(vector)) {
|
|
break;
|
|
}
|
|
|
|
T rightData = spatialData.getData(rightSortedIndex);
|
|
list.add(rightData);
|
|
}
|
|
|
|
this.root = this.getPooledNode(vector, list);
|
|
if (0 < left + 1) {
|
|
this.build0(spatialData, 0, left + 1);
|
|
}
|
|
|
|
if (right < spatialDataSize) {
|
|
this.build0(spatialData, right, spatialDataSize);
|
|
}
|
|
|
|
this.size = spatialDataSize;
|
|
}
|
|
}
|
|
|
|
@Nullable
|
|
@Override
|
|
public T closest(@Nonnull Vector3d point) {
|
|
KDTree.ClosestState<T> closestState = new KDTree.ClosestState<>(null, Double.MAX_VALUE);
|
|
this.closest0(closestState, this.root, point, 0);
|
|
return closestState.node == null ? null : closestState.node.data.getFirst();
|
|
}
|
|
|
|
@Override
|
|
public void collect(@Nonnull Vector3d center, double radius, @Nonnull List<T> results) {
|
|
double distanceSq = radius * radius;
|
|
this.collect0(results, this.root, center, distanceSq, 0);
|
|
}
|
|
|
|
@Override
|
|
public void collectCylinder(@Nonnull Vector3d center, double radius, double height, @Nonnull List<T> results) {
|
|
double radiusSq = radius * radius;
|
|
double halfHeight = height / 2.0;
|
|
this.collectCylinder0(results, this.root, center, radiusSq, halfHeight, radius, 0);
|
|
}
|
|
|
|
@Override
|
|
public void collectBox(@Nonnull Vector3d min, @Nonnull Vector3d max, @Nonnull List<T> results) {
|
|
this.collectBox0(results, this.root, min, max, 0);
|
|
}
|
|
|
|
@Override
|
|
public void ordered(@Nonnull Vector3d center, double radius, @Nonnull List<T> results) {
|
|
double distanceSq = radius * radius;
|
|
ObjectArrayList<KDTree.OrderedEntry<T>> entryResults = new ObjectArrayList();
|
|
this.ordered0(entryResults, this.root, center, distanceSq, 0);
|
|
entryResults.sort(Comparator.comparingDouble(o -> o.distanceSq));
|
|
ObjectListIterator var8 = entryResults.iterator();
|
|
|
|
while (var8.hasNext()) {
|
|
KDTree.OrderedEntry<T> entry = (KDTree.OrderedEntry<T>)var8.next();
|
|
int i = 0;
|
|
|
|
for (int bound = entry.values.size(); i < bound; i++) {
|
|
T data = entry.values.get(i);
|
|
if (this.collectionFilter.test(data)) {
|
|
results.add(data);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public void ordered3DAxis(@Nonnull Vector3d center, double xSearchRadius, double YSearchRadius, double zSearchRadius, @Nonnull List<T> results) {
|
|
ObjectArrayList<KDTree.OrderedEntry<T>> entryResults = new ObjectArrayList();
|
|
this._internal_ordered3DAxis(entryResults, this.root, center, xSearchRadius, YSearchRadius, zSearchRadius, 0);
|
|
entryResults.sort(Comparator.comparingDouble(o -> o.distanceSq));
|
|
ObjectListIterator var10 = entryResults.iterator();
|
|
|
|
while (var10.hasNext()) {
|
|
KDTree.OrderedEntry<T> entry = (KDTree.OrderedEntry<T>)var10.next();
|
|
int i = 0;
|
|
|
|
for (int bound = entry.values.size(); i < bound; i++) {
|
|
T data = entry.values.get(i);
|
|
if (this.collectionFilter.test(data)) {
|
|
results.add(data);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
@Nonnull
|
|
@Override
|
|
public String dump() {
|
|
return "KDTree(size=" + this.size + ")\n" + (this.root == null ? null : this.root.dump(0));
|
|
}
|
|
|
|
@Nonnull
|
|
private KDTree.Node<T> getPooledNode(Vector3d vector, List<T> data) {
|
|
if (this.nodePoolIndex < this.nodePool.size()) {
|
|
KDTree.Node<T> node = this.nodePool.get(this.nodePoolIndex++);
|
|
node.reset(vector, data);
|
|
return node;
|
|
} else {
|
|
KDTree.Node<T> node = new KDTree.Node<>(vector, data);
|
|
this.nodePool.add(node);
|
|
this.nodePoolIndex++;
|
|
return node;
|
|
}
|
|
}
|
|
|
|
private List<T> getPooledDataList() {
|
|
if (this.dataListPoolIndex < this.dataListPool.size()) {
|
|
return this.dataListPool.get(this.dataListPoolIndex++);
|
|
} else {
|
|
ObjectArrayList<T> set = new ObjectArrayList(1);
|
|
this.dataListPool.add(set);
|
|
this.dataListPoolIndex++;
|
|
return set;
|
|
}
|
|
}
|
|
|
|
private void build0(@Nonnull SpatialData<T> spatialData, int start, int end) {
|
|
int mid = (start + end) / 2;
|
|
int sortedIndex = spatialData.getSortedIndex(mid);
|
|
Vector3d vector = spatialData.getVector(sortedIndex);
|
|
T data = spatialData.getData(sortedIndex);
|
|
List<T> list = this.getPooledDataList();
|
|
list.add(data);
|
|
|
|
int left;
|
|
for (left = mid - 1; left >= start; left--) {
|
|
int leftSortedIndex = spatialData.getSortedIndex(left);
|
|
Vector3d leftVector = spatialData.getVector(leftSortedIndex);
|
|
if (!leftVector.equals(vector)) {
|
|
break;
|
|
}
|
|
|
|
T leftData = spatialData.getData(leftSortedIndex);
|
|
list.add(leftData);
|
|
}
|
|
|
|
int right;
|
|
for (right = mid + 1; right < end; right++) {
|
|
int rightSortedIndex = spatialData.getSortedIndex(right);
|
|
Vector3d rightVector = spatialData.getVector(rightSortedIndex);
|
|
if (!rightVector.equals(vector)) {
|
|
break;
|
|
}
|
|
|
|
T rightData = spatialData.getData(rightSortedIndex);
|
|
list.add(rightData);
|
|
}
|
|
|
|
this.put0(this.root, vector, list, 0);
|
|
if (start < left + 1) {
|
|
this.build0(spatialData, start, left + 1);
|
|
}
|
|
|
|
if (right < end) {
|
|
this.build0(spatialData, right, end);
|
|
}
|
|
}
|
|
|
|
private void put0(@Nonnull KDTree.Node<T> node, @Nonnull Vector3d vector, @Nonnull List<T> list, int axis) {
|
|
if (compare(node.vector, vector, axis) < 0) {
|
|
if (node.one == null) {
|
|
node.one = this.getPooledNode(vector, list);
|
|
} else {
|
|
this.put0(node.one, vector, list, (axis + 1) % 3);
|
|
}
|
|
} else if (node.two == null) {
|
|
node.two = this.getPooledNode(vector, list);
|
|
} else {
|
|
this.put0(node.two, vector, list, (axis + 1) % 3);
|
|
}
|
|
}
|
|
|
|
private void closest0(@Nonnull KDTree.ClosestState<T> closestState, @Nullable KDTree.Node<T> node, @Nonnull Vector3d vector, int depth) {
|
|
if (node != null) {
|
|
if (vector.equals(node.vector)) {
|
|
closestState.distanceSq = 0.0;
|
|
closestState.node = node;
|
|
} else {
|
|
int axis = depth % 3;
|
|
int compare = compare(node.vector, vector, axis);
|
|
double distanceSq = node.vector.distanceSquaredTo(vector);
|
|
if (distanceSq < closestState.distanceSq) {
|
|
closestState.node = node;
|
|
closestState.distanceSq = distanceSq;
|
|
}
|
|
|
|
int newDepth = depth + 1;
|
|
if (compare < 0) {
|
|
this.closest0(closestState, node.one, vector, newDepth);
|
|
} else {
|
|
this.closest0(closestState, node.two, vector, newDepth);
|
|
}
|
|
|
|
double plane = get(node.vector, axis);
|
|
double component = get(closestState.node.vector, axis);
|
|
double planeDistance = Math.abs(component - plane);
|
|
if (planeDistance * planeDistance < closestState.distanceSq) {
|
|
if (compare < 0) {
|
|
this.closest0(closestState, node.two, vector, newDepth);
|
|
} else {
|
|
this.closest0(closestState, node.one, vector, newDepth);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private void collect0(@Nonnull List<T> results, @Nullable KDTree.Node<T> node, @Nonnull Vector3d vector, double distanceSq, int depth) {
|
|
if (node != null) {
|
|
int axis = depth % 3;
|
|
int compare = compare(node.vector, vector, axis);
|
|
double nodeDistanceSq = node.vector.distanceSquaredTo(vector);
|
|
if (nodeDistanceSq < distanceSq) {
|
|
int i = 0;
|
|
|
|
for (int bound = node.data.size(); i < bound; i++) {
|
|
T data = node.data.get(i);
|
|
if (this.collectionFilter.test(data)) {
|
|
results.add(data);
|
|
}
|
|
}
|
|
}
|
|
|
|
int newDepth = depth + 1;
|
|
if (compare < 0) {
|
|
this.collect0(results, node.one, vector, distanceSq, newDepth);
|
|
} else {
|
|
this.collect0(results, node.two, vector, distanceSq, newDepth);
|
|
}
|
|
|
|
double plane = get(node.vector, axis);
|
|
double component = get(vector, axis);
|
|
double planeDistance = Math.abs(component - plane);
|
|
if (planeDistance * planeDistance < distanceSq) {
|
|
if (compare < 0) {
|
|
this.collect0(results, node.two, vector, distanceSq, newDepth);
|
|
} else {
|
|
this.collect0(results, node.one, vector, distanceSq, newDepth);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private void collectCylinder0(
|
|
@Nonnull List<T> results, @Nullable KDTree.Node<T> node, @Nonnull Vector3d center, double radiusSq, double halfHeight, double radius, int depth
|
|
) {
|
|
if (node != null) {
|
|
int axis = depth % 3;
|
|
int compare = compare(node.vector, center, axis);
|
|
double dy = node.vector.y - center.y;
|
|
if (Math.abs(dy) <= halfHeight) {
|
|
double dx = node.vector.x - center.x;
|
|
double dz = node.vector.z - center.z;
|
|
double xzDistanceSq = dx * dx + dz * dz;
|
|
if (xzDistanceSq <= radiusSq) {
|
|
int i = 0;
|
|
|
|
for (int bound = node.data.size(); i < bound; i++) {
|
|
T data = node.data.get(i);
|
|
if (this.collectionFilter.test(data)) {
|
|
results.add(data);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
int newDepth = depth + 1;
|
|
if (compare < 0) {
|
|
this.collectCylinder0(results, node.one, center, radiusSq, halfHeight, radius, newDepth);
|
|
} else {
|
|
this.collectCylinder0(results, node.two, center, radiusSq, halfHeight, radius, newDepth);
|
|
}
|
|
|
|
double plane = get(node.vector, axis);
|
|
double component = get(center, axis);
|
|
double axisRadius = axis == 2 ? halfHeight : radius;
|
|
if (Math.abs(component - plane) <= axisRadius) {
|
|
if (compare < 0) {
|
|
this.collectCylinder0(results, node.two, center, radiusSq, halfHeight, radius, newDepth);
|
|
} else {
|
|
this.collectCylinder0(results, node.one, center, radiusSq, halfHeight, radius, newDepth);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private void collectBox0(@Nonnull List<T> results, @Nullable KDTree.Node<T> node, @Nonnull Vector3d min, @Nonnull Vector3d max, int depth) {
|
|
if (node != null) {
|
|
int axis = depth % 3;
|
|
if (node.vector.x >= min.x
|
|
&& node.vector.x <= max.x
|
|
&& node.vector.y >= min.y
|
|
&& node.vector.y <= max.y
|
|
&& node.vector.z >= min.z
|
|
&& node.vector.z <= max.z) {
|
|
int i = 0;
|
|
|
|
for (int bound = node.data.size(); i < bound; i++) {
|
|
T data = node.data.get(i);
|
|
if (this.collectionFilter.test(data)) {
|
|
results.add(data);
|
|
}
|
|
}
|
|
}
|
|
|
|
int newDepth = depth + 1;
|
|
double plane = get(node.vector, axis);
|
|
double minComponent = get(min, axis);
|
|
double maxComponent = get(max, axis);
|
|
if (maxComponent >= plane) {
|
|
this.collectBox0(results, node.one, min, max, newDepth);
|
|
}
|
|
|
|
if (minComponent <= plane) {
|
|
this.collectBox0(results, node.two, min, max, newDepth);
|
|
}
|
|
}
|
|
}
|
|
|
|
private void ordered0(@Nonnull List<KDTree.OrderedEntry<T>> results, @Nullable KDTree.Node<T> node, @Nonnull Vector3d vector, double distanceSq, int depth) {
|
|
if (node != null) {
|
|
int axis = depth % 3;
|
|
int compare = compare(node.vector, vector, axis);
|
|
double nodeDistanceSq = node.vector.distanceSquaredTo(vector);
|
|
if (nodeDistanceSq < distanceSq) {
|
|
results.add(new KDTree.OrderedEntry<>(nodeDistanceSq, node.data));
|
|
}
|
|
|
|
int newDepth = depth + 1;
|
|
if (compare < 0) {
|
|
this.ordered0(results, node.one, vector, distanceSq, newDepth);
|
|
} else {
|
|
this.ordered0(results, node.two, vector, distanceSq, newDepth);
|
|
}
|
|
|
|
double plane = get(node.vector, axis);
|
|
double component = get(vector, axis);
|
|
double planeDistance = Math.abs(component - plane);
|
|
if (planeDistance * planeDistance < distanceSq) {
|
|
if (compare < 0) {
|
|
this.ordered0(results, node.two, vector, distanceSq, newDepth);
|
|
} else {
|
|
this.ordered0(results, node.one, vector, distanceSq, newDepth);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
private void _internal_ordered3DAxis(
|
|
@Nonnull List<KDTree.OrderedEntry<T>> results,
|
|
@Nullable KDTree.Node<T> node,
|
|
@Nonnull Vector3d center,
|
|
double xSearchRadius,
|
|
double ySearchRadius,
|
|
double zSearchRadius,
|
|
int depth
|
|
) {
|
|
if (node != null) {
|
|
int axis = depth % 3;
|
|
boolean inCuboid = node.vector.x >= center.x - xSearchRadius
|
|
&& node.vector.x <= center.x + xSearchRadius
|
|
&& node.vector.y >= center.y - ySearchRadius
|
|
&& node.vector.y <= center.y + ySearchRadius
|
|
&& node.vector.z >= center.z - zSearchRadius
|
|
&& node.vector.z <= center.z + zSearchRadius;
|
|
if (inCuboid) {
|
|
double nodeDistanceSq = node.vector.distanceSquaredTo(center);
|
|
results.add(new KDTree.OrderedEntry<>(nodeDistanceSq, node.data));
|
|
}
|
|
|
|
int newDepth = depth + 1;
|
|
int compare = compare(node.vector, center, axis);
|
|
KDTree.Node<T> primary = compare < 0 ? node.one : node.two;
|
|
KDTree.Node<T> secondary = compare < 0 ? node.two : node.one;
|
|
this._internal_ordered3DAxis(results, primary, center, xSearchRadius, ySearchRadius, zSearchRadius, newDepth);
|
|
double plane = get(node.vector, axis);
|
|
double component = get(center, axis);
|
|
double radius = axis == 0 ? xSearchRadius : (axis == 1 ? ySearchRadius : zSearchRadius);
|
|
if (Math.abs(component - plane) <= radius) {
|
|
this._internal_ordered3DAxis(results, secondary, center, xSearchRadius, ySearchRadius, zSearchRadius, newDepth);
|
|
}
|
|
}
|
|
}
|
|
|
|
private static int compare(@Nonnull Vector3d v1, @Nonnull Vector3d v2, int axis) {
|
|
return switch (axis) {
|
|
case 0 -> Double.compare(v1.x, v2.x);
|
|
case 1 -> Double.compare(v1.z, v2.z);
|
|
case 2 -> Double.compare(v1.y, v2.y);
|
|
default -> throw new IllegalArgumentException("Invalid axis: " + axis);
|
|
};
|
|
}
|
|
|
|
private static double get(@Nonnull Vector3d v, int axis) {
|
|
return switch (axis) {
|
|
case 0 -> v.x;
|
|
case 1 -> v.z;
|
|
case 2 -> v.y;
|
|
default -> throw new IllegalArgumentException("Invalid axis: " + axis);
|
|
};
|
|
}
|
|
|
|
private static class ClosestState<T> {
|
|
private KDTree.Node<T> node;
|
|
private double distanceSq;
|
|
|
|
public ClosestState(KDTree.Node<T> node, double distanceSq) {
|
|
this.node = node;
|
|
this.distanceSq = distanceSq;
|
|
}
|
|
}
|
|
|
|
private static class Node<T> {
|
|
private Vector3d vector;
|
|
private List<T> data;
|
|
@Nullable
|
|
private KDTree.Node<T> one;
|
|
@Nullable
|
|
private KDTree.Node<T> two;
|
|
|
|
public Node(Vector3d vector, List<T> data) {
|
|
this.vector = vector;
|
|
this.data = data;
|
|
}
|
|
|
|
public void reset(Vector3d vector, List<T> data) {
|
|
this.vector = vector;
|
|
this.data = data;
|
|
this.one = null;
|
|
this.two = null;
|
|
}
|
|
|
|
@Nonnull
|
|
public String dump(int depth) {
|
|
int nextDepth = depth + 1;
|
|
return "vector="
|
|
+ this.vector
|
|
+ ", data="
|
|
+ this.data
|
|
+ ",\n"
|
|
+ " ".repeat(depth)
|
|
+ "one="
|
|
+ (this.one == null ? null : this.one.dump(nextDepth))
|
|
+ ",\n"
|
|
+ " ".repeat(depth)
|
|
+ "two="
|
|
+ (this.two == null ? null : this.two.dump(nextDepth));
|
|
}
|
|
}
|
|
|
|
private static class OrderedEntry<T> {
|
|
private final double distanceSq;
|
|
private final List<T> values;
|
|
|
|
public OrderedEntry(double distanceSq, List<T> values) {
|
|
this.distanceSq = distanceSq;
|
|
this.values = values;
|
|
}
|
|
}
|
|
}
|