spatial-grid3.ts

// SPDX-License-Identifier: Apache-2.0
import type { Fn, Nullable, Pair } from "@thi.ng/api";
import { isNumber } from "@thi.ng/checks/is-number";
import { Heap } from "@thi.ng/heaps/heap";
import { clamp } from "@thi.ng/math/interval";
import type { ReadonlyVec, Vec } from "@thi.ng/vectors";
import { addN3 } from "@thi.ng/vectors/addn";
import { distSq3 } from "@thi.ng/vectors/distsq";
import { subN3 } from "@thi.ng/vectors/subn";
import { ASpatialGrid } from "./aspatial-grid.js";
import { __addResults, CMP } from "./utils.js";

const TMP: Vec = [];

export class SpatialGrid3<K extends ReadonlyVec, V> extends ASpatialGrid<K, V> {
	protected _stride: ReadonlyVec;

	constructor(
		min: ReadonlyVec,
		size: ReadonlyVec,
		res: ReadonlyVec | number
	) {
		super(min, size, isNumber(res) ? [res, res, res] : res);
		this._cells = new Array(this._res[0] * this._res[1] * this._res[2]);
		this._stride = [this._res[0], this._res[0] * this._res[1]];
	}

	copy(): SpatialGrid3<K, V> {
		return <SpatialGrid3<K, V>>super.copy();
	}

	empty() {
		return new SpatialGrid3<K, V>(this._min, this._size, this._res);
	}

	protected doQuery<T>(
		fn: Fn<Pair<K, V>, T>,
		k: K,
		r: number,
		limit = Infinity,
		acc: T[] = []
	) {
		const id1 = this.findIndex(subN3(TMP, k, r));
		const id2 = this.findIndex(addN3(TMP, k, r));
		const [width, slice] = this._stride;
		const x1 = id1 % width;
		const x2 = id2 % width;
		const y1 = ((id1 / width) | 0) * width;
		const y2 = ((id2 / width) | 0) * width;
		const z1 = ((id1 / slice) | 0) * slice;
		const z2 = ((id2 / slice) | 0) * slice;
		const cells = this._cells;
		let c: Nullable<Pair<K, V>[]>;
		let x: number, y: number, z: number;
		r *= r;
		const heap = new Heap<[number, Nullable<Pair<K, V>>?]>([[r]], {
			compare: CMP,
		});
		const sel = heap.values;
		for (z = z1; z <= z2; z += slice) {
			for (y = y1; y <= y2; y += width) {
				for (x = x1; x <= x2; x++) {
					c = cells[z + y + x];
					c?.length && this.queryCell(distSq3, heap, c, k, limit);
				}
			}
		}
		return __addResults(fn, sel, acc);
	}

	protected findIndex(k: ReadonlyVec) {
		const { _min: min, _res1: res1, _invSize: invSize } = this;
		const kx = clamp((k[0] - min[0]) * invSize[0], 0, res1[0]);
		const ky = clamp((k[1] - min[1]) * invSize[1], 0, res1[1]);
		const kz = clamp((k[2] - min[2]) * invSize[2], 0, res1[2]);
		return (
			(kx | 0) + (ky | 0) * this._stride[0] + (kz | 0) * this._stride[1]
		);
	}
}