| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362 |
-
- const ARRAY_TYPES = [
- Int8Array, Uint8Array, Uint8ClampedArray, Int16Array, Uint16Array,
- Int32Array, Uint32Array, Float32Array, Float64Array
- ];
-
- /** @typedef {Int8ArrayConstructor | Uint8ArrayConstructor | Uint8ClampedArrayConstructor | Int16ArrayConstructor | Uint16ArrayConstructor | Int32ArrayConstructor | Uint32ArrayConstructor | Float32ArrayConstructor | Float64ArrayConstructor} TypedArrayConstructor */
- /** @typedef {Int8Array | Uint8Array | Uint8ClampedArray | Int16Array | Uint16Array | Int32Array | Uint32Array | Float32Array | Float64Array} TypedArray */
-
- const VERSION = 1; // serialized format version
- const HEADER_SIZE = 8;
-
- // Shared scratch stack for iterative DFS in range/within. Sized for the worst case:
- // 3 ints per frame * (treeHeight + 1), with treeHeight ≤ ceil(log2(2^32 / 3)) ≈ 31.
- const STACK = new Uint32Array(96);
-
- export default class KDBush {
-
- /**
- * Creates an index from raw `ArrayBuffer` data.
- * @param {ArrayBufferLike} data
- */
- static from(data) {
- // @ts-expect-error duck typing array buffers
- if (!data || data.byteLength === undefined || data.buffer) {
- throw new Error('Data must be an instance of ArrayBuffer or SharedArrayBuffer.');
- }
- const [magic, versionAndType] = new Uint8Array(data, 0, 2);
- if (magic !== 0xdb) {
- throw new Error('Data does not appear to be in a KDBush format.');
- }
- const version = versionAndType >> 4;
- if (version !== VERSION) {
- throw new Error(`Got v${version} data when expected v${VERSION}.`);
- }
- const ArrayType = ARRAY_TYPES[versionAndType & 0x0f];
- if (!ArrayType) {
- throw new Error('Unrecognized array type.');
- }
- const [nodeSize] = new Uint16Array(data, 2, 1);
- const [numItems] = new Uint32Array(data, 4, 1);
-
- return new KDBush(numItems, nodeSize, ArrayType, undefined, data);
- }
-
- /**
- * Creates an index that will hold a given number of items.
- * @param {number} numItems
- * @param {number} [nodeSize=64] Size of the KD-tree node (64 by default).
- * @param {TypedArrayConstructor} [ArrayType=Float64Array] The array type used for coordinates storage (`Float64Array` by default).
- * @param {ArrayBufferConstructor | SharedArrayBufferConstructor} [ArrayBufferType=ArrayBuffer] The array buffer type used for storage (`ArrayBuffer` by default).
- * @param {ArrayBufferLike} [data] (For internal use only)
- */
- constructor(numItems, nodeSize = 64, ArrayType = Float64Array, ArrayBufferType = ArrayBuffer, data) {
- if (isNaN(numItems) || numItems < 0) throw new Error(`Unexpected numItems value: ${numItems}.`);
-
- this.numItems = +numItems;
- this.nodeSize = Math.min(Math.max(+nodeSize, 2), 65535);
- this.ArrayType = ArrayType;
- this.IndexArrayType = numItems < 65536 ? Uint16Array : Uint32Array;
-
- const arrayTypeIndex = ARRAY_TYPES.indexOf(this.ArrayType);
- const coordsByteSize = numItems * 2 * this.ArrayType.BYTES_PER_ELEMENT;
- const idsByteSize = numItems * this.IndexArrayType.BYTES_PER_ELEMENT;
- const padCoords = (8 - idsByteSize % 8) % 8;
-
- if (arrayTypeIndex < 0) {
- throw new Error(`Unexpected typed array class: ${ArrayType}.`);
- }
-
- if (data) { // reconstruct an index from a buffer
- this.data = data;
- // @ts-expect-error TS can't handle SharedArrayBuffer overloads
- this.ids = new this.IndexArrayType(data, HEADER_SIZE, numItems);
- // @ts-expect-error TS can't handle SharedArrayBuffer overloads
- this.coords = new ArrayType(data, HEADER_SIZE + idsByteSize + padCoords, numItems * 2);
- this._pos = numItems * 2;
- this._finished = true;
-
- } else { // initialize a new index
- const data = this.data = new ArrayBufferType(HEADER_SIZE + coordsByteSize + idsByteSize + padCoords);
- // @ts-expect-error TS can't handle SharedArrayBuffer overloads
- this.ids = new this.IndexArrayType(data, HEADER_SIZE, numItems);
- // @ts-expect-error TS can't handle SharedArrayBuffer overloads
- this.coords = new ArrayType(data, HEADER_SIZE + idsByteSize + padCoords, numItems * 2);
- this._pos = 0;
- this._finished = false;
-
- // set header
- new Uint8Array(data, 0, 2).set([0xdb, (VERSION << 4) + arrayTypeIndex]);
- new Uint16Array(data, 2, 1)[0] = nodeSize;
- new Uint32Array(data, 4, 1)[0] = numItems;
- }
- }
-
- /**
- * Add a point to the index.
- * @param {number} x
- * @param {number} y
- * @returns {number} An incremental index associated with the added item (starting from `0`).
- */
- add(x, y) {
- const index = this._pos >> 1;
- this.ids[index] = index;
- this.coords[this._pos++] = x;
- this.coords[this._pos++] = y;
- return index;
- }
-
- /**
- * Perform indexing of the added points.
- */
- finish() {
- const numAdded = this._pos >> 1;
- if (numAdded !== this.numItems) {
- throw new Error(`Added ${numAdded} items when expected ${this.numItems}.`);
- }
- // kd-sort both arrays for efficient search
- sort(this.ids, this.coords, this.nodeSize, 0, this.numItems - 1, 0);
-
- this._finished = true;
- return this;
- }
-
- /**
- * Search the index for items within a given bounding box.
- * @param {number} minX
- * @param {number} minY
- * @param {number} maxX
- * @param {number} maxY
- * @returns {number[]} An array of indices correponding to the found items.
- */
- range(minX, minY, maxX, maxY) {
- if (!this._finished) throw new Error('Data not yet indexed - call index.finish().');
-
- const {ids, coords, nodeSize} = this;
- STACK[0] = 0;
- STACK[1] = ids.length - 1;
- STACK[2] = 0;
- let sp = 3;
- const result = [];
-
- // recursively search for items in range in the kd-sorted arrays
- while (sp > 0) {
- const axis = STACK[--sp];
- const right = STACK[--sp];
- const left = STACK[--sp];
-
- // if we reached "tree node", search linearly
- if (right - left <= nodeSize) {
- for (let i = left; i <= right; i++) {
- const x = coords[2 * i];
- const y = coords[2 * i + 1];
- if (x >= minX && x <= maxX && y >= minY && y <= maxY) result.push(ids[i]);
- }
- continue;
- }
-
- // otherwise find the middle index
- const m = (left + right) >> 1;
-
- // include the middle item if it's in range
- const x = coords[2 * m];
- const y = coords[2 * m + 1];
- if (x >= minX && x <= maxX && y >= minY && y <= maxY) result.push(ids[m]);
-
- // queue search in halves that intersect the query
- if (axis === 0 ? minX <= x : minY <= y) {
- STACK[sp++] = left;
- STACK[sp++] = m - 1;
- STACK[sp++] = 1 - axis;
- }
- if (axis === 0 ? maxX >= x : maxY >= y) {
- STACK[sp++] = m + 1;
- STACK[sp++] = right;
- STACK[sp++] = 1 - axis;
- }
- }
-
- return result;
- }
-
- /**
- * Search the index for items within a given radius.
- * @param {number} qx
- * @param {number} qy
- * @param {number} r Query radius.
- * @returns {number[]} An array of indices correponding to the found items.
- */
- within(qx, qy, r) {
- const result = /** @type {number[]} */ ([]);
- this.withinInto(qx, qy, r, result);
- return result;
- }
-
- /**
- * Search the index for items within a given radius, writing matching ids into `out`
- * via indexed assignment (`out[i] = id`). Accepts any indexed-writable container —
- * a typed array sized to the expected upper bound (allocation-free, fast) or a plain
- * `Array` (which will grow as needed). Returns the number of matches written.
- * @param {number} qx
- * @param {number} qy
- * @param {number} r Query radius.
- * @param {number[] | TypedArray} out Container to write matching ids into.
- * @returns {number} The number of matches written to `out`.
- */
- withinInto(qx, qy, r, out) {
- if (!this._finished) throw new Error('Data not yet indexed - call index.finish().');
-
- const {ids, coords, nodeSize} = this;
- STACK[0] = 0;
- STACK[1] = ids.length - 1;
- STACK[2] = 0;
- let sp = 3;
- let count = 0;
- const r2 = r * r;
-
- // recursively search for items within radius in the kd-sorted arrays
- while (sp > 0) {
- const axis = STACK[--sp];
- const right = STACK[--sp];
- const left = STACK[--sp];
-
- // if we reached "tree node", search linearly
- if (right - left <= nodeSize) {
- for (let i = left; i <= right; i++) {
- if (sqDist(coords[2 * i], coords[2 * i + 1], qx, qy) <= r2) out[count++] = ids[i];
- }
- continue;
- }
-
- // otherwise find the middle index
- const m = (left + right) >> 1;
-
- // include the middle item if it's in range
- const x = coords[2 * m];
- const y = coords[2 * m + 1];
- if (sqDist(x, y, qx, qy) <= r2) out[count++] = ids[m];
-
- // queue search in halves that intersect the query
- if (axis === 0 ? qx - r <= x : qy - r <= y) {
- STACK[sp++] = left;
- STACK[sp++] = m - 1;
- STACK[sp++] = 1 - axis;
- }
- if (axis === 0 ? qx + r >= x : qy + r >= y) {
- STACK[sp++] = m + 1;
- STACK[sp++] = right;
- STACK[sp++] = 1 - axis;
- }
- }
-
- return count;
- }
- }
-
- /**
- * @param {Uint16Array | Uint32Array} ids
- * @param {TypedArray} coords
- * @param {number} nodeSize
- * @param {number} left
- * @param {number} right
- * @param {number} axis
- */
- function sort(ids, coords, nodeSize, left, right, axis) {
- if (right - left <= nodeSize) return;
-
- const m = (left + right) >> 1; // middle index
-
- // sort ids and coords around the middle index so that the halves lie
- // either left/right or top/bottom correspondingly (taking turns)
- select(ids, coords, m, left, right, axis);
-
- // recursively kd-sort first half and second half on the opposite axis
- sort(ids, coords, nodeSize, left, m - 1, 1 - axis);
- sort(ids, coords, nodeSize, m + 1, right, 1 - axis);
- }
-
- /**
- * Custom Floyd-Rivest selection algorithm: sort ids and coords so that
- * [left..k-1] items are smaller than k-th item (on either x or y axis)
- * @param {Uint16Array | Uint32Array} ids
- * @param {TypedArray} coords
- * @param {number} k
- * @param {number} left
- * @param {number} right
- * @param {number} axis
- */
- function select(ids, coords, k, left, right, axis) {
-
- while (right > left) {
- if (right - left > 600) {
- const n = right - left + 1;
- const m = k - left + 1;
- const z = Math.log(n);
- const s = 0.5 * Math.exp(2 * z / 3);
- const sd = 0.5 * Math.sqrt(z * s * (n - s) / n) * (m - n / 2 < 0 ? -1 : 1);
- const newLeft = Math.max(left, Math.floor(k - m * s / n + sd));
- const newRight = Math.min(right, Math.floor(k + (n - m) * s / n + sd));
- select(ids, coords, k, newLeft, newRight, axis);
- }
-
- const t = coords[2 * k + axis];
- let i = left;
- let j = right;
-
- swapItem(ids, coords, left, k);
- if (coords[2 * right + axis] > t) swapItem(ids, coords, left, right);
-
- while (i < j) {
- swapItem(ids, coords, i, j);
- i++;
- j--;
- while (coords[2 * i + axis] < t) i++;
- while (coords[2 * j + axis] > t) j--;
- }
-
- if (coords[2 * left + axis] === t) swapItem(ids, coords, left, j);
- else {
- j++;
- swapItem(ids, coords, j, right);
- }
-
- if (j <= k) left = j + 1;
- if (k <= j) right = j - 1;
- }
- }
-
- /**
- * @param {Uint16Array | Uint32Array} ids
- * @param {TypedArray} coords
- * @param {number} i
- * @param {number} j
- */
- function swapItem(ids, coords, i, j) {
- swap(ids, i, j);
- swap(coords, 2 * i, 2 * j);
- swap(coords, 2 * i + 1, 2 * j + 1);
- }
-
- /**
- * @param {TypedArray} arr
- * @param {number} i
- * @param {number} j
- */
- function swap(arr, i, j) {
- const tmp = arr[i];
- arr[i] = arr[j];
- arr[j] = tmp;
- }
-
- /**
- * @param {number} ax
- * @param {number} ay
- * @param {number} bx
- * @param {number} by
- */
- function sqDist(ax, ay, bx, by) {
- const dx = ax - bx;
- const dy = ay - by;
- return dx * dx + dy * dy;
- }
|