🦊

「作ればわかる」と思って簡易Zustandを作っていたらSolid.jsが脳裏にチラついた

に公開

Web Componentsっていいですよね(唐突)。

Web標準ゆえ大体のブラウザで動作が保証されており、いろいろな事情でNode.jsやTypeScriptを使うことができない環境でも、ある程度新しいブラウザさえあれば基本的には使えますし。
Markdown中にコンポーネントを書いたり、静的配信したりするのも簡単ですし。
「このscriptタグだけ書いておいて!」って頼めばHTMLしか触れない/触りたくない人でも比較的抵抗感なく使ってもらえますし。

あれ、でもなんだか状態管理したくなってきたぞ……

というわけで、Reactとともによく使われている印象があるシンプルな状態管理ライブラリ「Zustand」の実装を読んで、簡易Zustandを作ってみることにしました。

状態管理とは

「ハンバーガーメニューの表示・非表示」「現在ログインしているユーザー情報」「入力欄に入力された値」など、現代のWebサイトにはあらゆる状態が存在します。
これらを適切に更新したり、変更を追跡したりするには各状態をしっかり管理する必要があります。

例として、ラベルを持つ単純なチェックボックス2つと、それらの値を基に有効化/無効化されるボタンを考えてみましょう。

<input
  type="checkbox"
  id="question01"
>
<label for="question01">自分は新世界の神だと思う</label>
<input
  type="checkbox"
  id="question02"
>
<label for="question02">協力者には毎日ノートをチェックさせている</label>
<button
  type="button"
  id="declareButton"
>勝利を宣言する</button>

チェックボックスの値が変更される度にチェックボックスの状態を確認してボタンの状態を書き換える必要があります。

// それぞれの要素を取得
const question01 = document.querySelector("#question_01")
const question02 = document.querySelector("#question_02")
const declareButton = document.querySelector("#declareButton")

// チェックボックスの状態に基いてボタンの状態を変更する関数
const updateDeclareButtonAvailability = () => {
  const q01Checked = question01.checked ?? false
  const q02Checked = question02.checked ?? false
  // 両方がチェックされている場合のみボタンは有効
  const enabled = q01Checked && q02Checked
  declareButton.disabled = !enabled
}

// チェックボックスの変更がある度にボタンの状態を書き換える
question01.addEventListener("change", (ev) => {
  updateDeclareButtonAvailability()
}, false)
question02.addEventListener("change", (ev) => {
  updateDeclareButtonAvailability()
}, false)

// 最初の一回は自分で呼ぶ
updateDeclareButtonAvailability()

さらに入力欄やボタンが増え、管理すべき状態が増えると破綻しそうな匂いがしますね。
状態管理ライブラリを導入していい感じに状態を扱えるといいなあ、となります。

Zustandとは

状態管理ライブラリです。TypeScriptで開発されており、Reactに依存しています。

import { create } from 'zustand'

const useBearStore = create((set) => ({
  bears: 0,
  increasePopulation: () => set((state) => ({ bears: state.bears + 1 })),
  removeAllBears: () => set({ bears: 0 }),
}))
function BearCounter() {
  const bears = useBearStore((state) => state.bears)
  return <h1>{bears} around here ...</h1>
}

function Controls() {
  const increasePopulation = useBearStore((state) => state.increasePopulation)
  return <button onClick={increasePopulation}>one up</button>
}

https://github.com/pmndrs/zustand

しかし、Reactに依存しているのはuse〜系のHooks APIです。use〜系の関数は状態とその状態を管理する関数(状態の取得や更新など)を作成するcreateStore関数をラップして使いやすくした関数たちです。
コアの機能であるcreateStorecreateから呼ばれます)はReactに非依存であり、その使い方と同じように仕組みも比較的シンプルなので、作ってみようというわけです。

createStoreを作る

createStoreは初期値を与える関数createStateを引数に取って状態を作ったあと、その状態を管理する関数を1つのオブジェクトにまとめて返します。

// createStoreで状態を作る
// storeは状態の取得や更新をするための関数を詰め込んだオブジェクト
const store = createStore((set) => ({
  bears: 0,
}))

// 現在の状態を表示
console.log(store.getState())

// 状態が変更されたら呼ばれる関数(subscriber)を登録しておく
store.subscribe((state) => {
  console.log(state)
})

// 状態を書き換えるとsubscriberが呼ばれる
setTimeout(() => {
  store.setState((state) => ({ bears: state.bears + 1 }))
}, 3000)

まずは値の取得と更新ができるように

