Data structure for handling multiple bitmasks within one stack

MultiClassBitmask.ts

import _noop from "lodash/noop";

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

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

type OnPixelChanged = (mask: SparseBitmask | undefined, index: number) => void;

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

  constructor(length: number) {
    this.length = length;
    this.applyToRun = this.applyToRun.bind(this);
    this.deferredFlush = this.deferredFlush.bind(this);
  }

  reset(
    length: number = this.length,
    callback?: (mask: SparseBitmask | undefined, index: number) => void
  ): this {
    if (callback) {
      this.queuedFlushCallback = callback;
      this.queuedForForceUpdate = new Set(this.layerIndexToMask.keys());
      this.deferredFlush();
    }
    this.length = length;
    this.data.clear();
    this.children.clear();
    this.layerIndexToMask.clear();
    this.minLayerIndex = 0;
    this.maxLayerIndex = -1;
    this.queuedFlushCallback = undefined;
    this.queuedForForceUpdate.clear();
    this.queuedForRemove.clear();
    return this;
  }

  // flush is batched so that this.data is iterated over only once across multiple operations
  private flush(callback?: OnPixelChanged) {
    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,
      applyToRun,
    } = this;

    let prevState = new OrderedSet<number>(layerIndex => -layerIndex);
    let runStart = -1;
    let beforeTop: number | undefined;
    let afterTop: number | undefined;
    const acc = new Set<number>();
    const toRemove: number[] = [];
    for (const [runEnd, node] of data) {
      if (runStart >= 0) {
        if (callback) applyToRun(runStart, runEnd, afterTop, callback);
        runStart = -1;
      }

      const { currState, diff, pending, newFlag } = node;
      for (const layerIndex of diff) {
        if (queuedForRemove.has(layerIndex)) {
          if (pending.has(layerIndex)) pending.delete(layerIndex);
          else pending.add(layerIndex);
        }
      }
      for (const layerIndex of pending) {
        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) {
        for (const layerIndex of prevState) {
          currState.add(layerIndex);
        }
        for (const layerIndex of diff) {
          if (currState.has(layerIndex)) currState.delete(layerIndex);
          else currState.add(layerIndex);
        }
        afterTop = currState.peek();
        if (beforeTop !== afterTop) {
          runStart = runEnd;
        }
        node.newFlag = false;
      } else {
        beforeTop = currState.peek();
        if (acc.size > 0) {
          for (const layerIndex of acc) {
            if (currState.has(layerIndex)) {
              currState.delete(layerIndex);
            } else {
              currState.add(layerIndex);
            }
          }
          afterTop = currState.peek();
          if (beforeTop !== afterTop) {
            runStart = runEnd;
          }
        } else {
          // update afterTop still because need to check against forced
          afterTop = beforeTop;
        }
      }

      if (runStart < 0 && queuedForForceUpdate.has(afterTop)) {
        runStart = runEnd;
      }
      if (diff.size === 0) {
        toRemove.push(runEnd);
      }
      prevState = currState;
    }
    if (callback && runStart >= 0) {
      applyToRun(runStart, this.length, afterTop, callback);
    }

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

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

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

  private addToData(mask: SparseBitmask, layerIndex: number): void {
    for (const index of mask.asRunEndIndices().getIndices()) {
      const existing = this.data.get(index) ?? {
        currState: new OrderedSet<number>(layerIndex => -layerIndex),
        diff: new Set(),
        pending: new Set(),
        newFlag: true,
      };
      const { pending } = existing;
      if (pending.has(layerIndex)) pending.delete(layerIndex);
      else pending.add(layerIndex);
      this.data.set(index, existing);
    }
    this.children.set(mask, layerIndex);
    this.layerIndexToMask.set(layerIndex, mask);
  }

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

  addToTop(mask: SparseBitmask, onPixelChanged?: OnPixelChanged): 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): 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
  ): (onPixelChanged?: OnPixelChanged) => 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
  ): (onPixelChanged?: OnPixelChanged) => 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,
    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);
      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
      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;
    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
    const [_, node] = this.data.at(narrowed)!;
    return this.layerIndexToMask.get(node.currState.peek());
  }

  private applyToRun(
    runStart: number,
    runEnd: number,
    layerIndex: number | undefined,
    callback: (mask: SparseBitmask | undefined, index: number) => void
  ): void {
    const mask = this.layerIndexToMask.get(layerIndex);
    for (let i = runStart; i < runEnd; i += 1) {
      callback(mask, i);
    }
  }
}