recoil-jotai-compat.ts

import { useRef, useLayoutEffect, useCallback } from "react";
import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai/react";
import { atom as jotaiAtom } from "jotai/vanilla";
import { atomFamily as jotaiAtomFamily } from "jotai/utils";
import deepEqual from "fast-deep-equal";
import type { Atom as JotaiReadbleAtom, WritableAtom as _JotaiWritableAtom } from "jotai";

function useEvent<F extends Function>(handler: F): F {
  const handlerRef = useRef<F>(null);

  useLayoutEffect(() => {
    (handlerRef.current as any) = handler;
  });

  return useCallback((...args: any): any => {
    return handlerRef.current?.(...args);
  }, []) as any;
}

type JotaiWritableAtom<T> = _JotaiWritableAtom<T, [nextValue: T | ((prev: T) => T)], void>;
export type CompatAtomState<T> = CompatReadbleAtomState<T> | CompatWritableAtomState<T>;

export type CompatReadbleAtomState<T> = {
  key: string;
  type: "readble";
  jotai: JotaiReadbleAtom<T>;
  default: () => T;
};

type UpdateHook = () => void;
export type CompatWritableAtomState<T> = {
  key: string;
  type: "writable";
  jotai: JotaiWritableAtom<T>;
  default: () => T;
  updateHooks: UpdateHook[];
};

export type AtomEffect<T> = (param: {
  onSet: (fn: () => void) => void;
}) => void;

export type AtomOptions<T> = {
  // For compatibility with Recoil
  key: string;
  effects?: readonly AtomEffect<T>[];
  default: T;
};
export function atom<T>(options: AtomOptions<T>): CompatWritableAtomState<T> {
  const updateHooks: UpdateHook[] = [];
  options.effects?.forEach((effect) => {
    effect({
      onSet: (fn) => {
        updateHooks.push(fn);
      },
    });
  });
  const jotaiState = jotaiAtom(options.default, (get, set, nextValue: T | ((prev: T) => T)) => {
    if (nextValue instanceof Function) {
      set(jotaiState, nextValue(get(jotaiState)));
    } else {
      set(jotaiState, nextValue);
    }
  });
  return {
    key: options.key,
    jotai: jotaiState,
    default: () => options.default,
    type: "writable",
    updateHooks,
  };
}

const _atomFamilyStateCache = new Map<string, CompatWritableAtomState<any>>();
export type AtomFamilyOptions<T, P extends string> = {
  // For compatibility with Recoil
  key: string;
  default: (param: P) => T;
  effects?: (param: P) => readonly AtomEffect<T>[];
};
export function atomFamily<T, P extends string>(options: AtomFamilyOptions<T, P>) {
  const _family = jotaiAtomFamily((param: P) => {
    const jotaiState = jotaiAtom(options.default(param), (get, set, nextValue) => {
      if (nextValue instanceof Function) {
        set(jotaiState, nextValue(get(jotaiState)));
      } else {
        set(jotaiState, nextValue);
      }
    });
    return jotaiState;
  }, deepEqual);
  return (param: P): CompatWritableAtomState<T> => {
    const key = `${options.key}:${param}`;
    const cached = _atomFamilyStateCache.get(key);
    if (cached) {
      return cached;
    }
    const updateHooks: UpdateHook[] = [];
    options.effects?.(param).forEach((effect) => {
      effect({
        onSet: (fn) => {
          updateHooks.push(fn);
        },
      });
    });
    const state: CompatWritableAtomState<T> = {
      key,
      jotai: _family(param),
      default: () => options.default(param),
      type: "writable",
      updateHooks,
    };
    _atomFamilyStateCache.set(key, state);
    return state;
  };
}

export type GetCompatValue = <T>(val: CompatAtomState<T>) => T;
export type SelectorOptions<T> = {
  key: string;
  get: (opts: {
    get: GetCompatValue;
  }) => T;
};

export function selector<T>(options: SelectorOptions<T>): CompatReadbleAtomState<T> {
  return {
    key: options.key,
    jotai: jotaiAtom((jotaiGet) => {
      return options.get({
        get: (val) => {
          return jotaiGet(val.jotai);
        },
      });
    }),
    default: () => options.get({ get: (recoilVal) => recoilVal.default() }),
    type: "readble",
  };
}

const _selectorFamilyStateCache = new Map<string, CompatReadbleAtomState<any>>();
export type SelectorFamilyOptions<T, P extends string> = {
  key: string;
  get: (param: P) => (opts: {
    get: GetCompatValue;
  }) => T;
};
export function selectorFamily<T, P extends string>(options: SelectorFamilyOptions<T, P>) {
  const _family = jotaiAtomFamily((param: P) => {
    const jotaiState = jotaiAtom((jotaiGet) => {
      return options.get(param)({
        get: (val) => {
          return jotaiGet(val.jotai);
        },
      });
    });
    return jotaiState;
  }, deepEqual);
  return (param: P): CompatReadbleAtomState<T> => {
    const key = `${options.key}:${param}`;
    const cached = _selectorFamilyStateCache.get(key);
    if (cached) {
      return cached;
    }
    const state: CompatReadbleAtomState<T> = {
      key,
      jotai: _family(param),
      default: () => options.get(param)({ get: (recoilVal) => recoilVal.default() }),
      type: "readble",
    };
    _selectorFamilyStateCache.set(key, state);
    return state;
  };
}

