Simple Octree implementation in Java (Immutable OOP)

Bounds.java

package tk.zorgatone.simpleoctree;

public interface Bounds {
  Point getTopRight();
  Point getBottomLeft();
}

Cube.java

package tk.zorgatone.simpleoctree;

public interface Cube {
  Octant getOctant();
  Bounds getBounds();
}

ElementWrapper.java

package tk.zorgatone.simpleoctree;

class ElementWrapper {

  private final OctreePoint point;
  private final Object object;

  public ElementWrapper(Point point, Object object) {
    this.point = new OctreePoint(point);
    this.object = object;
  }

  public OctreePoint getPoint() {
    return point;
  }

  public Object getObject() {
    return object;
  }

  @Override
  public String toString() {
    if (object == null) {
      return "null";
    }

    return object.toString();
  }
}

Octant.java

package tk.zorgatone.simpleoctree;

public interface Octant {
  byte MAX_ELEMENTS = OctantSlot.MAXIMUM_INDEX + 1;

  Object getSlot(OctantSlot slot);
}

OctantSlot.java

package tk.zorgatone.simpleoctree;

import java.io.Serial;
import java.io.Serializable;

public enum OctantSlot implements Serializable {

  UP_NORTH_EAST((byte) 0),
  UP_NORTH_WEST((byte) 1),
  UP_SOUTH_WEST((byte) 2),
  UP_SOUTH_EAST((byte) 3),
  DOWN_SOUTH_EAST((byte) 4),
  DOWN_SOUTH_WEST((byte) 5),
  DOWN_NORTH_WEST((byte) 6),
  DOWN_NORTH_EAST((byte) 7);

  @Serial
  private static final long serialVersionUID = -3114131496647548770L;

  public static final byte MAXIMUM_INDEX = (byte) 7;
  private final byte index;

  OctantSlot(byte index) {
    this.index = index;
  }

  public static OctantSlot fromIndex(byte index) throws SlotConversionException {
    if (index < 0 || index > MAXIMUM_INDEX) {
      throw new SlotConversionException(
        new IndexOutOfBoundsException("Index out of [0-7] positions!")
      );
    }

    for (OctantSlot p : values()) {
      if (p.getIndex() == index) {
        return p;
      }
    }

    throw new SlotConversionException(
      new IndexOutOfBoundsException("Index out of [0-7] positions!")
    );
  }

  public byte getIndex() {
    return index;
  }

}

Octree.java

package tk.zorgatone.simpleoctree;

import java.util.Iterator;

public final class Octree<T> {

  private OctreeCube head;

  public Octree() {
    head = new OctreeCube();
  }

  public Cube getHead() {
    return head;
  }

  public void insert(Point point, T object) {
    head = head.insert(point, object);
  }

  public T find(Point point) {
    if (point == null) {
      return null;
    }

    var result = head.find(point);

    if (result == null) {
      return null;
    }

    @SuppressWarnings("unchecked")
    var tmp = (T) result;

    return tmp;
  }

  public Iterable<T> searchRadius(Point sphereCenter, double sphereRadius) {
    Iterable<Object> result = head.searchRadius(sphereCenter, sphereRadius);

    return () -> {
      Iterator<Object> resultIterator = result.iterator();

      return new Iterator<>() {
        @Override
        public boolean hasNext() {
          return resultIterator.hasNext();
        }

        @Override
        public T next() {
          @SuppressWarnings("unchecked")
          T tmp = (T) resultIterator.next();
          return tmp;
        }
      };
    };
  }

}

OctreeBounds.java

package tk.zorgatone.simpleoctree;

import static java.lang.Long.MAX_VALUE;
import static java.lang.Long.MIN_VALUE;


final class OctreeBounds implements Bounds {

  private final OctreePoint topRight;
  private final OctreePoint bottomLeft;

  OctreeBounds() {
    this(
      new OctreePoint(MAX_VALUE, MAX_VALUE, MAX_VALUE),
      new OctreePoint(MIN_VALUE, MIN_VALUE, MIN_VALUE)
    );
  }

  OctreeBounds(Point topRight, Point topLeft) {
    if (topRight instanceof OctreePoint) {
      this.topRight = (OctreePoint) topRight;
    } else {
      this.topRight = new OctreePoint(topRight);
    }

    if (topLeft instanceof OctreePoint) {
      this.bottomLeft = (OctreePoint) topLeft;
    } else {
      this.bottomLeft = new OctreePoint(topLeft);
    }
  }

