Немного предыстории

Недавно давно я смотрел ничем не примечательный техническое интервью и услышал фразу от интервьюируемого: «Ну можно написать свой useReducer или useState». Мне врезалась эта фраза в голову, ибо я никогда в серьез не задумывался как они работают под капотом и в исходниках особо не копался, максимум в типах. Из-за этого задача оказалась довольно сложной и интересной для меня ибо много получил много новой информации за довольно короткий срок и ее было сложно переварить и осознать.

Начало изысканий

Решил я начать эту задачу с хука useReducer так как он по сути основа для написания useState. Оглядываясь назад я думаю что этот подход был не верным ибо идти от верхнего уровня вниз постепенно заменяя зависимости на свои. К этому подходу я позже и пришел так как слишком сложно сразу писать код избегая всех зависимостей. Поэтому погружаться мы с вами будем постепенно.

useState

В моем понимании useState - это обертка над useReducer с уже заданным редюсером, который принимает callback функцию для подсчета нового состояние или просто новое значение. Поэтому единственное что придется сделать это подменить зависимости на свои.

import { Dispatch, Reducer, useReducer } from "react";

export type SetStateAction<State> = State | ((prev: State) => State);

export const useState = <State>(
    initialValue: State
  ): [State, Dispatch<SetStateAction<State>>] => {
    const reducer: Reducer<State, SetStateAction<State>> = (
      state,
      payload
    ): State => {
      if (typeof payload === "function")
        return (payload as (prev: State) => State)(state);
      return payload;
    };

    const [state, dispatch] = useReducer(
      reducer,
      initialValue
    );

    return [state, dispatch];
  };

Думаю тут все просто и задерживаться тут нет смысла.

useReducer

По началу я решил написать просто функцию, принимающую редюсер и исходное состояние и возвращающую актуальное состояние и функцию диспетчер, которая применяет ее к актуальному состоянию изменяя его.

function useReducer(reducer, initialValue) {
  let state = initialValue;

  const dispatch = (action) => {
    if (typeof state === "object")
      state = Object.assign(state, reducer(state, action));
    else state = reducer(state, action);
  };

  return [state, dispatch];
}

// Этот редюсер будет везде
const reducer = (state, action) => {
  switch (action.type) {
    case "increment":
      return { count: state.count + 1 };
    case "decrement":
      return { count: state.count - 1 };
    case "reset":
      return { count: 0 };
    default:
      return state;
  }
};

const [state, dispatch] = useReducer(reducer, { count: 0 });

console.log(state); // 0
dispatch({ type: "increment" });
dispatch({ type: "increment" });
console.log(state); //ispatch({ type: "decrement" });
dispatch({ type: "decrement" });
console.log(state); // 0

Да в логах конечно функция работает нормально, но для того что бы наш возвращаемый state менялся, он должен быть реактивным, либо косить под эту самую реактивность через паттерны Pub/Sub или Observer.

Мне стало интересно как работает в реакте эта самая реактивность, и я полез в исходники, реализация всех хуков находится в пакете react-reconciler в файле ReactFiberHooks.js. Реакт сохраняет состояние всех хуков в глобальных переменных currentHook и workInProgressHook в формате связанных списков. И сами хуки работают на двух функциях mountWorkInProgressHook и mountWorkInProgressHook которые создают и обновляют информацию о хуках включая состояния между рендерами.

export type Hook = {
  memoizedState: any,
  baseState: any,
  baseQueue: Update<any, any> | null,
  queue: any,
  next: Hook | null,
};

// Hooks are stored as a linked list on the fiber's memoizedState field. The
// current hook list is the list that belongs to the current fiber. The
// work-in-progress hook list is a new list that will be added to the
// work-in-progress fiber.
let currentHook: Hook | null = null;
let workInProgressHook: Hook | null = null;

Так как я не могу засунуть свой хук в этот связанный список, что бы скопировать принцип работы под капотом любого встроенного хука я решил посмотреть как работают другие библиотеки с хуками хранящими состояниями в реакт.

