axis-index-to-vector3-property.ts

const AXES = ["X" as const, "Y" as const, "Z" as const];

export default function axisIndexToVector3Property(axisIndex: number) {
	return AXES[axisIndex];
}

k-d-node.ts

import axisIndexToVector3Property from "./axis-index-to-vector3-property";

/**
 * @hidden
 */
export default class KDNode<T> {
	public left?: KDNode<T>;
	public right?: KDNode<T>;

	public readonly object: T;
	public readonly position: Vector3;

	public constructor(position: Vector3, object: T) {
		this.object = object;
		this.position = position;
	}

	public insert(position: Vector3, object: T, depth: number) {
		const axis = axisIndexToVector3Property(depth % 3);
		if (position[axis] < this.position[axis]) {
			if (!this.left) this.left = new KDNode(position, object);
			else this.left.insert(position, object, depth + 1);
		} else {
			if (!this.right) this.right = new KDNode(position, object);
			else this.right.insert(position, object, depth + 1);
		}
	}

	public destroy() {
		this.left?.destroy();
		this.right?.destroy();
		table.clear(this);
		setmetatable(this, undefined!);
	}
}

kd-tree.ts

import KDNode from "./k-d-node";
import axisIndexToVector3Property from "./axis-index-to-vector3-property";

export default class KDTree<ObjectType extends defined> {
	public insert(position: Vector3, data: ObjectType) {
		if (!this.root) this.root = new KDNode(position, data);
		else this.root.insert(position, data, 0);
	}

	public getNearest(position: Vector3, maxDistance: number): KDNode<ObjectType> | undefined {
		if (!this.root) return undefined;

		let bestNode: KDNode<ObjectType> | undefined = undefined;
		let bestDistance = math.huge;

		function search(node: KDNode<ObjectType> | undefined, depth: number) {
			if (!node) return;
			const distance = position.sub(node.position).Magnitude * 2;
			if (distance < bestDistance) {
				bestDistance = distance;
				bestNode = node;
			}

			const axis = axisIndexToVector3Property(depth % 3);
			const axisDistance = position[axis] - node.position[axis];

			if (axisDistance < 0) {
				search(node.left, depth + 1);
				if (axisDistance * axisDistance < bestDistance) search(node.right, depth + 1);
			} else {
				search(node.right, depth + 1);
				if (axisDistance * axisDistance < bestDistance) search(node.left, depth + 1);
			}
		}

		search(this.root, 0);
		if (bestDistance > maxDistance * maxDistance) return undefined;
		return bestNode;
	}

	public getNearestObject(position: Vector3, maxDistance: number): ObjectType | undefined {
		return this.getNearest(position, maxDistance)?.object;
	}

	public destroy() {
		this.root?.destroy();
		table.clear(this);
		setmetatable(this, undefined!);
	}

	private root?: KDNode<ObjectType>;
}