  static long distance(long one, long two) {
    try {
      if (one > two) {
        return Math.subtractExact(one, two);
      }

      return Math.subtractExact(two, one);
    } catch (Exception exception) {
      return -1;
    }
  }

  static double distance(Point one, Point two) {
    final long x1 = one.getX();
    final long y1 = one.getY();
    final long z1 = one.getZ();

    final long x2 = two.getX();
    final long y2 = two.getY();
    final long z2 = two.getZ();

    final double distX = distance(x2, x1);
    final double distY = distance(y2, y1);
    final double distZ = distance(z2, z1);

    // We are using multiplications because is faster than calling Math.pow
    return Math.sqrt(distX * distX +
      distY * distY +
      distZ * distZ);
  }

  static boolean isPointInsideSphere(
    Point point,
    Point sphereCenter,
    double sphereRadius
  ) {
    final double dist = distance(point, sphereCenter);

    if (Double.isNaN(dist) || dist < 0) {
      return false;
    }

    return dist < sphereRadius;
  }

  static long middlePoint(long one, long two) {
    // NOTE: divide first to prevent overflow
    return one / 2 + two / 2;
  }

  static OctreePoint middlePoint(Point topRight, Point bottomLeft) {
    if (topRight == null || bottomLeft == null) {
      return null;
    }

    final long x = middlePoint(topRight.getX(), bottomLeft.getX());
    final long y = middlePoint(topRight.getY(), bottomLeft.getY());
    final long z = middlePoint(topRight.getZ(), bottomLeft.getZ());

    return new OctreePoint(x, y, z);
  }

  @Override
  public OctreePoint getTopRight() {
    return topRight;
  }

  @Override
  public OctreePoint getBottomLeft() {
    return bottomLeft;
  }

  boolean within(Point point) {
    final long x = point.getX();
    final long y = point.getY();
    final long z = point.getZ();

    return bottomLeft.getX() <= x
      && topRight.getX() >= x
      && bottomLeft.getY() <= y
      && topRight.getY() >= y
      && bottomLeft.getZ() <= z
      && topRight.getZ() >= z;
  }

  boolean intersectsRadius(final Point sphereCenter, final double sphereRadius) {
    if (sphereCenter == null || sphereRadius < 0) {
      return false;
    }

    final long cX = sphereCenter.getX();
    final long cY = sphereCenter.getY();
    final long cZ = sphereCenter.getZ();

    // Get box closest point to sphere center by clamping
    final long x = Math.max(bottomLeft.getX(), Math.min(cX, topRight.getX()));
    final long y = Math.max(bottomLeft.getY(), Math.min(cY, topRight.getY()));
    final long z = Math.max(bottomLeft.getZ(), Math.min(cZ, topRight.getZ()));

    // We are using multiplications because is faster than calling Math.pow
    final double dist = distance(new OctreePoint(x, y, z), sphereCenter);

    if (Double.isNaN(dist) || dist < 0) {
      return false;
    }

    return dist < sphereRadius;
  }

  OctantSlot directionOf(Point point) {
    if (point == null) {
      return null;
    }

    if (within(point)) {
      final long x = point.getX();
      final long y = point.getY();
      final long z = point.getZ();

      var center = middlePoint(topRight, bottomLeft);

      if (x <= center.getX()) {
        if (y <= center.getY()) {
          if (z <= center.getZ()) {
            return OctantSlot.DOWN_SOUTH_WEST;
          }

          return OctantSlot.UP_SOUTH_WEST;
        } else if (z <= center.getZ()) {
          return OctantSlot.DOWN_NORTH_WEST;
        }

        return OctantSlot.UP_NORTH_WEST;
      } else if (y <= center.getY()) {
        if (z <= center.getZ()) {
          return OctantSlot.DOWN_SOUTH_EAST;
        }

        return OctantSlot.UP_SOUTH_EAST;
      } else if (z <= center.getZ()) {
        return OctantSlot.DOWN_NORTH_EAST;
      }

      return OctantSlot.UP_NORTH_EAST;
    }

    return null;
  }

}

OctreeCube.java

package tk.zorgatone.simpleoctree;

import static tk.zorgatone.simpleoctree.OctreeBounds.middlePoint;
import static tk.zorgatone.simpleoctree.OctreeOctant.copyCubes;

