手搓简易版jotai（react-atom）,100行代码了解核心

jotai

个人实用感受：

优点：

• api设计比较巧妙（多组件状态交互场景-比如表单））

用法普及

const store = createStore()
const atomStateA = atom('state A')
const atomStateC = atom((get)=>{
    return `${get(atomStateA)} 派生出来的C`
})

function InputA() {
  const [a, setStateA] = useAtom(atomStateA)
  return (
    <button onClick={() => {
      setStateA('new state A')
    }}
    >
      {a}
    </button>
  )
}

function InputB() {
  const a = useAtomValue(atomStateA)
  return <>{a}</>
}
function InputC() {
  const c = useAtomValue(atomStateC)
  return <>{c}</>
}


export default () => {
  return (
    <Provider store={store}>
      <>
        <InputA />
        <InputB />
        <InputC />
      </>
    </Provider>

  )
}

基本功能实现

export function atom(initState){
    return {
        init:initState
    }
}

export function createStore()  {
  // 创建状态缓存
  const atomStateMap = new WeakMap()
  
  // 获取实例状态
  function readAtom(atomEntity) {
    if(atomStateMap.has(atomEntity)){
      return atomStateMap.get(atomEntity)
    }
    // 如果缓存没有值，则返回 初始状态
    return atomEntity.init
  }
  
  // 设置实例状态
  function setAtomState(atomEntity, state) {   
    atomStateMap.set(atomEntity, state)
  }
 
  return {
    get: readAtom, set: setAtomState,
  }
}

const store = createStore()
export function useAtom(atomEntity){
  const setFn = useCallback((newState)=>{
    store.set(atomEntity,newState)
  },[atomEntity])
  // bug:不能主动获取最新状态
  return [setFn,store.get(atomEntity)]
}

状态实时更新-升级1

export function createStore()  {
  // 状态缓存
  const atomStateMap = new WeakMap()
  //  状态设置 事件监听缓存  ！！！对比👆新增！！！
  const listenersMap = new WeakMap()
  
  function readAtom(atomEntity) {
    if(atomStateMap.has(atomEntity)){
      return atomStateMap.get(atomEntity)
    }
    return atomEntity.init
  }

  function setAtomState(atomEntity, state) {   
    atomStateMap.set(atomEntity, state)
    // 设置状态时候 主动触发atom实例身上的监听事件
    publishAtom(atomEntity)
  }
  
  // 触发atom实例身上所有监听方法    ！！！对比👆新增！！！
  function publishAtom(atomEntity) {
    const listenerSet = listenersMap.get(atomEntity)
    if (listenerSet) {
      listenerSet.forEach((listener) => {
        listener()
      })
    }
  }
 
 // 给atom实例 添加一个监听方法      ！！！对比👆新增！！！
  function subscribeAtom(atomEntity, listener) {
    if (!listenersMap.has(atomEntity)) {
      listenersMap.set(atomEntity, new Set())
    }
    (listenersMap.get(atomEntity)!).add(listener)

    return () => {
      (listenersMap.get(atomEntity)!).delete(listener)
    }
  }

 
  return {
    get: readAtom, set: setAtomState,sub:subscribeAtom
  }
}

const store = createStore()
export function useAtom(atomEntity){
  const setFn = useCallback((newState)=>{
    store.set(atomEntity,newState)
  },[atomEntity])

  //  ！！！对比👆新增！！！
  const [state, setState] = useState(() => {
    return store.get(atomEntity)
  })
  
  // 添加atom实例的 监听方法      ！！！对比👆新增！！！
  useEffect(() => {
    return store.sub(atomEntity, () => {
      setState(store.get(atomEntity))
    })
  }, [store, atomEntity])
  

  return [setFn,state]
}

支持派生状态 - 升级2

const atomStateA = atom('state A')
const atomStateC = atom((get)=>{
    return `${get(atomStateA)} 派生出来的C`
})

export function atom(read){
  const entity = {};
  // ！！！对比👆新增！！！
  if (typeof read === 'function') {
    entity.read = read 
  }else{
    entity.init = read
  }
  return entity
}

export function createStore()  {
  const atomStateMap = new WeakMap()
  const listenersMap = new WeakMap()
  // 新增反向依赖 ！！！对比👆新增！！！
  const backDependenciesMap = new WeakMap()
  
  function readAtom(atomEntity) {
    if(atomStateMap.has(atomEntity)){
      return atomStateMap.get(atomEntity)
    }
    if('init' in atomEntity){
      return atomEntity.init
    }
    // 新增反向依赖 ！！！对比👆新增！！！
    function getter(depAtomEntity){
      if (!backDependenciesMap.has(atom)) {
        backDependenciesMap.set(atom, new Set())
      }
      (backDependenciesMap.get(atom)!).add(atomEntity)
      return readAtom(depAtomEntity)
    }
    return atomEntity.read(getter)
  }

  function setAtomState(atomEntity, state) {   
    atomStateMap.set(atomEntity, state)
    publishAtom(atomEntity)
    // 反向依赖出发  ！！！对比👆新增！！！
    backDependenciesMap.get(atomEntity).forEach((backEntity) => {
      publishAtom(backEntity)
    })
  }

  function publishAtom(atomEntity) {
    const listenerSet = listenersMap.get(atomEntity)
    if (listenerSet) {
      listenerSet.forEach((listener) => {
        listener()
      })
    }
  }

  function subscribeAtom(atomEntity, listener) {
    if (!listenersMap.has(atomEntity)) {
      listenersMap.set(atomEntity, new Set())
    }
    (listenersMap.get(atomEntity)!).add(listener)

    return () => {
      (listenersMap.get(atomEntity)!).delete(listener)
    }
  }

 
  return {
    get: readAtom, set: setAtomState,sub:subscribeAtom
  }
}

const store = createStore()
export function useAtom(atomEntity){
  const setFn = useCallback((newState)=>{
    store.set(atomEntity,newState)
  },[atomEntity])


  const [state, setState] = useState(() => {
    return store.get(atomEntity)
  })

  useEffect(() => {
    return store.sub(atomEntity, () => {
      setState(store.get(atomEntity))
    })
  }, [store, atomEntity])

  return [setFn,state]
}

本文章参照jotai实现，并非照搬，简化+改写了很多代码，让其简单易懂些。更多详情，请访问github-jotai