Data structure for handling multiple bitmasks within one stack

MultiClassBitmask.ts

import _noop from "lodash/noop";

import SparseBitmask from "../../utils/SparseBitmask";
import { OrderedSet, OrderedMap } from "../../utils/OrderedCollections";
import { range } from "../../utils/base";

interface RunNode<T> {
  currState: OrderedSet<number>;
  diff: Set<number>;
  pending: Set<number>;
  newFlag: boolean;
  lastPixel?: T;
}

type OnPixelChanged<T> = (value: T | undefined, index: number) => void;
export type PixelShader<T> = (
  layers: OrderedSet<number>,
  layerToMask: Map<number, SparseBitmask>
) => T | undefined;

const DEFAULT_PIXEL_SHADER: PixelShader<any> = (layers, layerToMask) =>
  layerToMask.get(layers.peek()!);

export default class MultiClassBitmask<T = SparseBitmask> {
  public length: number;
  public pixelShader: PixelShader<T>;
  private data: OrderedMap<number, RunNode<T>> = new OrderedMap();
  public children: OrderedMap<SparseBitmask, number> = new OrderedMap();
  private layerIndexToMask: Map<number, SparseBitmask> = new Map();
  private minLayerIndex = 0;
  private maxLayerIndex = -1;
  private flushQueued = false;
  private queuedFlushCallback?: OnPixelChanged<T>;
  private queuedForForceUpdate = new Set<number>();
  private queuedForRemove = new Set<number>();

  constructor(
    length: number,
    pixelShader: PixelShader<T> = DEFAULT_PIXEL_SHADER
  ) {
    this.length = length;
    this.pixelShader = pixelShader;
    this.deferredFlush = this.deferredFlush.bind(this);
  }

  reset(length: number = this.length, callback?: OnPixelChanged<T>): this {
    if (callback) {
      this.queuedFlushCallback = callback;
      this.deferredFlush();
    }
    this.length = length;
    this.data.clear();
    this.children.clear();
    this.layerIndexToMask.clear();
    this.minLayerIndex = 0;
    this.maxLayerIndex = -1;
    this.queuedFlushCallback = undefined;
    this.queuedForRemove.clear();
    return this;
  }

  // flush is batched so that this.data is iterated over only once across multiple operations
  private flush(callback?: OnPixelChanged<T>) {
    if (this.flushQueued && callback !== this.queuedFlushCallback) {
      this.deferredFlush();
      this.flushQueued = true;
    }
    this.queuedFlushCallback = callback;
    if (!this.flushQueued) {
      this.flushQueued = true;
      queueMicrotask(this.deferredFlush);
    }
  }

  private deferredFlush() {
    const {
      data,
      queuedForRemove,
      queuedForForceUpdate,
      queuedFlushCallback: callback,
      layerIndexToMask,
      pixelShader,
    } = this;

    const shouldForceUpdate = (layers: OrderedSet<number>) => {
      // eslint-disable-next-line no-restricted-syntax
      for (const layer of layers) {
        if (queuedForForceUpdate.has(layer)) return true;
      }
      return false;
    };

    let prevState = new OrderedSet<number>(layerIndex => -layerIndex);
    let runStart = -1;
    let beforePixel: T | undefined;
    let afterPixel: T | undefined;
    const acc = new Set<number>();
    const toRemove: number[] = [];
    data.forEach((node, runEnd) => {
      if (runStart >= 0) {
        if (callback) {
          // eslint-disable-next-line no-restricted-syntax
          for (const i of range(runStart, runEnd)) {
            callback(afterPixel, i);
          }
        }
        runStart = -1;
      }

      const { currState, diff, pending, newFlag } = node;
      diff.forEach(layerIndex => {
        if (queuedForRemove.has(layerIndex)) {
          if (pending.has(layerIndex)) pending.delete(layerIndex);
          else pending.add(layerIndex);
        }
      });
      pending.forEach(layerIndex => {
        if (diff.has(layerIndex)) diff.delete(layerIndex);
        else diff.add(layerIndex);
        if (acc.has(layerIndex)) acc.delete(layerIndex);
        else acc.add(layerIndex);
      });
      pending.clear();

      if (newFlag) {
        prevState.forEach(layerIndex => {
          currState.add(layerIndex);
        });
        diff.forEach(layerIndex => {
          if (currState.has(layerIndex)) currState.delete(layerIndex);
          else currState.add(layerIndex);
        });
        node.lastPixel = beforePixel;
        node.newFlag = false;
        afterPixel = pixelShader(currState, layerIndexToMask);
      } else {
        beforePixel = node.lastPixel;
        if (acc.size > 0) {
          acc.forEach(layerIndex => {
            if (currState.has(layerIndex)) {
              currState.delete(layerIndex);
            } else {
              currState.add(layerIndex);
            }
          });
          afterPixel = pixelShader(currState, layerIndexToMask);
        } else if (shouldForceUpdate(currState)) {
          afterPixel = pixelShader(currState, layerIndexToMask);
        } else {
          afterPixel = beforePixel;
        }
      }

      if (afterPixel !== beforePixel) {
        runStart = runEnd;
        node.lastPixel = afterPixel;
      }
      if (diff.size === 0) {
        toRemove.push(runEnd);
      }
      prevState = currState;
    });
    if (callback && runStart >= 0) {
      // eslint-disable-next-line no-restricted-syntax
      for (const i of range(runStart, this.length)) {
        callback(afterPixel, i);
      }
    }

    toRemove.forEach(index => {
      data.delete(index);
    });

    this.flushQueued = false;
    this.queuedFlushCallback = undefined;
    queuedForRemove.clear();
  }