Порывшись в исходниках zustand, я обнаружил что их подход создания хранилища следующий: они делают свой стор в обычном js (в исходниках файл vanila.ts), а с реактом их стор работает через хук useSyncExternalStore, который подписывает реакт на события сторонних api в обычном js.

При помощи этого хука можно создать не только хранилища, но и хуки взаимодействия с браузерным api подробнее можно почитать в этой статье или в документации.

Но перед тем как делать его в react я сначала попробовал сделать что-то подобное в обычном js и пришел к такому решению.

После реализации в обычном js я решил добавить уровень абстракции что бы можно было реализовать не только useReducer, но и стор который будет чем-то средним между redux стором и zustand, естественно в упрощенном формате. Я написал функцию createStore, которая принимает те же самые аргументы что и useReducer, а возвращает все входные параметры для useSyncExternalStore и функцию dispatch для мутации стора.

export type ActionDefault<State> = { type: string; payload?: Partial<State> };
export type Reducer<State, Action = ActionDefault<State>> = (
  state: State,
  action: Action
) => State;
export type Dispatch<Action> = (action: Action) => void;

type Listener = () => void;

export type StoreApi<State, Action = ActionDefault<State>> = {
  getState: () => State;
  getInitialState: () => State;
  subscribe: (listener: () => void) => () => boolean;
  dispatch: (action: Action) => void;
};

export const createStore = <State, InitState extends State, Action = ActionDefault<State>>(
  reducer: Reducer<State, Action>,
  initialArg: InitState,
  init?: (init: InitState) => State
): StoreApi<State, Action> => {
  
  let state: State;
  
  const getInitialState = () => {
    if (init !== undefined) return init(initialArg);
    return initialArg;
  }

  state = getInitialState();

  const listeners: Set<() => void> = new Set();

  const getState = (): State => {
    return state;
  };

  const subscribe = (listener: Listener) => {
    listeners.add(listener);
    return () => listeners.delete(listener);
  };

  const dispatch = (action: Action) => {
    const nextState = reducer(state, action);
    
    if (!Object.is(nextState, state)) {
      state = nextState;
      listeners.forEach((listener) => listener());
    }
  };

  return { getState, getInitialState, subscribe, dispatch };
};

export const Counter = () => {
  useHookContext(Counter);

  const [count, dispatch] = useReducer((state, action) => state + action, 0);

  return (
    <button onClick={() => dispatch(1)}>
      Count: {count}
    </button>
  );
}

После этого я столкнулся с проблемой: нужно хранить состояние между вызовами useReducer, иначе при вызове dispatch хук инициализируется заново и значение возвращается к начальному ибо создается новый стор. Сначала я решил эту проблему при помощи useMemo, так же можно ее решить при помощи useRef.

import { useMemo, useSyncExternalStore } from "react";
import {
  createStore,
  ActionDefault,
  Dispatch,
  Reducer,
} from "../lib/create-store";

export const useReducer = <State = unknown, Action = ActionDefault<State>>(
  reducer: Reducer<State, Action>,
  initialValue: State
): [State, Dispatch<Action>] => {
  const store = useMemo(() => createStore(reducer, initialValue), []);

  const state = useSyncExternalStore(
    store.subscribe,
    store.getState,
    store.getInitialState
  ) as State;

  return [state, store.dispatch];
};

Но я не захотел останавливаться и решил избавиться от useMemo и сохранять все хуки в своей глобальной переменной реализовать что-то вроде механизма fiber. Для начала я сделал хук, который задавал fiber контекст для того что бы к нодам прикреплять хуки, которые в них содержатся.

import { useSyncExternalStore } from "react";
import {
  createStore,
  ActionDefault,
  Dispatch,
  Reducer,
} from "../store/create-store";

// eslint-disable-next-line @typescript-eslint/no-explicit-any
const hookInstances = new WeakMap<object, any[]>();

let currentHookIndex = 0;
let currentFiber: object | null = null;

export const useHookContext = (fiber: object) => {
  currentFiber = fiber;
  currentHookIndex = 0;
};

