Reatom computed map

mapAtom.ts

import { Atom, atom, Ctx, CtxSpy } from "@reatom/framework";

export type MapAtom<K, V> = Atom<Map<K, V>>;

export type AnyKeysShape<Key> = Key[] | Map<Key, any>;

export const buildMapAtom = <
  Key,
  Value,
  KeysAtom extends Atom<AnyKeysShape<Key>>,
>(
  keysAtom: KeysAtom,
  builder: (key: Key, name: string, mapAtomCtx: Ctx) => Value,
  name: string,
): MapAtom<Key, Value> =>
  atom((ctx) => {
    const anyKeysShape = ctx.spy(keysAtom);
    const keys =
      anyKeysShape instanceof Map ? anyKeysShape.keys() : anyKeysShape;

    const map = new Map<Key, Value>();

    for (const key of keys) {
      map.set(key, builder(key, `${name}.${String(key)}`, ctx));
    }

    return map;
  }, `${name}.buildMapAtom`);

const mapAtomForceGet = <K, V>(
  mapAtomName: string | undefined,
  map: Map<K, V>,
  key: K,
): V => {
  if (!map.has(key)) {
    throw new Error(
      `Unable to lookup ${String(key)} in ${mapAtomName ?? "<unnamed atom>"}`,
    );
  }
  return map.get(key)!;
};

export const ctxSpyMap = <K, V>(
  ctx: CtxSpy,
  mapAtom: MapAtom<K, V>,
  key: K,
): V => {
  const map = ctx.spy(mapAtom);
  return mapAtomForceGet(mapAtom.__reatom.name, map, key);
};

export const ctxGetMap = <K, V>(
  ctx: Ctx,
  mapAtom: MapAtom<K, V>,
  key: K,
): V => {
  const map = ctx.get(mapAtom);
  return mapAtomForceGet(mapAtom.__reatom.name, map, key);
};