RecoilをRxJSで再実装する #React

動機

Recoilを見ていて、Atom/Selectorという単位を用意した上での非同期更新伝播の仕組みが『自由度をうまく狭めたRx』だなあ、という印象を受けた。
APIとしては React & Hooks環境が大前提になってるようだけど、上記の印象があるので、理想的には特にReact & Hooksべったりじゃなくてもいいのでは（Reduxと同様にstandaloneで利用できるのでは）と考えた。
現状のRecoilのソースを見たら、State管理のライブラリにいきなりReactDOMが出てきたりして辛いので、あまり見てはいけない気がした（型定義だけはとりあえず信用する）。

基本方針

atom/selectorの結果得られるRecoilStateを、RxJSのstreamをwrapするオブジェクトとして実装する。各stateの値の更新はstreamに流される
RecoilState はhooks以外でも使えるよう、内包するstreamをsubscribeできるようにしている
useRecoilState/useRecoilValueなどのhooksのAPIから、RecoilStateの現在の値の取得および非同期未解決の場合はPromise値をthrowする関係上、解決された値を取得できるだけでなく非同期処理(Promise)が実行中か実行後かを分かる必要がある。
- streamにはLoadableという名前の「状態(isLoading/hasValue/hasError)と値/Promiseのペア」を流す
- RxJSではBehaviorSubjectを使うことで最後にstreamに流れた値をいつでも取得できる
各hooks(useRecoilState/useRecoilValue)は、useState()でローカルのstateを保持し、useEffect()内でstreamをsubscribeする。streamから値が流れてきたらその値でローカルstateを更新する（これによりReactで再描画が発生する）

実装

recoil-rxjs/recoil.ts

import { SetStateAction } from "react";
import { BehaviorSubject } from "rxjs";
import { skip } from "rxjs/operators";

/**
 *
 */
function getNextState<T>(action: SetStateAction<T>, getPrevState: () => T) {
  if (typeof action === "function") {
    const prevState = getPrevState();
    return (action as ((prevState: T) => T))(prevState);
  }
  return action;
}

/**
 * 
 */
function isPromise<T>(value: T | Promise<T>): value is Promise<T> {
  return typeof (value as any)?.then === 'function';
}

/**
 * this interface is not in official recoil, but introduced for the ease of implementation.
 */
export interface RecoilState<T> {
  key: string;
  getValue: () => T;
  setValue: (action: SetStateAction<T>) => void;
  getPromise: () => Promise<T>;
  subscribe: (listener: (value: T) => void) => () => void;
}

/**
 * global states
 */
const states = new Map<string, RecoilState<any>>();

/**
 *
 */
type AtomInput<T> = {
  key: string;
  default: T;
};

/**
 *
 */
export function atom<T>(input: AtomInput<T>): RecoilState<T> {
  const { key, default: defaultValue } = input;
  let state: RecoilState<T> | undefined = states.get(key);
  if (state) {
    return state;
  }
  const stream$ = new BehaviorSubject<T>(defaultValue);
  state = {
    key,
    getValue() {
      return stream$.value;
    },
    setValue(action: SetStateAction<T>) {
      const value = getNextState(action, () => this.getValue());
      stream$.next(value);
    },
    getPromise() {
      return Promise.resolve(stream$.value);
    },
    subscribe(listener: (value: T) => void) {
      const sbscr = stream$.pipe(skip(1)).subscribe(listener);
      return () => sbscr.unsubscribe();
    }
  };
  states.set(key, state);
  return state;
}

/**
 *
 */
type SelectorGetHelper = {
  get<T>(state: RecoilState<T>): T;
  getPromise<T>(state: RecoilState<T>): Promise<T>;
};

type SelectorSetHelper = {
  get<T>(state: RecoilState<T>): T;
  set<T>(state: RecoilState<T>, value: T): void;
};

type SelectorInput<T> = {
  key: string;
  get: (helper: SelectorGetHelper) => T | Promise<T>;
  set?: (helper: SelectorSetHelper, value: T) => void;
};

type LoaderState<T> =
  | {
      status: "isLoading";
      promise: Promise<T>;
    }
  | {
      status: "hasValue";
      contents: T;
    }
  | {
      status: "hasError";
      contents: any;
    };

/**
 *
 */
