【Jotai】Reactの状態管理はコレで決まり！！！

この記事は Optimind Advent Calendar 2023 の3日目の記事となります。

はじめに

おそらく、この記事を開いていただいた方は「Reactの状態管理ライブラリって何を使えばいいんだろう？」とお悩みなのではないでしょうか？

まぁ、ちょっとググっただけでもいっぱい出てきますしね…。

私もここ最近しっくり来るものを探していて、ようやく素敵なライブラリに出会えたので、その紹介をしようと思います。

本編

選定の基準

状態管理ライブラリを選定するにあたって、以下の要件を満たせるかどうかを基準にしました。

軽量であること

フロントエンドアプリケーションにおいて、初期ロードの速さはとても重要です。
サイズが軽いに越したことはありません。

シンプルかつミニマムに記述できること

多機能過ぎてごちゃごちゃしていたり、導入が大変だったり、ボイラープレートが多いライブラリは様々な悲しみを生みます （あぁ…辛い記憶が…）

計算回数を最小限にできること

参照していないstateが変更された時に、無駄な再計算・再レンダリングが起きてしまうのは です。

また、共通のselectorを色んなコンポーネントで使う場合に、参照しているコンポーネント全てで同じ計算が走ってしまう…みたいなことも避けたいです。

Jotaiのココがいい！

探索の結果、一番しっくり来たのが Jotai というライブラリでした。
名前からも分かる通り、日本人の方が開発してくださっているようです。

軽量であること

コアAPIのみであれば 2kb という軽量さ…すごい！

シンプルかつミニマムに記述できること

Atom という最小単位でstateを管理できます。Recoilを使ったことがある方には馴染み深いでしょうか。

初期設定については、Providerですら基本不要という…。マジ？

計算回数を最小限にできること

Atom単位で管理するので、計算・レンダリングコストも最小化しやすいですし、selectorという概念はほぼ不要になります。

また、stateに対して共通の計算を複数のコンポーネントで使う場合も、依存stateの変更に対し1回の計算で済む点もポイント高いです。

Vite + SWCに対応している

（公式がサポートしてくれていたらいいなぁ程度に思っていたのに）
Viteに対応しているだけでもありがたいのですが、SWCにも対応しているとか神ですか…？

テストの書きやすさ

後述する store という機能を使うことで、初期値のセットや変更後の値を取得したりが簡単にできます。mockの手間が省けるのはありがたい…

Let's Jotai

それでは、実際に使ってみましょう！
今回のコードは stackblitz で公開していますので、Forkして触ってみてください！

インストール・初期設定

npm i jotai

基本的にはコレだけでOKです。

react-refreshも対応したい方

デフォルトではホットリロードするとstateの値が毎回揮発してしまいます。
react-refreshを導入すると、stateを保持するようにできて便利です。

この記事ではVite + SWCを想定しているので、SWC用のものをインストールします。
Babel用が必要な方は公式ドキュメントを参照してください。

npm i -D @swc-jotai/react-refresh

vite.config.ts

