Spatial index

/**
 * A spatial index over rectangles: which shapes does this point, this marquee,
 * this viewport touch? A packed R-tree, built once from all the rectangles,
 * sorted along a Hilbert curve so that rectangles near each other on the canvas
 * are near each other in memory.
 *
 * The whole tree is two typed arrays. A search touches only the nodes whose
 * boxes overlap the query, and allocates nothing but its results: the nodes it
 * has still to open wait on one stack that every search shares.
 *
 * A rectangle or a query may be given with either corner first, may have no
 * area, and may reach to infinity; a side that is not a number is a mistake in
 * the caller and throws a `RangeError`. An index holds up to 2^28 rectangles,
 * so that every position in it is a 32-bit integer.
 */

export interface SpatialIndex {
  /** Every node's box, `[minX, minY, maxX, maxY]`: leaves first, the root last. */
  readonly boxes: Float64Array;
  /** For a leaf, the rectangle's number; for a branch, where its children start in `boxes`. */
  readonly indices: Uint32Array;
  /** Where each level of the tree ends in `boxes`. */
  readonly levels: readonly number[];
  readonly nodeSize: number;
  readonly count: number;
}

const HILBERT_MAX = 65535;
const MAX_COUNT = 2 ** 28;

// Positions of nodes still to open, for every search in progress: one begun
// inside another works above what the outer one has pushed.
let stack = new Uint32Array(64);
let top = 0;

/** The distance along a Hilbert curve of order 16 at which the cell (x, y) lies. */
function hilbert(x: number, y: number): number {
  let distance = 0;
  for (let side = 32768; side > 0; side = Math.floor(side / 2)) {
    const right = (x & side) > 0 ? 1 : 0;
    const up = (y & side) > 0 ? 1 : 0;
    distance += side * side * ((3 * right) ^ up);
    if (up === 0) {
      if (right === 1) {
        x = HILBERT_MAX - x;
        y = HILBERT_MAX - y;
      }
      const swap = x;
      x = y;
      y = swap;
    }
  }
  return distance;
}

/**
 * The index for rectangles given flat, `[minX0, minY0, maxX0, maxY0, minX1, …]`,
 * with at least two boxes to a node. O(n log n).
 */
export function spatialIndexFrom(rectangles: ArrayLike<number>, nodeSize = 16): SpatialIndex {
  if (!Number.isInteger(nodeSize) || nodeSize < 2) {
    // One box to a node would never narrow to a root.
    throw new RangeError(`A node holds a whole number of boxes, two or more, not ${nodeSize}.`);
  }
  const count = rectangles.length / 4;
  if (!Number.isInteger(count)) {
    throw new RangeError(`Rectangles are four numbers each; ${rectangles.length} is not a multiple of four.`);
  }
  if (count > MAX_COUNT) throw new RangeError(`A spatial index holds at most 2^28 rectangles, not ${count}.`);
  // Each level holds a node per `nodeSize` of the level below, up to one root.
  const levels: number[] = [count * 4];
  let nodes = count;
  let width = count;
  while (width > 1) {
    width = Math.ceil(width / nodeSize);
    nodes += width;
    levels.push(nodes * 4);
  }
  const boxes = new Float64Array(Math.max(nodes, 1) * 4);
  const indices = new Uint32Array(Math.max(nodes, 1));
  if (count === 0) {
    boxes.fill(Number.NaN);
    return { boxes, indices, levels, nodeSize, count };
  }

  // The extent of the centres that are somewhere: a rectangle reaching to
  // infinity has none, and must not flatten the curve for the rest.
  let minX = Infinity;
  let minY = Infinity;
  let maxX = -Infinity;
  let maxY = -Infinity;
  for (let i = 0; i < rectangles.length; i += 4) {
    const x0 = rectangles[i]!;
    const y0 = rectangles[i + 1]!;
    const x1 = rectangles[i + 2]!;
    const y1 = rectangles[i + 3]!;
    // A box with such a side would be passed by no query, and so would its parents.
    if (x0 !== x0 || y0 !== y0 || x1 !== x1 || y1 !== y1) {
      throw new RangeError(`Rectangle ${i / 4} has a side that is not a number.`);
    }
    const x = (x0 + x1) / 2;
    const y = (y0 + y1) / 2;
    // Only a finite number less itself is zero.
    if (x - x === 0) {
      minX = Math.min(minX, x);
      maxX = Math.max(maxX, x);
    }
    if (y - y === 0) {
      minY = Math.min(minY, y);
      maxY = Math.max(maxY, y);
    }
  }

  // Order the rectangles by where their centres fall along the curve.
  const scaleX = HILBERT_MAX / (maxX - minX || 1);
  const scaleY = HILBERT_MAX / (maxY - minY || 1);
  const distances = new Float64Array(count);
  for (let item = 0; item < count; item += 1) {
    const x = Math.floor(scaleX * ((rectangles[item * 4]! + rectangles[item * 4 + 2]!) / 2 - minX));
    const y = Math.floor(scaleY * ((rectangles[item * 4 + 1]! + rectangles[item * 4 + 3]!) / 2 - minY));
    // A centre that is nowhere becomes cell 0 under `&`.
    distances[item] = hilbert(x, y);
  }
  // Sorted in place: the array was made here, and a copy would only be thrown away.
  // oxlint-disable-next-line unicorn/no-array-sort
  const sorted = Uint32Array.from({ length: count }, (_, item) => item).sort(
    (a, b) => distances[a]! - distances[b]!,
  );
  for (let position = 0; position < count; position += 1) {
    const item = sorted[position]!;
    indices[position] = item;
    // Either corner may come first; a box is stored smallest first.
    boxes[position * 4] = Math.min(rectangles[item * 4]!, rectangles[item * 4 + 2]!);
    boxes[position * 4 + 1] = Math.min(rectangles[item * 4 + 1]!, rectangles[item * 4 + 3]!);
    boxes[position * 4 + 2] = Math.max(rectangles[item * 4]!, rectangles[item * 4 + 2]!);
    boxes[position * 4 + 3] = Math.max(rectangles[item * 4 + 1]!, rectangles[item * 4 + 3]!);
  }

  // Each level's nodes are the unions of runs of `nodeSize` boxes below them.
  let write = count * 4;
  let read = 0;
  for (let level = 0; level < levels.length - 1; level += 1) {
    const end = levels[level]!;
    while (read < end) {
      const first = read;
      let nodeMinX = Infinity;
      let nodeMinY = Infinity;
      let nodeMaxX = -Infinity;
      let nodeMaxY = -Infinity;
      for (let child = 0; child < nodeSize && read < end; child += 1, read += 4) {
        nodeMinX = Math.min(nodeMinX, boxes[read]!);
        nodeMinY = Math.min(nodeMinY, boxes[read + 1]!);
        nodeMaxX = Math.max(nodeMaxX, boxes[read + 2]!);
        nodeMaxY = Math.max(nodeMaxY, boxes[read + 3]!);
      }
      indices[write / 4] = first;
      boxes[write] = nodeMinX;
      boxes[write + 1] = nodeMinY;
      boxes[write + 2] = nodeMaxX;
      boxes[write + 3] = nodeMaxY;
      write += 4;
    }
  }
  return { boxes, indices, levels, nodeSize, count };
}

