fqtree/src/quadtree.ts

a flexible quadtree for JavaScript/TypeScript

import { Option, is_none } from 'onsreo';
import { PointLike, nullify } from 'evjkit';
import { IRange, RectangleRange } from './range';
import { DataCache, get_data_cache, get_qt_cache } from './caches';

export declare interface QtSpatial extends PointLike {
  id: number;
}
export declare type QtDataTuple<T> = [data: T, qt_id: number];

let _next_qt_id = 0;

export class QuadTree<Spatial extends QtSpatial> {
  id = -1;
  boundary: RectangleRange;
  capacity = 4;
  data: Array<number>;
  divided = false;
  maxDepth: number;
  root: Option<number>;
  ne!: Option<number>;
  nw!: Option<number>;
  se!: Option<number>;
  sw!: Option<number>;

  /**
   * A Quadtree for spatial partitioning and searching.
   * @param boundary - this QuadTree<Spatial>'s bounds
   * @param capacity - how many points this quadtree can hold
   * @param root - QuadTree<Spatial>'s root Instance (only used for subdivisions)
   */
  constructor(
    boundary: RectangleRange,
    capacity = 4,
    maxDepth = 10,
    root?: Option<number>,
  ) {
    this.boundary = boundary;
    this.capacity = capacity;
    this.id = _next_qt_id++;
    this.root = root ?? this.id;
    this.maxDepth = maxDepth;
    this.data = [];
    const qt_cache = get_qt_cache();
    qt_cache.add(this, this.id);
  }

  /**
   * clears this quadtree and its children (if any).
   */
  _clear(): void {
    this.data.forEach(id => this.data_cache.remove(id));
    this.data = [];
    this.qt_cache.remove(this.id);
    this.root = null;
    if (this.divided) {
      this.qt_cache.get(this.ne)?._clear();
      this.qt_cache.get(this.nw)?._clear();
      this.qt_cache.get(this.se)?._clear();
      this.qt_cache.get(this.sw)?._clear();
    }
    this.ne = null;
    this.nw = null;
    this.se = null;
    this.sw = null;
    this.divided = false;
    nullify(this, 'boundary');
  }

  destroy() {
    this._clear();
    _next_qt_id = 0;
    this.data_cache.destroy();
    this.qt_cache.destroy();
  }

  get qt_cache(): DataCache<QuadTree<Spatial>> {
    return get_qt_cache();
  }
  get data_cache(): DataCache<QtDataTuple<Spatial>> {
    return get_data_cache();
  }

  // /**
  //  * resets all the internal data for this tree and its children.
  //  */
  // _reset(): void {
  //   this.data.forEach(id => this.data_cache.remove(id));
  //   this.data = [];
  //   this.qt_cache.remove(this.id);
  //   this.root = null;
  //   if (this.divided) {
  //     this.qt_cache.get(this.ne)?._reset();
  //     this.qt_cache.get(this.nw)?._reset();
  //     this.qt_cache.get(this.se)?._reset();
  //     this.qt_cache.get(this.sw)?._reset();
  //   }
  //   this.ne = null;
  //   this.nw = null;
  //   this.se = null;
  //   this.sw = null;
  //   this.divided = false;
  //   nullify(this, 'boundary');
  // }

  /**
   * resets the quadtree and optionally sets new bounds.
   * @param boundary - this QuadTree<Spatial>'s bounds
   */
  reset(bounds: RectangleRange = this.boundary): void {
    // store bounds temporarily so they can be persisted
    const tempBounds = bounds;
    this._clear();
    _next_qt_id = 0;
    this.id = _next_qt_id++;
    this.root = this.id;
    this.qt_cache.add(this, this.id);
    this.boundary = tempBounds;
  }

  /**
   * subdivide this tree into 4 new quadrants.
   */
  subdivide(): void {
    const { x, y, w, h } = this.boundary.bounds;
    const hw = Math.ceil(w / 2),
      hh = Math.ceil(h / 2);
    const ne = new RectangleRange(x + hw, y - hh, hw, hh);
    const nw = new RectangleRange(x - hw, y - hh, hw, hh);
    const se = new RectangleRange(x + hw, y + hh, hw, hh);
    const sw = new RectangleRange(x - hw, y + hh, hw, hh);
    // we insert things this way so that the points will cause the
    // appropriate number of subdivisions, that way ranges and lines
    // will be inserted most efficiently. We don't do this separation
    // in insertMany as it would have too large of a performance
    // penalty
    this.ne = new QuadTree<Spatial>(
      ne,
      this.capacity,
      this.maxDepth - 1,
      this.root,
    ).id;
    this.nw = new QuadTree<Spatial>(
      nw,
      this.capacity,
      this.maxDepth - 1,
      this.root,
    ).id;
    this.se = new QuadTree<Spatial>(
      se,
      this.capacity,
      this.maxDepth - 1,
      this.root,
    ).id;
    this.sw = new QuadTree<Spatial>(
      sw,
      this.capacity,
      this.maxDepth - 1,
      this.root,
    ).id;
    // now we've successfully divided
    this.divided = true;
    // now collect and re-insert this quadrants data
    // into its sub-quadrants
    this.getData().forEach(spatial => {
      this.data_cache.remove(spatial.id);
      this.insert(spatial);
    });
    // clear this node's internal data as we only want it
    // to exist on the leaves
    this.data = [];
  }

  _insert(spatial: Spatial) {
    // insert and divide based on type of spatial
    this.data.push(spatial.id);
    // if we've exceeded capacity, subdivide
    if (this.data.length > this.capacity && this.maxDepth) {
      this.subdivide();
    } else {
      this.data_cache.add([spatial, this.id], spatial.id);
    }
  }