export function selector<T>(input: SelectorInput<T>): RecoilState<T> {
  const { key, get: getValue, set: setValue } = input;
  let state: RecoilState<T> | undefined = states.get(key);
  if (state) {
    return state;
  }
  let stream$: BehaviorSubject<LoaderState<T>> | undefined = undefined;
  const toLoaderState = (value: T | Promise<T>): LoaderState<T> => {
    return isPromise(value) ? {
      status: "isLoading",
      promise: value
    } : {
      status: "hasValue",
      contents: value,
    };
  };
  const isLoaderStateInSamePromise = (promise: Promise<T>) => {
    if (!stream$) {
      return false;
    }
    const currentLoader = stream$.value;
    return currentLoader.status === 'isLoading' && currentLoader.promise === promise;
  };
  const handleLoaderStateUpdate = (loader: LoaderState<T>) => {
    stream$?.next(loader);
    if (loader.status === 'isLoading') {
      const promise = loader.promise;
      promise.then(
        value => {
          if (isLoaderStateInSamePromise(promise)) {
            stream$?.next({
              status: "hasValue",
              contents: value
            });
          }
        },
        error => {
          if (isLoaderStateInSamePromise(promise)) {
            stream$?.next({
              status: "hasError",
              contents: error
            });
          }
        }
      );
    }
  };
  const createLoaderStateStream = () => {
    const value = getValue(getHelper);
    const loader = toLoaderState(value);
    return new BehaviorSubject<LoaderState<T>>(loader);
  };
  const updateState = () => {
    if (!stream$) {
      stream$ = createLoaderStateStream();
    } else {
      const value = getValue(getHelper);
      const loader = toLoaderState(value);
      handleLoaderStateUpdate(loader);
    }
  };
  const deps: Set<string> = new Set();
  const getHelper = {
    get<U>(state: RecoilState<U>) {
      if (!deps.has(state.key)) {
        state.subscribe(updateState);
        deps.add(state.key);
      }
      return state.getValue();
    },
    getPromise<U>(state: RecoilState<U>) {
      if (!deps.has(state.key)) {
        state.subscribe(updateState);
        deps.add(state.key);
      }
      return state.getPromise();
    }
  };
  const setHelper = {
    get<U>(state: RecoilState<U>) {
      if (!deps.has(state.key)) {
        console.log("subscribing", key, state.key);
        state.subscribe(updateState);
        deps.add(state.key);
      }
      return state.getValue();
    },
    set<U>(state: RecoilState<U>, value: U) {
      return state.setValue(value);
    }
  };
  state = {
    key,
    getPromise() {
      if (!stream$) {
        stream$ = createLoaderStateStream();
      }
      const loader = stream$.value;
      if (loader.status === "isLoading") {
        return loader.promise;
      }
      if (loader.status === "hasError") {
        return Promise.reject(loader.contents);
      }
      return Promise.resolve(loader.contents);
    },
    getValue() {
      if (!stream$) {
        stream$ = createLoaderStateStream();
      }
      const loader = stream$.value;
      if (loader.status === "isLoading") {
        throw loader.promise;
      }
      if (loader.status === "hasError") {
        throw loader.contents;
      }
      return loader.contents;
    },
    setValue(action: SetStateAction<T>) {
      if (setValue) {
        const value = getNextState(action, () => this.getValue());
        setValue(setHelper, value);
      }
    },
    subscribe(listener: (value: T) => void) {
      if (!stream$) {
        stream$ = createLoaderStateStream();
      }
      const sbscr = stream$.subscribe(loader => {
        if (loader.status === "hasValue") {
          listener(loader.contents);
        }
      });
      return () => sbscr.unsubscribe();
    }
  };
  states.set(key, state);
  return state;
}

recoil-rxjs/hooks.ts

import {
  useState,
  useCallback,
  useEffect,
  useMemo,
  SetStateAction
} from "react";
import { RecoilState } from "./recoil";

/**
 *
 */
export function useRecoilState<T>(state: RecoilState<T>) {
  const [value, setRawValue] = useState(state.getValue());
  const setValue = useCallback(
    (action: SetStateAction<T>) => {
      state.setValue(action);
    },
    [state]
  );
  useEffect(() => {
    return state.subscribe(setRawValue);
  }, [state]);
  return [value, setValue] as const;
}

/**
 *
 */
export function useRecoilValue<T>(state: RecoilState<T>) {
  const [value, setRawValue] = useState(state.getValue());
  useEffect(() => {
    return state.subscribe(setRawValue);
  }, [state]);
  return value;
}

/**
 *
 */
export function useSetRecoilState<T>(state: RecoilState<T>) {
  const setValue = useCallback(
    (action: SetStateAction<T>) => {
      state.setValue(action);
    },
    [state]
  );
  return setValue;
}

/**
 *
 */
type CallbackInterface = {
  getPromise<U>(state: RecoilState<U>): Promise<U>;
  set<U>(state: RecoilState<U>, action: SetStateAction<U>): void;
};

export function useRecoilCallback<Args extends any[], R>(
  callback: (helper: CallbackInterface, ...args: Args) => R,
  deps: any[]
) {
  const helper = useMemo(
    () => ({
      getPromise<U>(state: RecoilState<U>) {
        return state.getPromise();
      },
      set<U>(state: RecoilState<U>, action: SetStateAction<U>) {
        return state.setValue(action);
      }
    }),
    []
  );
  return useCallback(
    (...args: Args) => {
      return callback(helper, ...args);
    },
    [helper, ...deps] // eslint-disable-line
  );
}

コード

Codesandboxで実際に動くやつ。アプリはuhyoさんの記事でつかっていたやつから拝借。
https://codesandbox.io/s/stoic-drake-2f64z

所感

各selectorでmemoizationの必要があるとおもうけど、省いている（ただしそこもRxJSでできるはず）。
RecoilRootはこの実装では必要なかったやつ。そもそもglobalにatom/selectorを定義しちゃうのに階層でstoreを局所化する意味とは何だ？storeの構成は共通でもstoreのstateは別々にしたい、というユースケースがあるのかどうか？
なお久々にRxJS使ったらstreamにメソッドが生えているのではなく.pipe()でoperationをつなぐようになっていて、バンドルサイズ小さくなるよう工夫されてるなと思った