/** The index for rectangles given as corner and size, as an interface holds them. O(n log n). */
export function spatialIndexFromRects(
  rects: ArrayLike<{
    readonly x: number;
    readonly y: number;
    readonly width: number;
    readonly height: number;
  }>,
  nodeSize = 16,
): SpatialIndex {
  const flat = new Float64Array(rects.length * 4);
  for (let item = 0; item < rects.length; item += 1) {
    const { x, y, width, height } = rects[item]!;
    flat[item * 4] = x;
    flat[item * 4 + 1] = y;
    flat[item * 4 + 2] = x + width;
    flat[item * 4 + 3] = y + height;
  }
  return spatialIndexFrom(flat, nodeSize);
}

/**
 * The numbers of the rectangles that touch the query box, appended to `into`,
 * in no particular order. Touching edges count, so a point finds what it lies
 * on the edge of.
 */
export function spatialSearch(
  { boxes, indices, levels, nodeSize, count }: SpatialIndex,
  minX: number,
  minY: number,
  maxX: number,
  maxY: number,
  into: number[] = [],
): number[] {
  if (minX > maxX) {
    const swap = minX;
    minX = maxX;
    maxX = swap;
  }
  if (minY > maxY) {
    const swap = minY;
    minY = maxY;
    maxY = swap;
  }
  // Only a side that is not a number is left out of order by the swaps.
  if (!(minX <= maxX && minY <= maxY)) throw new RangeError("A query has a side that is not a number.");
  if (count === 0) return into;
  const leaves = count * 4;
  const base = top;
  // Every level may leave a node's worth of boxes waiting.
  const room = base + levels.length * nodeSize;
  if (stack.length < room) {
    const grown = new Uint32Array(Math.max(room, stack.length * 2));
    grown.set(stack);
    stack = grown;
  }
  // A search begun inside this one may replace `stack` with a longer copy; this
  // one has made its room and keeps to the array it has.
  const pending = stack;
  let node = boxes.length - 4;
  // The root's own box is tested first, as a node of one.
  let end = boxes.length;
  try {
    for (;;) {
      for (let position = node; position < end; position += 4) {
        if (
          maxX < boxes[position]! ||
          maxY < boxes[position + 1]! ||
          minX > boxes[position + 2]! ||
          minY > boxes[position + 3]!
        ) {
          continue;
        }
        if (position < leaves) into.push(indices[position / 4]!);
        else pending[top++] = indices[position / 4]!;
      }
      if (top === base) return into;
      node = pending[--top]!;
      // A node's children run to `nodeSize` boxes, or to the end of their level.
      let level = 0;
      while (levels[level]! <= node) level += 1;
      end = Math.min(node + nodeSize * 4, levels[level]!);
    }
  } finally {
    // Whatever was thrown from `into` or `boxes`, later searches start clean.
    top = base;
  }
}