getters.tsx

import { AppLoadContext, DataFunctionArgs } from "react-router";

type GetterFn<T> = (args: DataFunctionArgs) => T;

export function createGetter<T>(fn: GetterFn<T>) {
  let weakMap = new WeakMap<
    AppLoadContext,
    { result: T; thrown?: undefined } | { result?: undefined; thrown: any }
  >();

  return function (args: DataFunctionArgs) {
    const key = args.context;

    if (weakMap.has(key)) {
      const { thrown, result } = weakMap.get(key)!;
      if (thrown) throw thrown;

      return result!;
    }

    try {
      const result = fn(args);
      weakMap.set(key, { result });

      return result;
    } catch (e) {
      weakMap.set(key, { thrown: e });
      throw e;
    }
  };
}
type KeyedGetterFunction<T, K> = (args: DataFunctionArgs, key: K) => T;

export function createKeyedGetter<T, K>(fn: KeyedGetterFunction<T, K>) {
  let weakMap = new WeakMap<
    AppLoadContext,
    Map<
      K,
      { result: T; thrown?: undefined } | { result?: undefined; thrown: any }
    >
  >();

  return function (args: DataFunctionArgs, key: K) {
    const context = args.context;

    if (!weakMap.has(context)) {
      weakMap.set(context, new Map());
    }

    const map = weakMap.get(context)!;
    if (map.has(key)) {
      const { thrown, result } = map.get(key)!;
      if (thrown) throw thrown;

      return result!;
    }

    try {
      const result = fn(args, key);
      map.set(key, { result });

      return result;
    } catch (e) {
      map.set(key, { thrown: e });
      throw e;
    }
  };
}
type BatchGetterFunction<T, K> = (
  args: DataFunctionArgs,
  keys: K[],
) => Promise<Map<K, T>>;

export function createBatchedGetter<T, K>(fn: BatchGetterFunction<T, K>) {
  let weakMap = new WeakMap<
    AppLoadContext,
    {
      isDispatchQueued: boolean;
      map: Map<
        K,
        {
          isDispatched: boolean;
        } & ReturnType<typeof withResolvers<T>>
      >;
    }
  >();

  function get(args: DataFunctionArgs, key: K) {
    const context = args.context;

    if (!weakMap.has(context)) {
      weakMap.set(context, {
        isDispatchQueued: false,
        map: new Map(),
      });
    }

    const entry = weakMap.get(context)!;
    if (entry.map.has(key)) {
      return entry.map.get(key)!.promise;
    }

    let result = withResolvers<T>();
    entry.map.set(key, {
      isDispatched: false,
      ...result,
    });

    if (!entry.isDispatchQueued) {
      entry.isDispatchQueued = true;
      queueMicrotask(() => dispatch(args));
    }

    return result.promise;
  }

  function dispatch(args: DataFunctionArgs) {
    const context = args.context;
    const entry = weakMap.get(context);
    if (!entry) return;

    entry.isDispatchQueued = false;
    const { map } = entry;

    const keys = Array.from(map.entries())
      .filter(([_, { isDispatched }]) => isDispatched === false)
      .map(([key]) => key);

    if (keys.length === 0) return;

    const results = fn(args, keys);

    keys.forEach((key) => {
      const entry = map.get(key)!;

      entry.isDispatched = true;

      entry.resolve(
        results.then((map) => {
          return map.get(key)!;
        }),
      );
    });
  }

  return get;
}
function withResolvers<T>() {
  let resolve: (value: T | PromiseLike<T>) => void = () => {};
  let reject: (reason?: any) => void = () => {};

  const promise = new Promise<T>((a, b) => {
    resolve = a;
    reject = b;
  });

  return { promise, resolve, reject };
}