export type CompatTransactionInterface = {
  get: <T>(a: CompatAtomState<T>) => T;
  set: <T>(a: CompatWritableAtomState<T>, u: ((currVal: T) => T) | T) => void;
  reset: (a: CompatWritableAtomState<any>) => void;
};
export type TransactionInterface_UNSTABLE = CompatTransactionInterface;
export type SerializableParam = string;
export const RecoilEnv: any = {
  RECOIL_DUPLICATE_ATOM_KEY_CHECKING_ENABLED: true,
};
export type GetRecoilValue = <T>(a: CompatAtomState<T>) => T;

export function useRecoilTransaction_UNSTABLE<Args extends ReadonlyArray<unknown>>(
  fn: (tInterface: CompatTransactionInterface) => (...args: Args) => void,
  // For compatibility with Recoil
  deps?: ReadonlyArray<unknown>,
) {
  const store = useStore();
  return useEvent((...args: Args): void => {
    const localStore = new Map<
      string,
      {
        atom: CompatAtomState<any>;
        state: any;
        updated?: boolean;
      }
    >();

    const get: CompatTransactionInterface["get"] = (a) => {
      const localValue = localStore.get(a.key);
      if (localValue) {
        return localValue.state;
      }
      const value = store.get(a.jotai);
      localStore.set(a.key, {
        atom: a,
        state: value,
      });
      return value;
    };

    const tInterface: CompatTransactionInterface = {
      get,
      set: (a, u) => {
        if (u instanceof Function) {
          localStore.set(a.key, { atom: a, state: u(get(a)), updated: true });
        } else {
          localStore.set(a.key, { atom: a, state: u, updated: true });
        }
      },
      reset: (a) => {
        localStore.set(a.key, { atom: a, state: a.default(), updated: true });
      },
    };

    fn(tInterface)(...args);

    localStore.forEach((value) => {
      if (value.atom.type === "writable" && value.updated) {
        store.set(value.atom.jotai, value.state);
        value.atom.updateHooks.forEach((hook) => {
          hook();
        });
      }
    });
  });
}

export type CompatSnapshot = {
  getPromise: <T>(a: CompatAtomState<T>) => Promise<T>;
};
export type CompatCallbackInterface = {
  set: <T>(a: CompatWritableAtomState<T>, u: ((currVal: T) => T) | T) => void;
  reset: (a: CompatWritableAtomState<any>) => void;
  snapshot: CompatSnapshot;
  transact_UNSTABLE: (fn: (tInterface: CompatTransactionInterface) => void) => void;
};
export function useRecoilCallback<Args extends ReadonlyArray<unknown>, Return>(
  fn: (cInterface: CompatCallbackInterface) => (...args: Args) => Return,
  deps?: ReadonlyArray<unknown>,
) {
  const callTransaction = useRecoilTransaction_UNSTABLE((tInterface) => {
    return (fn: (tInterface: CompatTransactionInterface) => void) => {
      fn(tInterface);
    };
  });
  const store = useStore();
  return useEvent((...args: Args): Return => {
    const cInterface: CompatCallbackInterface = {
      set: (a, u) => {
        store.set(a.jotai, u instanceof Function ? u(store.get(a.jotai)) : u);
        a.updateHooks.forEach((hook) => {
          hook();
        });
      },
      reset: (a) => {
        store.set(a.jotai, a.default());
        a.updateHooks.forEach((hook) => {
          hook();
        });
      },
      snapshot: {
        getPromise: async (a) => {
          return store.get(a.jotai);
        },
      },
      transact_UNSTABLE: callTransaction,
    };

    return fn(cInterface)(...args);
  });
}

export function useRecoilValue<T>(a: CompatAtomState<T>) {
  return useAtomValue(a.jotai);
}

export function useSetRecoilState<T>(a: CompatWritableAtomState<T>) {
  const _setter = useSetAtom(a.jotai);
  return useEvent((value: T | ((prev: T) => T)) => {
    _setter(value);
    a.updateHooks.forEach((hook) => {
      hook();
    });
  });
}

export function useRecoilState<T>(a: CompatWritableAtomState<T>) {
  const [value, _setter] = useAtom(a.jotai);
  const setter = useEvent((value: T | ((prev: T) => T)) => {
    _setter(value);
    a.updateHooks.forEach((hook) => {
      hook();
    });
  });

  return [value, setter] as const;
}