export default defineConfig({
  plugins: [
    // reactプラグインの引数に以下のようにオプションを追加
    react({
      plugins: [['@swc-jotai/react-refresh', {}]],
    }),
  ],
  ....

はい、本当にこれだけで設定完了です。

既存のuseStateを置き換えてみる

ReactテンプレートのカウンターをJotaiで置き換えてみます。

state.ts

import { atom } from 'jotai';

// 引数が初期値になります
export const countState = atom(0);

App.tsx

// import { useState } from 'react';
import { useAtom } from 'jotai';
import { countState } from './state';

const App = () => {
  // const [count, setCount] = useState(0);
  const [count, setCount] = useAtom(countState);

  return (
    ....

コレだけでOKです。
storeやactions、selectorの定義も不要ですね。

Read-Write Atom

JotaiのAtomには、先ほどのような初期値のみを指定する primitive Atom の他にも3種類あります。

Read-Only Atom

計算や、複数のAtomを一つにまとめたりできます。

state.ts

export const doubleCountState = atom((get) => get(countState) * 2);

Write-Only Atom

複数のAtomに一括で値をセットしたりできます。

state.ts

export const incrementCountAction = atom(
  null,
  (get, set) => set(countState, get(countState) + 1),
);

Read-Write Atom

ReadとWriteを両方使うことも可能です。

// 必ず一緒に値を変更する値は1つのAtomにまとめると、無駄な再レンダリングを防げます
export const fetchDataState = atom({
  data: null,
  isLoading: false,
});

export const fetchDataAction = atom(
  // 同じコンポーネント内でデータも扱うなら一緒にした方がimportを減らせます
  (get) => get(fetchDataState),
  // 非同期もOK
  async (get, set, id: string) => {
    set(fetchDataState, {
      data: get(fetchDataState).data,
      isLoading: true,
    });
    
    const data = await queryData(id);
    
    set(fetchDataState, {
      data,
      isLoading: false,
    });
  },
);

const [{ data, isLoading }, fetchData] = useAtom(fetchDataAction);

これを使用して最初のコードをリファクタすると以下のようになります。

state.ts

export const incrementCountAction = atom(
  (get) => get(countState),
  (get, set) => set(countState, get(countState) + 1),
);

App.tsx

import { useAtom } from 'jotai';
// import { countState } from './state';
import { incrementCountAction } from './state';

const App = () => {
  // const [count, setCount] = useAtom(countState);
  const [count, incrementCount] = useAtom(incrementCountAction);

  return (
    ....

      {/* <button onClick={() => setCount(count + 1)}>count is {count}</button> */}
      <button onClick={incrementCount}>count is {count}</button>
    ....

stateを更新する処理がコンポーネントから切り出せました。

Write Atomの検証

Jotai の強みとして

また、stateに対して共通の計算を複数のコンポーネントで使う場合も、依存stateの変更に対し1回の計算で済む点もポイント高いです。

と書いていましたが、それが本当かどうか検証してみます。

例の countState を2倍にして返すAtom doubleCountState を使います。
呼ばれた回数が分かるようにログを仕込んでおきましょう。

state.ts

export const doubleCountState = atom((get) => {
  const result = get(countState) * 2;

  console.log('doubleCountState called!', result);

  return result;
});

これを適当に作ったコンポーネントから参照するようにします。

DoubleCount.tsx

import { useAtomValue } from 'jotai';
import { doubleCountState } from './state';

export const DoubleCount = () => {
  // Read部分のみを使いたい場合には、useAtomValueを使うと良いです
  const doubleCount = useAtomValue(doubleCountState);

  return <div>doubleCount: {doubleCount}</div>;
};

さらに、このコンポーネントを App.tsx へ大量に配置します。

App.tsx

   ....
      <h1>Vite + React</h1>
      <div className="card">
        <button onClick={() => setCount(count + 1)}>count is {count}</button>
        <DoubleCount />
        <DoubleCount />
        <DoubleCount />
   ....

これでもし各コンポーネント毎に計算が走っていれば、increment毎に3回ログが書き込まれてしまうはずです。

５回クリックしました。果たして……

おお！問題なさそうですね〜。increment1回に対し1回しか計算が行われていません！
（一番上はホットリロードによるものです）

Provider と store

前章までで Jotai の基本的な使い方を見てきましたが、Provider は一度も登場していません。それは、1つのAtomを複数箇所で別々に扱う必要がなかったからです。
例えば、コンポーネント毎にカウンターを独立して持たせたい場合などにProviderが必要になってきます。

試しに、カウント部分を Counter.tsx に切り出して、Providerの有無でどんな変化が起きるか確認してみましょう。

Counter.tsx

import { useAtom } from 'jotai';
import { incrementCountAction } from './state';
import { DoubleCount } from './DoubleCount';

export const Counter = () => {
  const [count, incrementCount] = useAtom(incrementCountAction);

  return (
    <div className="card">
      <button onClick={incrementCount}>count is {count}</button>
      <DoubleCount />
    </div>
  );
};

App.tsx

   ....
      </div>
      <Counter />
      <Counter />
      <p className="read-the-docs">
   ....

Providerを使っていない状態では、全てのコンポーネントで同じ値が参照されていることが分かるかと思います。

では、<Counter /> の一つをProviderで囲ってみましょう。

App.tsx

   ....
      </div>
      <Counter />
      <Provider>
        <Counter />
      </Provider>
      <p className="read-the-docs">
   ....

Providerで囲った部分のカウントが独立しているのが分かるかと思います。
このように、同じAtomでも別々のProvider配下では別々のStateとして扱うことができます。
これは、ReactのProvider・Contextと同じ挙動ですね。

また、storeをProviderに渡すことで初期値を新たに設定することができます。

App.tsx

....

const counterStore = createStore();
// カウンターの初期値を1000に設定している
counterStore.set(countState, 1000);

const App = () => {
....
      <Provider store={counterStore}>
        <Counter />
      </Provider>
....

ちゃんと、初期値が1000になっていますね！
（手動で1000回クリックしたわけじゃないですよ？）

テストを書いてみる

せっかくなので、テストも書いてみましょう。
先ほどのProvider+storeを使うと簡単に書けます。

__tests__/Counter.test.tsx

import '@testing-library/jest-dom';
import { render } from '@testing-library/react';
import { userEvent } from '@testing-library/user-event';
import { Provider, createStore } from 'jotai';

import { countState } from '../state';
import { Counter } from '../Counter';

describe('Counter', () => {
  test('ボタンをクリックしたときにincrementされること', async () => {
    // 初期値の設定
    const store = createStore();
    store.set(countState, 999);

    const { getByText, getByRole } = render(
      <Provider store={store}>
        <Counter />
      </Provider>
    );

    // 初期値がセットされていることを確認
    expect(getByText('count is 999')).toBeInTheDocument();

    await userEvent.click(getByRole('button'));

    // 描画される値なら描画されているかをテストする
    expect(getByText('count is 1000')).toBeInTheDocument();
    expect(getByText('doubleCount: 2000')).toBeInTheDocument();

    // もし描画されないstateの変更をテストしたい場合はstoreから取得できる
    expect(store.get(countState)).toBe(1000);
  });
});

実行してみます。 npm run test

無事に通りましたね。
storeの優秀なところは、変更後の値を取得できることです。

    // もし描画されないstateの変更をテストしたい場合はstoreから取得できる
    expect(store.get(countState)).toBe(1000);

これによって、reducerやdispathなんかをモックして mock.calls をテストするという面倒から解放されます。

Provider配下で一部のatomだけデフォルトを参照したい

前述の通り、Provider配下では完全にAtomが別々のものと扱われます。
しかし、例えば画面全体をローディング状態にするAtomなどは、Provider配下であろうと共通の値を参照したい場合もありそうです。

そんなときは、useAtom() の第二引数 options に { store: getDefaultStore() } を渡してあげることで実現できます。

ただ、毎回やるのは面倒なのでhooksに切り出しておくと便利かもしれません。

hooks/useDefaultStore.tsx

import { getDefaultStore, useAtom } from 'jotai';

const defaultStore = getDefaultStore();

type UseAtomParams = Parameters<typeof useAtom>;

export const useAtomDefault: typeof useAtom = (
  atom: UseAtomParams[0],
  options?: UseAtomParams[1]
) =>
  useAtom(atom, {
    store: defaultStore,
    ...options,
  });

Counter.tsx

  ....
  const [defaultCount] = useAtomDefault(countState);

  return (
  ....

      <DoubleCount />
      <div>defaultCount: {defaultCount}</div>
    </div>

  ....

これで、Provider配下でもデフォルトの値を参照することができました。

ブラウザのStorageに値を保存

Webアプリを開発していると、どうしてもリロード間で値を共有したい要件が出てきます。
例えば、ページ遷移後にトーストを出すためのフラグであったり、ダークモード切り替えフラグであったりです。

Jotai はそんな痒い所に手が届く機能もデフォルトで提供してくれています。

atomWithStorage(key: string, init: any) で定義したAtomの値はStorageに保持されるようになります。具体例は 公式サイト に載っています。

公式サイトには紹介しきれなかった面白い機能が沢山載っています。是非ご一読ください！

おわりに

この度は、この記事を読んでいただきありがとうございました。
少しでも Jotai の良さが伝わっていましたら幸いです

// TODO: 寝る
const myState = atom('tired');