import java.util.LinkedList;
import java.util.List;

final class OctreeCube implements Cube {

  private final OctreeBounds bounds;
  private final OctreeOctant octant;

  OctreeCube() {
    this(new OctreeBounds(), new OctreeOctant());
  }

  OctreeCube(OctreeBounds bounds, OctreeOctant octant) {
    this.bounds = bounds;
    this.octant = octant;
  }

  @SuppressWarnings("DuplicatedCode")
  static OctreeBounds divide(OctantSlot direction, OctreeBounds bounds) {
    var topRight = bounds.getTopRight();
    var bottomLeft = bounds.getBottomLeft();
    var center = middlePoint(topRight, bottomLeft);

    return switch (direction) {
      case UP_NORTH_EAST -> new OctreeBounds(topRight, center);
      case UP_NORTH_WEST -> new OctreeBounds(
        new OctreePoint(
          center.getX(),
          topRight.getY(),
          topRight.getZ()
        ),
        new OctreePoint(
          bottomLeft.getX(),
          center.getY() + 1,
          center.getZ() + 1
        )
      );
      case UP_SOUTH_WEST -> new OctreeBounds(
        new OctreePoint(
          center.getX(),
          center.getY(),
          topRight.getZ()
        ),
        new OctreePoint(
          bottomLeft.getX(),
          bottomLeft.getY(),
          center.getZ() + 1
        )
      );
      case UP_SOUTH_EAST -> new OctreeBounds(
        new OctreePoint(
          topRight.getX(),
          center.getY(),
          topRight.getZ()
        ),
        new OctreePoint(
          center.getX() + 1,
          bottomLeft.getY(),
          center.getZ() + 1
        )
      );
      case DOWN_SOUTH_EAST -> new OctreeBounds(
        new OctreePoint(
          topRight.getX(),
          center.getY(),
          center.getZ()
        ),
        new OctreePoint(
          center.getX() + 1,
          bottomLeft.getY(),
          bottomLeft.getZ()
        )
      );
      case DOWN_SOUTH_WEST -> new OctreeBounds(center, bottomLeft);
      case DOWN_NORTH_WEST -> new OctreeBounds(
        new OctreePoint(
          center.getX(),
          topRight.getY(),
          center.getZ()
        ),
        new OctreePoint(
          bottomLeft.getX(),
          center.getY() + 1,
          bottomLeft.getZ()
        )
      );
      case DOWN_NORTH_EAST -> new OctreeBounds(
        new OctreePoint(
          topRight.getX(),
          topRight.getY(),
          center.getZ()
        ),
        new OctreePoint(
          center.getX() + 1,
          center.getY() + 1,
          bottomLeft.getZ()
        )
      );
    };
  }

  @Override
  public Octant getOctant() {
    return octant;
  }

  @Override
  public Bounds getBounds() {
    return bounds;
  }

  OctreeCube insert(Point point, Object object) {
    if (object instanceof Cube) {
      throw new IllegalStateException("Cannot insert cube into another directly!");
    }

    if (object instanceof ElementWrapper) {
      throw new IllegalStateException("Cannot insert ElementWrapper into a Cube!");
    }

    OctantSlot slot = bounds.directionOf(point);

    if (slot == null) {
      return this;
    }

    var content = octant.getSlot(slot);

    if (content == null) {
      Object[] elements = copyCubes(octant);
      elements[slot.getIndex()] = new ElementWrapper(point, object);

      return new OctreeCube(bounds, new OctreeOctant(elements));
    }

    if (content instanceof OctreeCube) {
      Object[] elements = copyCubes(octant);
      elements[slot.getIndex()] = ((OctreeCube) content).insert(point, object);

      return new OctreeCube(bounds, new OctreeOctant(elements));
    }

    if (content instanceof ElementWrapper) {
      var contentWrapper = (ElementWrapper) content;

      OctreeBounds newBounds = divide(slot, bounds);
      OctreeCube emptyCube = new OctreeCube(newBounds, new OctreeOctant());

      Object[] elements = copyCubes(octant);
      elements[slot.getIndex()] = emptyCube.insert(
        contentWrapper.getPoint(),
        contentWrapper.getObject()
      ).insert(
        point,
        object
      );

      return new OctreeCube(bounds, new OctreeOctant(elements));
    }

    throw new IllegalStateException("Content element is invalid!");
  }

