Liveblocks hooks with suspense

liveblocksHooks.ts

import { LiveList, LiveMap, LiveObject, Presence } from '@liveblocks/client';
import { useRoom } from '@liveblocks/react';
import { useState, useEffect, useMemo, useRef } from 'react';
import invariant from 'tiny-invariant';
import { useAsset } from 'use-asset';

export function useRoot() {
  const room = useRoom();
  // suspend until root storage is ready
  const { root } = useAsset((r) => {
    return r.getStorage<any>();
  }, room);
  return root;
}

function useRerenderOnChange(
  // for some reason room.subscribe does NOT like a | type of
  // liveblocks objects.
  parent: any,
  child: any,
) {
  const room = useRoom();
  const [_, setCounter] = useState(0);
  useEffect(() => {
    if (!child) return;
    return room.subscribe(child, () => {
      setCounter((v) => v + 1);
    });
  }, [room, child]);
  useEffect(() => {
    if (!parent) return;
    return room.subscribe(parent, () => {
      setCounter((v) => v + 1);
    });
  }, [parent, room]);
}

function getOrCreateInParent<T>(
  parent: null | LiveObject | LiveMap<string, T> | LiveList<T>,
  key: string | number,
  creator: () => T,
) {
  if (parent === null) {
    return null;
  }
  if (parent instanceof LiveList) {
    invariant(
      typeof key === 'number',
      'Parent is a List, expected indexOrKey to be a number',
    );
    const existing = parent.get(key);
    if (existing) {
      return existing;
    }
    // TODO: upsert into list? not sure how it would work.
    throw new Error('Index out of bounds in list');
  } else {
    invariant(
      typeof key === 'string',
      'Parent is an Object or Map, expected indexOrKey to be a string',
    );
    const existing = parent.get(key);
    if (existing) {
      return existing;
    }
    const created = creator();
    parent.set(key, created);
    return created;
  }
}

export function useLiveList<TItem>(
  parent:
    | null
    | LiveObject
    | LiveMap<string, LiveList<TItem>>
    | LiveList<LiveList<TItem>>,
  key: number | string,
  initialEntries?: TItem[],
): LiveList<TItem> | null {
  const initialEntriesRef = useRef(initialEntries);
  const target = useMemo(
    () =>
      getOrCreateInParent(
        parent,
        key,
        () => new LiveList<TItem>(initialEntriesRef.current),
      ),
    [key, parent],
  );
  useRerenderOnChange(parent, target);

  return target;
}

export function useLiveObject<TShape>(
  parent:
    | null
    | LiveObject
    | LiveMap<string, LiveObject<TShape>>
    | LiveList<LiveObject<TShape>>,
  key: string | number,
  initialData?: TShape,
): LiveObject<TShape> {
  const initialDataRef = useRef(initialData);
  const target = useMemo(
    () =>
      getOrCreateInParent(
        parent,
        key,
        () => new LiveObject(initialDataRef.current),
      ),
    [key, parent],
  );
  useRerenderOnChange(parent, target);

  return target;
}

export function useLiveMap<TKey extends string, TValue>(
  parent:
    | null
    | LiveObject
    | LiveMap<string, LiveMap<TKey, TValue>>
    | LiveList<LiveMap<TKey, TValue>>,
  key: string | number,
  initialEntries?: [TKey, TValue][],
) {
  const initialEntriesRef = useRef(initialEntries);
  const target = useMemo(
    () =>
      getOrCreateInParent(
        parent,
        key,
        () => new LiveMap<TKey, TValue>(initialEntriesRef.current),
      ),
    [key, parent],
  );
  useRerenderOnChange(parent, target);

  return target;
}