export const useReducer = <State, InitState extends State, Action = ActionDefault<State>>(
  reducer: Reducer<State, Action>,
  initialState: InitState, 
  init?: (init: InitState) => State
): [State, Dispatch<Action>] => {
  if (!currentFiber) {
    throw new Error("useReducer must tracked by useHookContext hook");
  }

  let hooks = hookInstances.get(currentFiber);
  if (!hooks) {
    hooks = [];
    hookInstances.set(currentFiber, hooks);
  }

  if (!hooks[currentHookIndex]) {
    hooks[currentHookIndex] = createStore(reducer, initialState, init);
  }

  const store = hooks[currentHookIndex];
  currentHookIndex++;


  const state = useSyncExternalStore(
    store.subscribe,
    store.getState,
    store.getInitialState
  ) as State;

  return [state, store.dispatch];
};

Но данная реализация требует дополнительное действие в виде использования хука useHookContext.

export const Counter = () => {
  useHookContext(Counter);

  const [count, dispatch] = useReducer((state, action) => state + action, 0);

  return (
    <button onClick={() => dispatch(1)}>
      Count: {count}
    </button>
  );
}

Эту проблему можно решить если переопределить метод React.createElement, но по мне это уже перебор и я не стал уже заморачиваться по этому поводу.

State manager

Пока я пытался сделать useReducer я написал метод createStore, который представлен выше, он предоставляет api для того что бы можно было сделать его реактивным при помощи хука useSyncExternalStore. По сути функция которую я сделал далее - просто этот же самый createStore, но возвращается результат useSyncExternalStore и dispatch.

import { useSyncExternalStore } from "react";
import { createStore, ActionDefault, Dispatch, Reducer } from "./create-store";

export const createReactStore = <
  State extends object,
  Action = ActionDefault<State>
>(
  reducer: Reducer<State, Action>,
  initialValue: State
) => {
  const vanillaStore = createStore(reducer, initialValue);

  const useStore = <Selected = State>(
    selector: (state: State) => Selected = (state) =>
      state as unknown as Selected
  ): Selected & { dispatch: Dispatch<Action> } => {
    return {
      ...useSyncExternalStore(
        vanillaStore.subscribe,
        () => selector(vanillaStore.getState()) as Selected
      ),
      dispatch: vanillaStore.dispatch,
    };
  };

  return useStore;
};

Я тут сделал что-то похожее на zustand ибо мой стор возвращает хук, из которого уже можно получить стейт и диспатч, этот стор будет работать только для объектов, что видно из реализации. Единственное что я добавил возможность передать селектор в хук что бы было проще работать с вложенными структурами. Пример использования:

import { createReactStore } from "./store/create-react-store";

const reducer = (state: State, action: { type: string }) => {
  return {
    increment: { ...state, count: state.count + 1 },
    decrement: { ...state, count: state.count - 1 },
    reset: { count: 0 }
  }[action.type] ?? state
};

const useCounter = createReactStore(reducer, { count: 0 });

export const StoreComponent = () => {
  const { count } = useCounter();
  return <h1>Store in StoreComponent: {count}</h1>;
};

export const App = () => {
  const { count, dispatch } = useCounter();

  return (
    <div>
      <StoreComponent />
      <h1>Store counter in App: {count}</h1>
      <div className="buttons">
        <button onClick={() => dispatch({ type: "increment" })}>increment</button>
        <button onClick={() => dispatch({ type: "decrement" })}>decrement</button>
        <button onClick={() => dispatch({ type: "reset" })}>reset</button>
      </div>
    </div>
  );
}

Итог

Я провел довольно интересный для меня ресерч который по факту никакого практического результата не дал, кроме издевательства над самим собой, но мне после этого почему-то захотелось написать эту статью и поделиться своими результатами.

Мне интересно как можно улучшить перфоманс в этой небольшой задаче. Так же я уверен, что в моей работе могут быть ошибки, так как я считаю не достаточно опытен и в моих знаниях есть пробелы даже на базовом уровне. Поэтому буду рад получать какие-то правки или дополнительную информацию по теме.