  Object find(Point point) {
    if (point == null) {
      return null;
    }

    OctantSlot slot = bounds.directionOf(point);

    if (slot == null) {
      return null;
    }

    var content = octant.getSlot(slot);

    if (content == null) {
      return null;
    }

    if (content instanceof OctreeCube) {
      return ((OctreeCube) content).find(point);
    }

    if (content instanceof ElementWrapper) {
      return ((ElementWrapper) content).getObject();
    }

    throw new IllegalStateException("Content element is invalid!");
  }

  Iterable<Object> searchRadius(Point sphereCenter, double sphereRadius) {
    List<Object> list = new LinkedList<>();

    searchRadius(sphereCenter, sphereRadius, list);

    return list;
  }

  private void searchRadius(
    Point sphereCenter,
    double sphereRadius,
    List<Object> list
  ) {
    if (!bounds.intersectsRadius(sphereCenter, sphereRadius)) {
      return;
    }

    for (OctantSlot slot : OctantSlot.values()) {
      Object child = octant.getSlot(slot);

      if (child instanceof OctreeCube) {
        OctreeCube cube = (OctreeCube) child;

        cube.searchRadius(sphereCenter, sphereRadius, list);
      } else if (child instanceof ElementWrapper) {
        ElementWrapper wrapper = (ElementWrapper) child;

        assert bounds.within(wrapper.getPoint()) : "Child is in wrong node!";

        if (OctreeBounds.isPointInsideSphere(wrapper.getPoint(), sphereCenter, sphereRadius)) {
          list.add(wrapper.getObject());
        }
      } else if (child != null) {
        throw new IllegalStateException("Content element is invalid!");
      }
    }
  }

}

OctreeOctant.java

package tk.zorgatone.simpleoctree;

final class OctreeOctant implements Octant {

  private final Object[] cubes;

  static Object[] copyCubes(OctreeOctant octant) {
    Object[] cubes = new Object[MAX_ELEMENTS];

    if (octant == null || octant.cubes == null) {
      return cubes;
    }

    System.arraycopy(octant.cubes, 0, cubes, 0, MAX_ELEMENTS);

    return cubes;
  }

  OctreeOctant() {
    this(null);
  }

  OctreeOctant(Object[] elements) {
    if (elements == null) {
      cubes = null;
    } else if (elements.length > 0) {
      cubes = new Object[MAX_ELEMENTS];
      for (byte i = 0, len = (byte) Math.min(elements.length, MAX_ELEMENTS); i < len; ++i) {
        cubes[i] = elements[i];
      }
    } else {
      cubes = null;
    }
  }

  @Override
  public Object getSlot(OctantSlot slot) {
    if (cubes == null || slot == null) {
      return null;
    }

    return cubes[slot.getIndex()];
  }

}

OctreePoint.java

package tk.zorgatone.simpleoctree;

final class OctreePoint implements Point {

  private final long x;
  private final long y;
  private final long z;

  OctreePoint(long x, long y, long z) {
    this.x = x;
    this.y = y;
    this.z = z;
  }

  OctreePoint(Point point) {
    this.x = point.getX();
    this.y = point.getY();
    this.z = point.getZ();
  }

  @Override
  public long getX() {
    return x;
  }

  @Override
  public long getY() {
    return y;
  }

  @Override
  public long getZ() {
    return z;
  }

}

Point.java

package tk.zorgatone.simpleoctree;

public interface Point {
  long getX();
  long getY();
  long getZ();
}

SlotConversionException.java

package tk.zorgatone.simpleoctree;

import java.io.Serial;
import java.io.Serializable;

public class SlotConversionException extends Exception implements Serializable {

  @Serial
  private static final long serialVersionUID = -5804885729576243031L;

  public SlotConversionException() {
    super();
  }

  public SlotConversionException(String message) {
    super(message);
  }

  public SlotConversionException(String message, Throwable cause) {
    super(message, cause);
  }

  public SlotConversionException(Throwable cause) {
    super(cause);
  }

  public SlotConversionException(
    String message,
    Throwable cause,
    boolean enableSuppression,
    boolean writableStackTrace
  ) {
    super(message, cause, enableSuppression, writableStackTrace);
  }

  @Override
  public int hashCode() {
    return super.hashCode();
  }

  @Override
  public boolean equals(Object obj) {
    return obj instanceof SlotConversionException && super.equals(obj);
  }

}