  /**
   * insert a spatial into this tree.
   * @param spatial - spatial to insert.
   */
  insert(spatial: Spatial): void {
    if (!this.contains(spatial)) return;
    if (this.divided) {
      this.qt_cache.get(this.ne)?.insert(spatial);
      this.qt_cache.get(this.nw)?.insert(spatial);
      this.qt_cache.get(this.se)?.insert(spatial);
      this.qt_cache.get(this.sw)?.insert(spatial);
    } else {
      this._insert(spatial);
    }
  }

  /**
   * insert all the provided spatials.
   * @param spatials - spatials to insert.
   */
  insertMany(spatials: Set<Spatial> | Array<Spatial>): void {
    spatials.forEach((spatial: Spatial) => this.insert(spatial));
  }

  /**
   * tests if this quadtree's boundary completely contains the spatial.
   * @param spatial - spatial to test.
   * @returns true if contained within the bounds, else false.
   */
  contains(spatial: Spatial): boolean {
    return this.boundary.containsPoint(spatial);
  }

  /**
   * tests if this quadtree's boundary intersects with the boundary.
   * @param spatial - spatial to test.
   * @returns true if intersects, else false.
   */
  intersects(range: IRange): boolean {
    return this.boundary.intersects(range);
  }

  getData(): Spatial[] {
    return this.data.reduce<Spatial[]>((data, id) => {
      const maybe_data = this.data_cache.get(id);
      if (is_none(maybe_data)) return data;
      return data.addItem(maybe_data[0]);
    }, []);
  }

  /**
   * query insertain spatial using provided range.
   * @param range - range for query
   * @param found - set to fill [default: empty Set].
   * @returns query results as a set
   */
  query(range: IRange, found: Set<Spatial> = new Set()): Set<Spatial> {
    if (!this.intersects(range)) return found;
    if (this.divided) {
      this.qt_cache.get(this.ne)?.query(range, found);
      this.qt_cache.get(this.nw)?.query(range, found);
      this.qt_cache.get(this.se)?.query(range, found);
      this.qt_cache.get(this.sw)?.query(range, found);
    } else {
      this.getData().forEach(
        spatial => range.containsPoint(spatial) && found.add(spatial),
      );
    }
    return found;
  }

  _remove(id: number): void {
    const maybe_data = this.data_cache.get(id);
    if (is_none(maybe_data)) return;
    const [data, parent_id] = maybe_data;
    this.qt_cache.get(parent_id)?.remove(data);
  }

  /**
   * remove the given spatial from the quadtree.
   * @param spatial - spatial to remove.
   */
  remove(spatial: Spatial): void {
    this._remove(spatial.id);
  }

  /**
   * remove a spatial by id.
   * @param id - id of spatial to remove.
   */
  removeById(id: number): void {
    this._remove(id);
  }

  /**
   * remove all spatials with the provided IDs.
   * @param removedIds - array of uuids to remove.
   */
  removeByIds(removedIds: Array<number>): void {
    removedIds.forEach(id => this._remove(id));
  }

  /**
   * remove many spatials from the quadtree.
   * @param spatials - spatials to remove.
   */
  removeMany(spatials: Set<Spatial> | Array<Spatial>): void {
    spatials.forEach(spatial => this._remove(spatial.id));
  }

  /**
   * update a spatial point.
   * @param spatial - spatial point to update.
   */
  update(spatial: Spatial): void {
    const maybe_data = this.data_cache.get(spatial.id);
    if (is_none(maybe_data)) {
      // if no parents, then we only need to insert at root
      this.qt_cache.get(this.root)?.insert(spatial);
      return;
    }
    const [_data, parent_id] = maybe_data;

    const maybe_parent = this.qt_cache.get(parent_id);
    let invalid =
      is_none(maybe_parent) ||
      !maybe_parent.boundary ||
      !maybe_parent.boundary.containsPoint(spatial);
    // IF parent is NONE
    //   OR has a bad boundary
    //   OR is not contained within this boundary
    // THEN remove the point from its parent
    //   AND reinsert from root
    // ELSE do nothing, since no change requried
    if (invalid) {
      this._remove(spatial.id);
      this.qt_cache.get(this.root)?.insert(spatial);
    }
  }

  /**
   * returns all spatials within the tree.
   * @param spatials - optional starting set [default: empty Set].
   * @returns All contained spatials within the tree.
   */
  getAllSpatials(spatials: Set<Spatial> = new Set()): Set<Spatial> {
    if (this.divided) {
      spatials =
        this.qt_cache.get(this.ne)?.getAllSpatials(spatials) ?? spatials;
      spatials =
        this.qt_cache.get(this.nw)?.getAllSpatials(spatials) ?? spatials;
      spatials =
        this.qt_cache.get(this.se)?.getAllSpatials(spatials) ?? spatials;
      spatials =
        this.qt_cache.get(this.sw)?.getAllSpatials(spatials) ?? spatials;
    } else {
      this.getData().forEach(spatial => spatials.add(spatial));
    }
    return spatials;
  }

  /**
   * update the bounds of the current quadtree
   * @param newBounds - the bounds for the quadtree.
   * @param spatials - optional spatials to insert [default: all previously inserted spatials]
   */
  updateBounds(
    newBounds: RectangleRange,
    spatials: Set<Spatial> = this.getAllSpatials(),
  ): void {
    // clear structure, update bounds, and insert all the exising or
    // provided spatials while also clearing their parents
    this.reset(newBounds);
    spatials.forEach(spatial => {
      this.remove(spatial);
      this.insert(spatial);
    });
  }
}