rbush-3d-knn.js

(function (global, factory) {
    typeof exports === 'object' && typeof module !== 'undefined' ? module.exports = factory() :
    typeof define === 'function' && define.amd ? define(factory) :
    (global = typeof globalThis !== 'undefined' ? globalThis : global || self, global["rbush-3d-knn"] = factory());
})(this, (function () { 'use strict';

    class TinyQueue {
        constructor(data = [], compare = defaultCompare) {
            this.data = data;
            this.length = this.data.length;
            this.compare = compare;

            if (this.length > 0) {
                for (let i = (this.length >> 1) - 1; i >= 0; i--) this._down(i);
            }
        }

        push(item) {
            this.data.push(item);
            this.length++;
            this._up(this.length - 1);
        }

        pop() {
            if (this.length === 0) return undefined;

            const top = this.data[0];
            const bottom = this.data.pop();
            this.length--;

            if (this.length > 0) {
                this.data[0] = bottom;
                this._down(0);
            }

            return top;
        }

        peek() {
            return this.data[0];
        }

        _up(pos) {
            const {data, compare} = this;
            const item = data[pos];

            while (pos > 0) {
                const parent = (pos - 1) >> 1;
                const current = data[parent];
                if (compare(item, current) >= 0) break;
                data[pos] = current;
                pos = parent;
            }

            data[pos] = item;
        }

        _down(pos) {
            const {data, compare} = this;
            const halfLength = this.length >> 1;
            const item = data[pos];

            while (pos < halfLength) {
                let left = (pos << 1) + 1;
                let best = data[left];
                const right = left + 1;

                if (right < this.length && compare(data[right], best) < 0) {
                    left = right;
                    best = data[right];
                }
                if (compare(best, item) >= 0) break;

                data[pos] = best;
                pos = left;
            }

            data[pos] = item;
        }
    }

    function defaultCompare(a, b) {
        return a < b ? -1 : a > b ? 1 : 0;
    }

    function knn(tree, x, y, z, n, predicate, maxDistance) {
        var node = tree.data,
            result = [],
            i, child, dist, candidate;

        var queue = new TinyQueue(undefined, compareDist);

        while (node) {
            for (i = 0; i < node.children.length; i++) {
                child = node.children[i];
                dist = boxDist(x, y, z, node.children[i]);
                // console.log(dist);
                if (!maxDistance || dist <= maxDistance) {
                    queue.push({
                        node: child,
                        isItem: node.leaf,
                        dist: dist
                    });
                }
            }

            while (queue.length && queue.peek().isItem) {
                candidate = queue.pop().node;
                if (!predicate || predicate(candidate))
                    result.push(candidate);
                if (n && result.length === n) return result;
            }

            node = queue.pop();
            if (node) node = node.node;
        }

        return result;
    }

    function compareDist(a, b) {
        return a.dist - b.dist;
    }

    function boxDist(x, y, z, box) {
        var dx = axisDist(x, box.minX, box.maxX),
            dy = axisDist(y, box.minY, box.maxY),
            dz = axisDist(z, box.minZ, box.maxZ);
        return Math.sqrt(dx * dx + dy * dy + dz * dz);
    }

    function axisDist(k, min, max) {
        return k < min ? min - k : k <= max ? 0 : k - max;
    }

    return knn;

}));