Implementation of produce function that supports Maps and Sets

immer.ts

class Draft<T extends object> {
  private staged = false;
  private committed = false;
  private readonly immutable: Readonly<T>;
  public readonly mutable: T;
  private target: unknown;
  protected pool: Pool;

  public constructor (immutable: Readonly<T>, mutable: T, target: unknown, pool: Pool) {
    this.immutable = immutable;
    this.mutable = mutable;
    this.target = target;
    this.pool = pool;
  }

  public stage (): void {
    this.staged = true;
    this.pool.tryStage(this.target);
  }

  protected commit (): void { }

  public copy (): T {
    if (!this.staged) {
      return this.immutable;
    }

    if (!this.committed) {
      this.committed = true;
      this.commit();
    }

    return this.mutable;
  }
}

class ArrayDraft<T extends any[]> extends Draft<T> {
  protected commit (): void {
    for (let i = 0; i < this.mutable.length; ++i) {
      this.mutable[i] = this.pool.copyOr(this.mutable[i]);
    }
  }
}

const move = <T extends Map<any, any> | Set<any>> (original: T): T => {
  const Constructor = original.constructor as MapConstructor & SetConstructor;
  const copy = new Constructor(original) as T;
  original.clear();
  return copy;
};

class MapDraft<T extends Map<any, any>> extends Draft<T> {
  protected commit (): void {
    for (const [key, value] of move(this.mutable)) {
      this.mutable.set(this.pool.copyOr(key), this.pool.copyOr(value));
    }
  }
}

class SetDraft<T extends Set<any>> extends Draft<T> {
  protected commit (): void {
    for (const value of move(this.mutable)) {
      this.mutable.add(this.pool.copyOr(value));
    }
  }
}

class ObjectDraft<T extends Object> extends Draft<T> {
  protected commit (): void {
    for (const key of Object.keys(this.mutable)) {
      this.mutable[key] = this.pool.copyOr(this.mutable[key]);
    }
  }
}

class Handler {
  private pool: Pool;

  public constructor (pool: Pool) {
    this.pool = pool;
  }

  public apply (target: object, thisArg: object, args: object[]): unknown {
    const mutableThisArg = this.pool.mutableOr(thisArg);
    const mutableArgs = args.map(arg => this.pool.mutableOr(arg));
    const result = Reflect.apply(target as Function, mutableThisArg, mutableArgs);

    switch (target) {
      case Map.prototype.clear:
      case Map.prototype.delete:
      case Map.prototype.set:
      case Set.prototype.add:
      case Set.prototype.clear:
      case Set.prototype.delete:
        this.pool.tryStage(target);
    }

    return this.pool.proxyOr(result);
  }

  public defineProperty (): boolean {
    return false;
  }

  public deleteProperty (target: object, key: string): boolean {
    if (!Reflect.deleteProperty(target, key)) {
      return false;
    }

    this.pool.tryStage(target);
    return true;
  }

  public get (target: object, key: string): unknown {
    const value = Reflect.get(target, key);

    return this.pool.proxyOr(value, target);
  }

  public set (target: object, key: string, value: object): boolean {
    const mutable = this.pool.mutableOr(value);

    if (!Reflect.set(target, key, mutable)) {
      return false;
    }

    this.pool.tryStage(target);
    return true;
  }
}

class Nullopt { }

class Optional<T = void> {
  public static readonly nullopt = new Nullopt();

  public static of<U = void> (callback: () => U | Nullopt): Optional<U> {
    return new this(callback());
  }

  public readonly value: T | Nullopt;

  public constructor (value: T | Nullopt) {
    this.value = value;
  }

  public hasValue (): boolean {
    return !(this.value instanceof Nullopt);
  }

  public and<U = void> (callback: (value: T) => U | Nullopt): Optional<U> {
    if (!this.hasValue()) {
      return this as unknown as Optional<U>;
    }
    return Optional.of(() => callback(this.value as T));
  }

  public or (callback: () => T | Nullopt): Optional<T> {
    if (this.hasValue()) {
      return this;
    }
    return Optional.of(callback);
  }
}

class Pool {
  private static tryGet<T> (key: unknown, map: Map<unknown, T>): Optional<T> {
    return Optional.of<T>(() => map.has(key) ? map.get(key)! : Optional.nullopt);
  }

  private drafts = new Map<unknown, Draft<any>>();
  private revocables = new Map<unknown, { proxy: any, revoke: () => void; }>();
  private handler = new Handler(this);

  private addProxy<T extends object> (immutable: T, draft: Draft<T>): T {
    const { mutable } = draft;
    const revocable = Proxy.revocable<T>(mutable, this.handler);
    const { proxy } = revocable;

    this.drafts
      .set(immutable, draft)
      .set(mutable, draft)
      .set(proxy, draft);
    this.revocables
      .set(immutable, revocable)
      .set(mutable, revocable)
      .set(proxy, revocable);

    return proxy;
  }

  private tryAddProxy<T extends object> (immutable: T, target: unknown): Optional<T> {
    return Optional
      .of((): [typeof Draft, T] | Nullopt => Array.isArray(immutable) ? [ArrayDraft, [...immutable]] : Optional.nullopt)
      .or(() => immutable instanceof Map ? [MapDraft, new Map(immutable)] : Optional.nullopt)
      .or(() => immutable instanceof Set ? [SetDraft, new Set(immutable)] : Optional.nullopt)
      .or(() => typeof immutable === 'function' ? [Draft, immutable] : Optional.nullopt)
      .or(() => typeof immutable === 'object' && immutable !== null ? [ObjectDraft, { ...immutable }] : Optional.nullopt)
      .and(([Draft, mutable]) => this.addProxy(immutable, new Draft(immutable, mutable, target, this)));
  }

  public proxyOr<T extends object> (immutable: T, target: unknown = null): T {
    return Pool.tryGet(immutable, this.revocables)
      .and<T>(revocable => revocable.proxy)
      .or(() => this.tryAddProxy(immutable, target).value)
      .or(() => immutable)
      .value as T;
  }

  public mutableOr<T extends object> (key: T): T {
    return Pool.tryGet(key, this.drafts)
      .and(draft => draft.mutable)
      .or(() => key)
      .value as T;
  }

  public tryStage (key: unknown): void {
    Pool.tryGet(key as object, this.drafts)
      .and(draft => { draft.stage(); });
  }

  public copyOr<T extends object> (key: T): T {
    return Pool.tryGet(key, this.drafts)
      .and(draft => draft.copy())
      .or(() => key)
      .value as T;
  }

  public revoke (): void {
    for (const revocable of this.revocables.values()) {
      revocable.revoke();
    }

    this.revocables.clear();
    this.drafts.clear();
  }
}

const produce = <T extends object, TArgs extends any[]> (
  producer: (previous: T, ...args: TArgs) => void,
) => (previous: T, ...args: TArgs): T => {
  const pool = new Pool();

  try {
    const mutable = pool.proxyOr(previous);
    producer(mutable, ...args);
    return pool.copyOr(mutable);
  } finally {
    pool.revoke();
  }
};

export default produce;