React Context 完美替代品 Jotai

1. 前言

React 的属性透传场景虽然有很多方式可以实现，但能做到代码写的少、re-render 轻松处理的方式并不多。

而状态管理工具 Jotai 却可以很好的解决这些问题。

最近的业务和组件场景里也在用此方式实现。

2. React Context 的不足

常规解决数据透传通常使用 React 的 Context 来实现，但它却有一些不足的地方：

React Context.Provider 嵌套地狱问题。
Context value 修改时，触发消费 Context 组件 re-render 的问题（组件引用值未变化时）。

2.1 Provider 嵌套地狱

如下代码所示：

<Context1.Provider value={value1}>
  <Context2.Provider value={value2}>
    <Context3.Provider value={value3}>
      <Context4.Provider value={value4}>
        <Context5.Provider value={value5}>
          <Context6.Provider value={value6}>
            {children}
          </Context6.Provider>
        </Context5.Provider>
      </Context4.Provider>
    </Context3.Provider>
  </Context2.Provider>
</Context1.Provider>

2.2 Context re-render

此问题通常有三种方式来解决：

将包含多个属性的 Context 剥离成单个 Context 进行使用。但这样会导致 2.1 中嵌套地狱的问题发生。

function Button() {
  let theme = useContext(ThemeContext);  
  return <ExpensiveTree className={theme} />;
}

将组件剥离成两个，剥离出来的用 React.memo 进行包裹，再消费。

function Button() {
  let appContextValue = useContext(AppContext);
  let theme = appContextValue.theme; 
  return <ThemedButton theme={theme} />
}

const ThemedButton = memo(({ theme }) => { 
  return <ExpensiveTree className={theme} />;
});

在组件的 return 中，用 React.useMemo 包裹，将 Context 中消费的值，做为其依赖项。

function Button() {
  let appContextValue = useContext(AppContext);
  let theme = appContextValue.theme;  

  return useMemo(() => { 
    return <ExpensiveTree className={theme} />;
  }, [theme])
}

但这些方法，都会在使用时增加一些额外的工作量，变得繁琐和易出错。

3. 状态管理库 Jotai 的优势

刚好 Jotai 的其中一部分优势就能解决 React Context 中的这些问题。

Jotai 可以完美的解决嵌套地狱的问题，及精准 re-render 的问题。

简单写了一个例子，来看一下 Jotai re-render：[codesandbox] jotai 配合 React.memo re-render 测试。

例子中 Jotai 的 Atom 及组件消费关系如下：

4. Jotai 核心

Jotai 的学习成本很低，源码也比较易读。核心就三个 API：

atom：创建原子相关方法。
useAtom：视图关联原子相关方法。
Provider：隔离和存储原子实例相关组件。

只需要熟悉这三个 API 的使用就可以轻松上手。

4.1 atom

atom 创建的其实是一份原子描述，并不会进行相关的存储。

可以创建原始原子及派生原子，用法如下：

// primitive atom
function atom<Value>(initialValue: Value): PrimitiveAtom<Value>

// read-only atom
function atom<Value>(read: (get: Getter) => Value | Promise<Value>): Atom<Value>

// writable derived atom
function atom<Value, Update>(
  read: (get: Getter) => Value | Promise<Value>,
  write: (get: Getter, set: Setter, update: Update) => void | Promise<void>
): WritableAtom<Value, Update>

// write-only derived atom
function atom<Value, Update>(
  read: Value,
  write: (get: Getter, set: Setter, update: Update) => void | Promise<void>
): WritableAtom<Value, Update>

  
const primitiveAtom = atom(initialValue)
const derivedAtomWithRead = atom(read)
const derivedAtomWithReadWrite = atom(read, write)
const derivedAtomWithWriteOnly = atom(null, write)

4.2 useAtom

用来将 atom 关联到视图中，形成响应关系。（atom 必须保持引用不变，否则会无限循环）

基本用法：

// primitive or writable derived atom
function useAtom<Value, Update>(
  atom: WritableAtom<Value, Update>,
  scope?: Scope
): [Value, SetAtom<Update>]

// read-only atom
function useAtom<Value>(atom: Atom<Value>, scope?: Scope): [Value, never]

如果“只读”或“只写”，建议直接使用：

useAtomValue
useSetAtom

4.3 Provider

使用上和 React Provider 很相似，当然底层也是用 React Context.Provider 包装的。

主要用来隔离 atom 实例及存储 atom 实例用的。

有隔离需求可以包裹在视图中，无隔离需求可以不添加，和 React Context.Provider 一样。

基本用法：

const Provider: React.FC<{
  initialValues?: Iterable<readonly [AnyAtom, unknown]>
  scope?: Scope
}>

5. Jotai 实现原理

下面看一下内部实现的基本原理。

5.1 Provider

Provider 的创建源码：

const ScopeContainerContext = getScopeContext(scope)

return createElement(
  ScopeContainerContext.Provider,
  {
    // 创建 atoms 的 weakMap 等存储，挂载在 Provider 中
    value: scopeContainerRef.current,
  },
  children
)

const ScopeContextMap = new Map<Scope | undefined, ScopeContext>()

export const getScopeContext = (scope?: Scope) => {
  if (!ScopeContextMap.has(scope)) {
    ScopeContextMap.set(scope, createContext(createScopeContainer()))
  }

  return ScopeContextMap.get(scope) as ScopeContext
}