とはいえいきなりこれを実装するのは大変なので、まずは状態の取得と更新ができることを目指します。
つまり、状態を取得する関数getStateと状態を更新する関数setStateを定義して、それらをオブジェクトapiに詰めて返す関数createStoreを書く、ということです。
createStoreは初期状態を与える関数createStateを引数に取り、状態を初期化します。

// 引数のcreateStateは初期状態を与える関数
const createStore = (createState) => {
  // 状態を保持する変数
  let state;

  // 状態を取得する関数
  const getState = () => state
  // 状態を更新する関数
  const setState = (nextState) => {
    state = nextState
  }
  // オブジェクトに詰める
  const api = {
    getState,
    setState,
  }
  // createStateを呼び出して状態を初期化する
  state = createState(setState, getState, api)
  // 状態を操作する関数をオブジェクトにまとめたものを返す
  return api
}

このcreateStoreは以下のように使います。

// createStateを渡してstoreを作る
const store = createStore((setState, getState, api) => ({
  count: 0,
}))

// 状態を取得してログに出力
console.log(store.getState())

// 3秒後に状態を更新
setTimeout(() => {
  // 現在の状態を取得
  const state = store.getState()
  // 更新後の状態を作成
  const nextState = {
    ...state,
    count: state.count + 1,
  }
  // 状態を更新する
  store.setState(nextState)
  // 更新後の状態が現在の状態になっているはずなのでログに出力する
  console.log(store.getState())
}, 3000)

Subscriberを呼ぶ

状態の取得と書き換えはできたので、状態を書き換えたときに呼ばれる関数 subscriber を登録できるようにします。

const createStore = (createState) => {
  const subscribers = new Set()
  let state;
  // 現在の状態を返す関数
  const getState = () => state
  // 状態を変更する関数
  const setState = (nextState) => {
    const prevState = state
    state = nextState
    // 変更を通知する
    for (const subscriber of subscribers) {
      subscriber(state, prevState)
    }
  }
  // 状態が変更されたときに実行されるsubscriberを登録する関数
  const subscribe = (subscriber) => {
    subscribers.add(subscriber)
    const unsubscribe = () => {
      subscribers.delete(subscriber)
    }
    return unsubscribe
  }
  // 状態を操作する関数をまとめる
  const api = {
    getState,
    setState,
    subscribe,
  }
  // createStateを呼び出して状態を初期化する
  state = createState(setState, getState, api)
  // 状態を操作する関数をまとめたものを返す
  return api
}

しかしsetStateが呼ばれると問答無用でsubscriberが呼ばれるのはあまり賢くありません。
setStateする際に、状態の更新前後で変更があったかを調べるようにしましょう。

const setState = (nextState) => {
  if (Object.is(state, nextState )) {
    // 変更無しなら何もしない
    return
  }
  // あとは一緒
  const prevState = state
  state = nextState
  for (const subscriber of subscribers) {
    subscriber(state, prevState)
  }
}

subscriberを登録できるようになったので、更新時の処理を分けて書くことができます。

const store = createStore((setState, getState, api) => ({
  count: 0,
}))

// 状態を取得してログに出力
console.log(store.getState())

// 更新するごとに現在の状態をログに出力する
store.subscribe((state) => {
  console.log(state)
})

// 3秒後に状態を更新
setTimeout(() => {
  // 現在の状態を取得
  const state = store.getState()
  // 更新後の状態を作成
  const nextState = {
    ...state,
    count: state.count + 1,
  }
  // 状態を更新する
  store.setState(nextState)
}, 3000)

差分更新できるように

オブジェクトの更新をもっと楽に書けるよう、差分更新にします。

// 現在の状態を取得
const state = store.getState()
// 状態を更新する
store.setState({
  // 差分を与えればよしなにしてくれる
  count: state.count + 1,
})

そのためにまず、Objectかどうか調べる関数を定義します。

const isObject = (x) => (
  typeof x === 'object' &&
  x !== null
)

setStateで変更があった場合にObject.assign()で元の状態とマージしてあげます。
オマケでsetStateに関数を渡せるようにします。

// 状態を変更する関数
const setState = (partial) => {
  // partialが関数なら現在の状態を渡して実行する
  const nextState = (typeof partial === 'function') ? partial(state) : partial
  // 変更なしなら何もしない
  if (Object.is(state, nextState)) { return }
  const prevState = state
  // オブジェクトなら変更部分を現在の状態とマージする
  state = isObject(nextState) ? Object.assign({}, state, nextState) : nextState
  // あとは一緒なので省略
}