  private removeFromData(mask: SparseBitmask): void {
    const layerIndex = this.children.get(mask);
    if (layerIndex != null) {
      this.queuedForRemove.add(layerIndex);
      this.children.delete(mask);
      this.layerIndexToMask.delete(layerIndex);
    }
  }

  private addToData(mask: SparseBitmask, layerIndex: number): void {
    // eslint-disable-next-line no-restricted-syntax
    for (const index of mask.getIndices(true)) {
      this.data.update(index, (node = getDefaultNode()) => {
        const { pending } = node;
        if (pending.has(layerIndex)) pending.delete(layerIndex);
        else pending.add(layerIndex);
        return node;
      });
    }
    this.children.set(mask, layerIndex);
    this.layerIndexToMask.set(layerIndex, mask);
  }

  remove(mask: SparseBitmask, onPixelChanged?: OnPixelChanged<T>): void {
    this.removeFromData(mask);
    this.flush(onPixelChanged);
  }

  addToTop(mask: SparseBitmask, onPixelChanged?: OnPixelChanged<T>): void {
    if (this.children.has(mask)) {
      this.removeFromData(mask);
    }
    this.maxLayerIndex += 1;
    this.addToData(mask, this.maxLayerIndex);
    this.flush(onPixelChanged);
  }

  addToBottom(mask: SparseBitmask, onPixelChanged?: OnPixelChanged<T>): void {
    if (this.children.has(mask)) {
      this.removeFromData(mask);
    }
    this.minLayerIndex -= 1;
    this.addToData(mask, this.minLayerIndex);
    this.flush(onPixelChanged);
  }

  // returns a callback to revert the change
  bringToFront(
    mask: SparseBitmask,
    onPixelChanged?: OnPixelChanged<T>
  ): (onPixelChanged?: OnPixelChanged<T>) => void {
    const originalLayerIndex = this.children.get(mask);
    if (originalLayerIndex == null) return _noop;
    this.removeFromData(mask);
    this.addToTop(mask, onPixelChanged);
    return (_onPixelChanged = onPixelChanged) => {
      this.removeFromData(mask);
      this.addToData(mask, originalLayerIndex);
      this.flush(_onPixelChanged);
    };
  }

  // returns a callback to revert the change
  sendToBack(
    mask: SparseBitmask,
    onPixelChanged?: OnPixelChanged<T>
  ): (onPixelChanged?: OnPixelChanged<T>) => void {
    const originalLayerIndex = this.children.get(mask);
    if (originalLayerIndex == null) return _noop;
    this.removeFromData(mask);
    this.addToBottom(mask, onPixelChanged);
    return (_onPixelChanged = onPixelChanged) => {
      this.removeFromData(mask);
      this.addToData(mask, originalLayerIndex);
      this.flush(_onPixelChanged);
    };
  }

  /**
   * update mask without changing their layer indices
   * set forceUpdate to true to force callback to execute even if no pixel change
   * eg. when changing layer color
   */
  update(
    mask: SparseBitmask,
    onPixelChanged?: OnPixelChanged<T>,
    forceUpdate = false
  ): void {
    const layerIndex = this.children.get(mask);
    if (layerIndex == null) return;
    this.removeFromData(mask);
    this.addToData(mask, layerIndex);
    if (forceUpdate) {
      this.queuedForForceUpdate.add(layerIndex);
    }
    this.flush(onPixelChanged);
  }

  search(index: number): SparseBitmask | undefined {
    const binarySearch = (low: number, high: number): number => {
      if (low >= high) return high;
      const mid = Math.ceil((low + high) / 2);
      const [midIndex] = this.data.at(mid)!;
      if (midIndex === index) return mid;
      return index < midIndex
        ? binarySearch(low, mid - 1)
        : binarySearch(mid, high);
    };
    const narrowed = binarySearch(-1, this.data.size - 1);
    if (narrowed < 0) return undefined;
    const [_, node] = this.data.at(narrowed)!;
    return this.layerIndexToMask.get(node.currState.peek()!);
  }
}

function getDefaultNode(): RunNode<any> {
  return {
    currState: new OrderedSet<number>(layerIndex => -layerIndex),
    diff: new Set(),
    pending: new Set(),
    newFlag: true,
  };
}