const scopeContainerRef = useRef<ScopeContainer>()
// 从此处可以看到，Provider 仅作为隔离数据使用，后期不会发生任何 value 层面的变化，
// 所有的触发，都由 内部的 某个 atom 的 listeners 去调用渲染
if (!scopeContainerRef.current) {
  // lazy initialization
  const scopeContainer = createScopeContainer(
    // initialValues 仅作用一次，后期的任何变化，将不会有任何效果
    initialValues,
    unstable_createStore
  )
  if (unstable_enableVersionedWrite) {
    let retrying = 0
    scopeContainer.w = (write) => {
      setVersion((parentVersion) => {
        const nextVersion = retrying ? parentVersion : { p: parentVersion }
        write(nextVersion)
        return nextVersion
      })
    }
    scopeContainer.v = version
    scopeContainer.r = (fn) => {
      ++retrying
      fn()
      --retrying
    }
  }
  scopeContainerRef.current = scopeContainer
}

可以清楚的看到，主要就做了：

Provider 隔离。
在 Provider 的 value 中挂载 atom 实例。
通过 scope 缓存 createContext 实例。

另外，创建完 Provider 后，未来的所有变化都不会导致其内部的 value 发生变化，所以由 Provider 导致的渲染是不存在的。且 initialValues 也只会消费一次，未来变化也不会响应。

通过 scope 的设置，可以实现和原生 React Context.Provider 一样的效果，进行分割及多层嵌套。

5.2 atom

atom 的生成很简单，就是一份描述：

export function atom<Value, Args extends unknown[], Result>(
  read: Value | Read<Value, SetAtom<Args, Result>>,
  write?: Write<Args, Result>
) {
  const key = `atom${++keyCount}`
  const config = {
    toString: () => key,
  } as WritableAtom<Value, Args, Result> & { init?: Value }
  if (typeof read === 'function') {
    config.read = read as Read<Value, SetAtom<Args, Result>>
  } else {
    config.init = read
    config.read = (get) => get(config)
    config.write = ((get: Getter, set: Setter, arg: SetStateAction<Value>) =>
      set(
        config as unknown as PrimitiveAtom<Value>,
        typeof arg === 'function'
          ? (arg as (prev: Value) => Value)(get(config))
          : arg
      )) as unknown as Write<Args, Result>
  }
  if (write) {
    config.write = write
  }
  return config
}

生成完成后，会挂载在 Provider 的 store 中。store 由 WeakMap 维护，所以不用担心内存泄漏的问题。

store 中的派生 atom 的触发变化和渲染，主要由 dependencies 和 listeners 实现。

5.3 useAtom

export function useAtomValue<Value>(atom: Atom<Value>, scope?: Scope) {
  const ScopeContext = getScopeContext(scope)
  const scopeContainer = useContext(ScopeContext)
  // ...
}

可以看到 useAtom 中的信息也是通过读取 Context 实现，scope 用来标识与 Provider 匹配。

const [[version, valueFromReducer, atomFromReducer], rerenderIfChanged] =
  useReducer<
    Reducer<
      readonly [VersionObject | undefined, Awaited<Value>, Atom<Value>],
      VersionObject | undefined
    >,
    VersionObject | undefined
  >(
    (prev, nextVersion) => {
      const nextValue = getAtomValue(nextVersion)
      if (Object.is(prev[1], nextValue) && prev[2] === atom) {
        return prev // bail out
      }
      return [nextVersion, nextValue, atom]
    },
    versionFromProvider,
    (initialVersion) => {
      const initialValue = getAtomValue(initialVersion)
      return [initialVersion, initialValue, atom]
    }
  )

基于useAtom 创建的值，由 useReducer 实现，可以有效优化部分 re-render 的问题。

useEffect(() => {
  const { v: versionFromProvider } = scopeContainer
  if (versionFromProvider) {
    store[COMMIT_ATOM](atom, versionFromProvider)
  }
  // Call `rerenderIfChanged` whenever this atom is invalidated. Note
  // that derived atoms may not be recomputed yet.
  const unsubscribe = store[SUBSCRIBE_ATOM](
    atom,
    rerenderIfChanged,
    versionFromProvider
  )
  rerenderIfChanged(versionFromProvider)
  return unsubscribe
}, [store, atom, scopeContainer])

const subscribeAtom = (
  atom: AnyAtom,
  callback: (version?: VersionObject) => void,
  version?: VersionObject
) => {
  const mounted = addAtom(version, atom)
  const listeners = mounted.l
  listeners.add(callback)
  return () => {
    listeners.delete(callback)
    // TODO should version be `undefined` for delAtom?
    delAtom(version, atom)
  }
}

useRender 的 dispatch 会添加到 atom 的订阅列表中，当 atom 发生变化时，会执行订阅列表中的 dispatch 实现响应。

6. 开发中的注意

如果 Jotai 使用在应用级的全局状态中，不太需要注意。但在开发的组件中使用 Jotai 还是需要注意的。

6.1 开发组件的注意

最根本的一点就是 atom 的隔离问题，不然组件内的数据用的是同一份，会导致干扰。

所以需要额外包裹 Provider 进行隔离：

const ComponentA = () => {
  return <div>A</div>;
};

export default () => {
  return (
    <Provider>
      <ComponentA />
    </Provider>
  );
};

如果原组件有 memo，需要把它移到 Provider 这一层：

export default memo(() => {
  return (
    <Provider>
      <ComponentA />
    </Provider>
  );
});

如果有 forwardRef 之类的，基本要写两遍：

const ComponentA = forwardRef(({ onChange, defaultValue }, ref) => {
  useImperativeHandle(ref, () => ({ changeStyle }));
  return <div>A</div>;
});

export default memo(
  forwardRef((props, ref) => {
    useImperativeHandle(ref, () => ({
      changeStyle: _ref.current.changeStyle,
    }));

    const _ref = useRef();

    return (
      <Provider>
        <ComponentA ref={_ref} />
      </Provider>
    );
  })
);