subscribeWithSelectorsを作る

ここまでで基本的な機能はほぼできましたが、この手の状態管理ではどうしてもstoreが大きくなりがちです。

そこで、store.subscribe()でサブスクライブするときにstoreの一部のみに絞るselectorを指定できるようにします。

const store = createStore(
  // selectorを指定できるようにする準備
  subscribeWithSelectors((set, get) => ({
    bears: 0,
    rabbits: 0,
    total: () => get().bears + get().rabbits,
    incrBears: () => set(({ bears }) => ({ bears: bears + 1 })),
    incrRabbits: () => set(({ rabbits }) => ({ rabbits: rabbits + 1 })),
  }))
)

// selectorたち
const selectBears = ({ bears }) => bears
const selectRabbits = ({ rabbits }) => rabbits
const selectTotal = ({ total }) => total

console.log(store.getState())

// subscribeするときにselectorを指定できる
store.subscribe(selectBears, (bears) => {
  console.log(`bears = ${bears}`)
})
store.subscribe(selectRabbits, (rabbits) => {
  console.log(`rabbits = ${rabbits}`)
})
store.subscribe(selectTotal, (total) => {
  console.log(`total = ${total()}`)
})

// もちろんstore全体を対象にすることも可
store.subscribe((state) => {
  console.log(state)
})

setTimeout(() => {
  store.getState().incrRabbits()

  setTimeout(() => {
    store.getState().incrBears()
  }, 3000)
}, 3000)

subscribeWithSelectorsという関数がキモで、createStoreで作られたapi.subscribestore.subscribeとして呼び出されるメソッド)を置き換えています。

const subscribeWithSelectors = (createState) => (set, get, api) => {
  const orgSubscribe = api.subscribe
  api.subscribe = (selectorOrSubscriber, optSubscriber) => {
    const selector = optSubscriber ? selectorOrSubscriber : (x => x)
    const subscriber = optSubscriber ?? selectorOrSubscriber
    // selectorを適用したsubscriberに置き換える
    return orgSubscribe((state, prevState) => {
      const selectedState = selector(state)
      const selectedPrevState = selector(prevState)
      if (Object.is(selectedState, selectedPrevState)) { return }
      subscriber(selectedState, selectedPrevState)
    })
  }
  return createState(set, get, api)
}

Solid.js「こんにちは!」

余談ですが、実はcreateStoreの実装を見た時から脳裏でSolid.jsが顔を覗かせていました。

というのも以前Solid.jsの実装をちょっとだけ読んでみたことがあり、その際にも「クロージャ内に状態を保存してgetterとsetterを返す」という仕組みが使われていたからです。

https://zenn.dev/fj68/articles/da3458abab5d1c

上の記事で実装したSolid.js風のcreateSignal()をJavaScriptに書き直してちょっとcreateStoreに寄せてみるとこんなコードになります。
getterがちょっと特殊ですが、これは自動でsubscribeするためのトリックで、「クロージャに保存された現在の状態を返す」という機能は変わりません。

const createSignal = (initialState) => {
  const subscribers = new Set()
  let state = (typeof initialState === 'function') ? initialState() : initialState

  // 現在の状態を返す関数
  const getter = () => {
    const subscriber = gContexts.top()
    if (subscriber) {
      // コンテクストにsubscriberがあればsubscribeする
      subscribers.add(subscriber)
    }
    return state
  }
  // 状態を変更する関数
  const setter = (nextState) => {
    if (Object.is(state, nextState)) {
      // 変更されていなければ更新しない
      return
    }
    // 状態を変更してsubscriberに変更を通知する
    state = nextState
    for (const subscriber of subscribers) {
      subscriber()
    }
  }
  // 状態を操作するための関数を返す
  return [getter, setter]
}

使い方はこうです。

const [count, setCount] = createSignal(0)
createEffect(() => {    // createEffectの実装は省略
  console.log(`count = ${count()}`)
})
// 値を書き換えると、createEffectに渡した関数が勝手に実行される
setTimeout(() => {
  setCount(count() + 1)
}, 3000)

初期値の指定方法やgetter/setterの返し方、Object.assignで差分更新するかどうかなどの細かい違いはあれど、状態をクロージャに保存するのは一緒です。

ごん、お前だったのか……

おわり

実際のZustandではまだまだやっていることがあり、比較関数をObject.is以外にできるようにオプションで設定できるようになっていたりしますが、とりあえず基本的な仕組みはわかったのでここまでとします。

Discussion