Semi-mutable Set and Map, immutable wrappers over native mutable js Maps and Sets

SemiMutable.js

//@flow

export class SemiMap<Key, Val> {
  /*::
  ;+_: Map<Key, Val>
  ;+@@iterator: () => Iterator<[Key, Val]>
  mutable: boolean
  */
  constructor(parent: Map<Key, Val> = new Map()) {
    this._ = parent
    this.mutable = false
  }
  clone(): SemiMap<Key, Val> {
    if (this.mutable) return this
    return new SemiMap(this._)
  }
  get(key: Key): Val | void {
    return this._.get(key)
  }
  set(key: Key, val: Val): SemiMap<Key, Val> {
    if (this._.get(key) === val) return this
    this._.set(key, val)
    return this.clone()
  }
  has(key: Key): boolean {
    return this._.has(key)
  }
  delete(key: Key): SemiMap<Key, Val> {
    if (!this._.has(key)) return this
    this._.delete(key)
    return this.clone()
  }
  clear(): SemiMap<Key, Val> {
    if (this._.size === 0) return this
    this._.clear()
    return this.clone()
  }
  entries(): Iterator<[Key, Val]> {
    return this._.entries()
  }
  values(): Iterator<Val> {
    return this._.values()
  }
  keys(): Iterator<Key> {
    return this._.keys()
  }
  map<Res>(
    fn: (val: Val, key: Key, map: SemiMap<Key, Val>) => Res,
    acc?: SemiMap<Key, Res>,
  ): SemiMap<Key, Res> {
    return mapReduce(this, fn, acc)
  }
  forEach(
    callbackfn: (value: Val, index: Key, map: Map<Key, Val>) => mixed,
    thisArg?: any,
  ) {
    if (thisArg === undefined) this._.forEach(callbackfn)
    else this._.forEach(callbackfn, thisArg)
  }
  static from<Key, Val>(
    value: | Map<Key, Val>
      | $ReadOnlyArray<[Key, Val]>
      | SemiMap<Key, Val>
      | Iterable<[Key, Val]>,
  ): SemiMap<Key, Val> {
    if (value instanceof SemiMap) return new SemiMap(value._)
    if (value instanceof Map) return new SemiMap(value)
    return new SemiMap(new Map(value))
  }
  //$off
  [Symbol.iterator]() {
    return this._[Symbol.iterator]()
  }
  get size(): number {
    return this._.size
  }
}

export class SemiSet<Val> {
  /*::
  ;+_: Set<Val>
  ;+@@iterator: () => Iterator<Val>
  mutable: boolean
  */
  constructor(parent: Set<Val> = new Set()) {
    this._ = parent
    this.mutable = false
  }
  clone(): SemiSet<Val> {
    if (this.mutable) return this
    return new SemiSet(this._)
  }
  add(val: Val): SemiSet<Val> {
    if (this._.has(val)) return this
    this._.add(val)
    return this.clone()
  }
  has(val: Val): boolean {
    return this._.has(val)
  }
  delete(val: Val): SemiSet<Val> {
    if (!this._.has(val)) return this
    this._.delete(val)
    return this.clone()
  }
  clear(): SemiSet<Val> {
    if (this._.size === 0) return this
    this._.clear()
    return this.clone()
  }
  entries(): Iterator<[Val, Val]> {
    return this._.entries()
  }
  values(): Iterator<Val> {
    return this._.values()
  }
  keys(): Iterator<Val> {
    return this._.keys()
  }
  forEach(
    callbackfn: (value: Val, index: Val, set: Set<Val>) => mixed,
    thisArg?: any,
  ) {
    if (thisArg === undefined) this._.forEach(callbackfn)
    else this._.forEach(callbackfn, thisArg)
  }
  static from<Val>(
    value: Set<Val> | $ReadOnlyArray<Val> | SemiSet<Val> | Iterable<Val>,
  ): SemiSet<Val> {
    if (value instanceof SemiSet) return new SemiSet(value._)
    if (value instanceof Set) return new SemiSet(value)
    return new SemiSet(new Set(value))
  }
  //$off
  [Symbol.iterator]() {
    return this._[Symbol.iterator]()
  }
  get size(): number {
    return this._.size
  }
}

export function getIn<KeyA, KeyB, Val>(
  pair: [KeyA, KeyB],
  map: SemiMap<KeyA, SemiMap<KeyB, Val>>,
): Val | void {
  const subMap = map.get(pair[0])
  if (subMap === undefined) return
  return subMap.get(pair[1])
}

export function setIn<KeyA, KeyB, Val>(
  pair: [KeyA, KeyB],
  val: Val,
  map: SemiMap<KeyA, SemiMap<KeyB, Val>>,
): SemiMap<KeyA, SemiMap<KeyB, Val>> {
  return map.set(pair[0], (map.get(pair[0]) || new SemiMap()).set(pair[1], val))
}

export function hasIn<KeyA, KeyB, Val>(
  pair: [KeyA, KeyB],
  map: SemiMap<KeyA, SemiMap<KeyB, Val>>,
): boolean {
  const child = map.get(pair[0])
  if (!child) return false
  return child.has(pair[1])
}
export function deleteIn<KeyA, KeyB, Val>(
  pair: [KeyA, KeyB],
  map: SemiMap<KeyA, SemiMap<KeyB, Val>>,
): SemiMap<KeyA, SemiMap<KeyB, Val>> {
  let child = map.get(pair[0])
  if (!child) return map
  child = child.delete(pair[1])
  if (child.size === 0) return map.delete(pair[0])
  return map.set(pair[0], child)
}

export function mapReduce<Key, Val, Res>(
  source: SemiMap<Key, Val>,
  mapReducer: (
    val: Val,
    key: Key,
    source: SemiMap<Key, Val>,
    acc: SemiMap<Key, Res>,
  ) => Res,
  acc?: SemiMap<Key, Res>,
): SemiMap<Key, Res> {
  let current = acc || new SemiMap()
  const keys = new Set(current.keys())
  for (const [key, val] of source) {
    keys.delete(key)
    current = current.set(key, mapReducer(val, key, source, current))
  }
  //Remove unmatched keys
  //It happens when source delete items between calls
  current.mutable = true
  for (const key of keys) {
    current.delete(key)
  }
  current.mutable = false
  return current
}