【pinia源码】二、defineStore源码解析

前言

【pinia源码】系列文章主要分析pinia的实现原理。该系列文章源码参考pinia v2.0.14。

源码地址：https://github.com/vuejs/pinia

官方文档：https://pinia.vuejs.org

本篇文章将分析defineStore的实现。

使用

通过defineStore定义一个store。

const useUserStore = defineStore('counter', {
  state: () => ({
    count: 0
  }),
  actions: {
    increment() {
      this.count++
    }
  }
})

// or
const useUserStore = defineStore({
  id: 'counter',
  state: () => ({
    count: 0
  }),
  actions: {
    increment() {
      this.count++
    }
  }
})

// or
const useUserStore = defineStore('counter', () => {
  const count = ref(0)
  
  function increment() {
    count.value++
  }
  return { count, increment }
})

defineStore

export function defineStore(
  idOrOptions: any,
  setup?: any,
  setupOptions?: any
): StoreDefinition {
  let id: string
  let options:
    | DefineStoreOptions<
        string,
        StateTree,
        _GettersTree<StateTree>,
        _ActionsTree
      >
    | DefineSetupStoreOptions<
        string,
        StateTree,
        _GettersTree<StateTree>,
        _ActionsTree
      >

  const isSetupStore = typeof setup === 'function'
  if (typeof idOrOptions === 'string') {
    id = idOrOptions
    options = isSetupStore ? setupOptions : setup
  } else {
    options = idOrOptions
    id = idOrOptions.id
  }

  function useStore(pinia?: Pinia | null, hot?: StoreGeneric): StoreGeneric { // ... }

  useStore.$id = id

  return useStore
}

defineStore函数可以接收三个参数：idOrOptions、setup、setOptions，后两个参数为可选参数。下面是三个defineStore的函数类型定义。

export function defineStore<
  Id extends string,
  S extends StateTree = {},
  G extends _GettersTree<S> = {},
  A /* extends ActionsTree */ = {}
>(
  id: Id,
  options: Omit<DefineStoreOptions<Id, S, G, A>, 'id'>
): StoreDefinition<Id, S, G, A>

export function defineStore<
  Id extends string,
  S extends StateTree = {},
  G extends _GettersTree<S> = {},
  A /* extends ActionsTree */ = {}
  >(options: DefineStoreOptions<Id, S, G, A>): StoreDefinition<Id, S, G, A>

export function defineStore<Id extends string, SS>(
  id: Id,
  storeSetup: () => SS,
  options?: DefineSetupStoreOptions<
    Id,
    _ExtractStateFromSetupStore<SS>,
    _ExtractGettersFromSetupStore<SS>,
    _ExtractActionsFromSetupStore<SS>
    >
): StoreDefinition<
  Id,
  _ExtractStateFromSetupStore<SS>,
  _ExtractGettersFromSetupStore<SS>,
  _ExtractActionsFromSetupStore<SS>
  >

首先在defineStore中声明了三个变量：id、options、isSetupStore，其中id为定义的store的唯一id，options为定义store时的options，isSetupStore代表传入的setup是不是个函数。

然后根据传入的idOrOptions的类型，为id、otions赋值。紧接着声明了一个useStore函数，并将id赋给它，然后将其return。截止到此，我们知道defineStore会返回一个函数，那么这个函数具体是做什么的呢？我们继续看useStore的实现。

useStore

function useStore(pinia?: Pinia | null, hot?: StoreGeneric): StoreGeneric {
  // 获取当前实例
  const currentInstance = getCurrentInstance()
  // 测试环境下，忽略提供的参数，因为总是能使用getActivePinia()获取pinia实例
  // 非测试环境下，如果未传入pinia，则会从组件中使用inject获取pinia
  pinia =
    (__TEST__ && activePinia && activePinia._testing ? null : pinia) ||
    (currentInstance && inject(piniaSymbol))
  // 设置激活的pinia
  if (pinia) setActivePinia(pinia)

  // 如果没有activePinia，那么可能没有install pinia，开发环境下进行提示
  if (__DEV__ && !activePinia) {
    throw new Error(
      `[🍍]: getActivePinia was called with no active Pinia. Did you forget to install pinia?\n` +
        `\tconst pinia = createPinia()\n` +
        `\tapp.use(pinia)\n` +
        `This will fail in production.`
    )
  }

  // 设置pinia为激活的pinia
  pinia = activePinia!

  // 从pina._s中查找id否注册过，如果没有被注册，创建一个store并注册在pinia._s中
  if (!pinia._s.has(id)) {
    if (isSetupStore) {
      createSetupStore(id, setup, options, pinia)
    } else {
      createOptionsStore(id, options as any, pinia)
    }

    if (__DEV__) {
      useStore._pinia = pinia
    }
  }

  // 从pinia._s中获取id对应的store
  const store: StoreGeneric = pinia._s.get(id)!

  if (__DEV__ && hot) {
    const hotId = '__hot:' + id
    const newStore = isSetupStore
      ? createSetupStore(hotId, setup, options, pinia, true)
      : createOptionsStore(hotId, assign({}, options) as any, pinia, true)

    hot._hotUpdate(newStore)

    // cleanup the state properties and the store from the cache
    delete pinia.state.value[hotId]
    pinia._s.delete(hotId)
  }

  if (
    __DEV__ &&
    IS_CLIENT &&
    currentInstance &&
    currentInstance.proxy &&
    !hot
  ) {
    const vm = currentInstance.proxy
    const cache = '_pStores' in vm ? vm._pStores! : (vm._pStores = {})
    cache[id] = store
  }

  // 返回store
  return store as any
}

useStore接收两个可选参数：pinia、hot。pinia是个Pinia的实例，而hot只在开发环境下有用，它与模块的热更新有关。

在useStore中会首先获取当前组件实例，如果存在组件实例，使用inject(piniaSymbol)获取pinia（在install中会进行provide），并将其设置为activePinia，然后在activePinia._s中查找是否有被注册为id的store，如果没有则创建store，将其注册到activePinia._s中。最后返回activePinia._s中id对应的store。

现在我们知道useStore函数，最终会返回一个store。那么这个store是什么呢？它是如何创建的呢？在useStore中根据不同情况中有两中方式来创建store，分别是：createSetupStore、createOptionsStore。这两个方式的使用条件是：如果defineStore第二个参数是个function调用createSetupStore，相反调用createOptionsStore。

createSetupStore

createSetupStore函数代码过长，这里就不贴完整代码了。createSetupStore可接收参数如下：

createSetupStore代码有500多行，如果从头开始看的话，不容易理解。我们可以根据createSetupStore的用途，从其核心开始看。因为createSetupStore是需要创建store，并将store注册到pinia._s中，所以createSetupStore中可能需要创建store，我们找到创建store的地方。

const partialStore = {
  _p: pinia,
  // _s: scope,
  $id,
  $onAction: addSubscription.bind(null, actionSubscriptions),
  $patch,
  $reset,
  $subscribe(callback, options = {}) {
    const removeSubscription = addSubscription(
      subscriptions,
      callback,
      options.detached,
      () => stopWatcher()
    )
    const stopWatcher = scope.run(() =>
      watch(
        () => pinia.state.value[$id] as UnwrapRef<S>,
        (state) => {
          if (options.flush === 'sync' ? isSyncListening : isListening) {
            callback(
              {
                storeId: $id,
                type: MutationType.direct,
                events: debuggerEvents as DebuggerEvent,
              },
              state
            )
          }
        },
        assign({}, $subscribeOptions, options)
      )
    )!

    return removeSubscription
  },
  $dispose,
} as _StoreWithState<Id, S, G, A>

if (isVue2) {
  partialStore._r = false
}

const store: Store<Id, S, G, A> = reactive(
  assign(
    __DEV__ && IS_CLIENT
      ? // devtools custom properties
        {
          _customProperties: markRaw(new Set<string>()),
          _hmrPayload,
        }
      : {},
    partialStore
  )
) as unknown as Store<Id, S, G, A>

pinia._s.set($id, store)

store是用reactive包装的一个响应式对象，reactive所包装的对象是由partialStore通过Object.assign进行复制的。partialStore中定义了很多方法，这些方法都是暴露给用户操作store的一些接口，如$onAction可设置actions的回调、$patch可更新store中的state、$dispose可销毁store。

在调用完pinia._s.set($id, store)之后，会执行setup，获取所有的数据。setup的执行会在创建pinia实例时创建的effectScope中运行，而且会再单独创建一个effectScope，用来单独执行setup.

const setupStore = pinia._e.run(() => {
  scope = effectScope()
  return scope.run(() => setup())
})!

然后遍历setupStore的属性：如果prop（key对应的值）为ref（不为computed）或reactive，则将key及prop同步到pina.state.value[$id]中；如果prop为function，则会使用wrapAction包装prop，并将包装后的方法赋值给setupStore[key]，以覆盖之前的值，同时将包装后的方法存入optionsForPlugin.actions中。

for (const key in setupStore) {
  const prop = setupStore[key]

  // 如果prop是ref（但不是computed）或reactive
  if ((isRef(prop) && !isComputed(prop)) || isReactive(prop)) {
    if (__DEV__ && hot) {
      set(hotState.value, key, toRef(setupStore as any, key))
    } else if (!isOptionsStore) {
      if (initialState && shouldHydrate(prop)) {
        if (isRef(prop)) {
          prop.value = initialState[key]
        } else {
          mergeReactiveObjects(prop, initialState[key])
        }
      }

      // 将对应属性同步至pinia.state中
      if (isVue2) {
        set(pinia.state.value[$id], key, prop)
      } else {
        pinia.state.value[$id][key] = prop
      }
    }

    if (__DEV__) {
      _hmrPayload.state.push(key)
    }
  } else if (typeof prop === 'function') { // 如果prop是function
    // 使用wrapAction包装prop，在wrapAction会处理afeterCallback、errorCallback
    const actionValue = __DEV__ && hot ? prop : wrapAction(key, prop)

    // 将actionsValue添加到setupStore中，覆盖原来的function
    if (isVue2) {
      set(setupStore, key, actionValue)
    } else {
      setupStore[key] = actionValue
    }

    if (__DEV__) {
      _hmrPayload.actions[key] = prop
    }

    // 将function类型的prop存入optionsForPlugin.actions中
    optionsForPlugin.actions[key] = prop
  } else if (__DEV__) {
    if (isComputed(prop)) {
      _hmrPayload.getters[key] = isOptionsStore
        ? // @ts-expect-error
        options.getters[key]
        : prop
      if (IS_CLIENT) {
        const getters: string[] =
          setupStore._getters || (setupStore._getters = markRaw([]))
        getters.push(key)
      }
    }
  }
}

接下来我们看下wrapAction是如何进行包装function类型上的prop。

function wrapAction(name: string, action: _Method) {
  return function (this: any) {
    setActivePinia(pinia)
    const args = Array.from(arguments)

    const afterCallbackList: Array<(resolvedReturn: any) => any> = []
    const onErrorCallbackList: Array<(error: unknown) => unknown> = []
    function after(callback: _ArrayType<typeof afterCallbackList>) {
      afterCallbackList.push(callback)
    }
    function onError(callback: _ArrayType<typeof onErrorCallbackList>) {
      onErrorCallbackList.push(callback)
    }

    triggerSubscriptions(actionSubscriptions, {
      args,
      name,
      store,
      after,
      onError,
    })

    let ret: any
    try {
      ret = action.apply(this && this.$id === $id ? this : store, args)
    } catch (error) {
      triggerSubscriptions(onErrorCallbackList, error)
      throw error
    }

    // 如果结果是promise，在promise中触发afterCallbackList及onErrorCallbackList
    if (ret instanceof Promise) {
      return ret
        .then((value) => {
          triggerSubscriptions(afterCallbackList, value)
          return value
        })
        .catch((error) => {
          triggerSubscriptions(onErrorCallbackList, error)
          return Promise.reject(error)
        })
    }

    triggerSubscriptions(afterCallbackList, ret)
    return ret
  }
}

wrapAction首先返回一个函数，在这个函数中，首先将pinia设置为activePinia，触发actionSubscriptions中的函数，然后执行action函数，如果执行过程中出错，会执行onErrorCallbackList中的errorCallback，如果没有出错的话，执行afterCallbackList中的afterCallback，最后将action的返回结果return。

wrapAction中的actionSubscriptions是个什么呢？

其实actionSubscriptions中的callback就是是通过store.$onAction添加的回调函数；在执行actionSubscriptions中的callback过程中，会将对应callback添加到afterCallbackList或onErrorCallbackList中。例如：

store.$onAction(({ after, onError, name, store }) => {
  after((value) => {
    console.log(value)
  })
  
  onError((error) => {
    console.log(error)
  })
})

遍历完setupStore之后，会将setupStore合并至store和store的原始对对象中，以方便使用storeToRefs()检索响应式对象。

if (isVue2) {
  Object.keys(setupStore).forEach((key) => {
    set(
      store,
      key,
      setupStore[key]
    )
  })
} else {
  assign(store, setupStore)
  assign(toRaw(store), setupStore)
}

紧接着拦截store.$state的get、set方法：当调用store.$state时，能够从pinia.state.value找到对应的state；当使用store.$state = xxx去修改值时，则调用$patch方法修改值。

Object.defineProperty(store, '$state', {
  get: () => (__DEV__ && hot ? hotState.value : pinia.state.value[$id]),
  set: (state) => {
    /* istanbul ignore if */
    if (__DEV__ && hot) {
      throw new Error('cannot set hotState')
    }
    $patch(($state) => {
      assign($state, state)
    })
  },
})

截止到此，store就准备完毕。如果在Vue2环境下，会将store._r设置为true。

if (isVue2) {
  store._r = true
}

接下来就需要调用使用use方法注册的plugins：

pinia._p.forEach((extender) => {
  if (__DEV__ && IS_CLIENT) {
    const extensions = scope.run(() =>
      extender({
        store,
        app: pinia._a,
        pinia,
        options: optionsForPlugin,
      })
    )!
    Object.keys(extensions || {}).forEach((key) =>
      store._customProperties.add(key)
    )
    assign(store, extensions)
  } else {
    // 将plugin的结果合并到store中
    assign(
      store,
      scope.run(() =>
        extender({
          store,
          app: pinia._a,
          pinia,
          options: optionsForPlugin,
        })
      )!
    )
  }
})

最后返回store。

if (
  initialState &&
  isOptionsStore &&
  (options as DefineStoreOptions<Id, S, G, A>).hydrate
) {
  ;(options as DefineStoreOptions<Id, S, G, A>).hydrate!(
    store.$state,
    initialState
  )
}

isListening = true
isSyncListening = true
return store

接下来看下store中的几个方法：

$onAction

在每个action中添加回调函数。回调接收一个对象参数：该对象包含name（action的key值）、store（当前store）、after（添加action执行完之后的回调）、onError（添加action执行过程中的错误回调）、args（action的参数）属性。

示例：

// 统计add action的调用次数
let count = 0, successCount = 0, failCount = 0
store.$onAction(({ name, after, onError }) => {
  if (name === 'add') {
    count++
    after((resolveValue) => {
      successCount++
      console.log(resolveValue)
    })
  
    onError((error) => {
      failCount++
      console.log(error)
    })
  }
})

$onAction内部通过发布订阅模式实现。在pinia中有个专门的订阅模块subscriptions.ts，其中包含两个主要方法：addSubscription（添加订阅）、triggerSubscriptions（触发订阅）。

addSubscription可接收四个参数：subscriptions（订阅列表）、callback（添加的订阅函数）、detached（游离的订阅，如果为false在组件卸载后，自动移除订阅；如果为true，不会自动移除订阅）、onCleanup（订阅被移除时的回调）

triggerSubscriptions接收两个参数：subscriptions（订阅列表）、args（action的参数列表）

export function addSubscription<T extends _Method>(
  subscriptions: T[],
  callback: T,
  detached?: boolean,
  onCleanup: () => void = noop
) {
  subscriptions.push(callback)

  const removeSubscription = () => {
    const idx = subscriptions.indexOf(callback)
    if (idx > -1) {
      subscriptions.splice(idx, 1)
      onCleanup()
    }
  }

  if (!detached && getCurrentInstance()) {
    onUnmounted(removeSubscription)
  }

  return removeSubscription
}

export function triggerSubscriptions<T extends _Method>(
  subscriptions: T[],
  ...args: Parameters<T>
) {
  subscriptions.slice().forEach((callback) => {
    callback(...args)
  })
}

$onAction通过addSubscription.bind(null, actionSubscriptions)实现。

如何触发订阅？

首先在store的初始化过程中，会将action使用wrapAction函数进行包装，wrapAction返回一个函数，在这个函数中会先触发actionSubscriptions，这个触发过程中会将afterCallback、onErrorCallback添加到对应列表。然后调用action，如果调用过程中出错，则触发onErrorCallbackList，否则触发afterCallbackList。如果action的结果是Promise的话，则在then中触发onErrorCallbackList，在catch中触发onErrorCallbackList。然后会将包装后的action覆盖原始action，这样每次调用action时就是调用的包装后的action。

$patch

使用$patch可以更新state的值，可进行批量更新。$patch接收一个partialStateOrMutator参数，它可以是个对象也可以是个方法。

示例：

store.$patch((state) => {
  state.name = 'xxx'
  state.age = 14
})
// or
store.$patch({
  name: 'xxx',
  age: 14
})

$patch源码：

function $patch(
  partialStateOrMutator:
    | _DeepPartial<UnwrapRef<S>>
    | ((state: UnwrapRef<S>) => void)
): void {
  // 合并的相关信息
  let subscriptionMutation: SubscriptionCallbackMutation<S>
  // 是否触发状态修改后的回调，isListening代表异步触发，isSyncListening代表同步触发
  // 此处先关闭回调的触发，防止修改state的过程中频繁触发回调
  isListening = isSyncListening = false
  if (__DEV__) {
    debuggerEvents = []
  }
  // 如果partialStateOrMutator是个function，执行方法，传入当前的store
  if (typeof partialStateOrMutator === 'function') {
    partialStateOrMutator(pinia.state.value[$id] as UnwrapRef<S>)
    subscriptionMutation = {
      type: MutationType.patchFunction,
      storeId: $id,
      events: debuggerEvents as DebuggerEvent[],
    }
  } else { // 如果不是function，则调用mergeReactiveObjects合并state
    mergeReactiveObjects(pinia.state.value[$id], partialStateOrMutator)
    subscriptionMutation = {
      type: MutationType.patchObject,
      payload: partialStateOrMutator,
      storeId: $id,
      events: debuggerEvents as DebuggerEvent[],
    }
  }
  // 当合并完之后，将isListening、isSyncListening设置为true，意味着可以触发状态改变后的回调函数了
  const myListenerId = (activeListener = Symbol())
  nextTick().then(() => {
    if (activeListener === myListenerId) {
      isListening = true
    }
  })
  isSyncListening = true
  // 因为在修改pinia.state.value[$id]的过程中关闭（isSyncListening与isListening）了监听，所以需要手动触发订阅列表
  triggerSubscriptions(
    subscriptions,
    subscriptionMutation,
    pinia.state.value[$id] as UnwrapRef<S>
  )
}

$reset

通过构建一个新的state object将state重置为初始状态。只在options配置下生效。如果是setup配置，开发环境下报错。

store.$reset = function $reset() {
  // 重新执行state，获取一个新的state
  const newState = state ? state() : {}
  // 通过$patch，使用assign将newState合并到$state中
  this.$patch(($state) => {
    assign($state, newState)
  })
}

$subscribe

设置state改变后的回调，返回一个移除回调的函数。可接受两个参数：callback（添加的回调函数）、options:{detached, flush, ...watchOptions}（detached同addSubscription中的detached；flush代表是否同步触发回调，可取值：sync）。

示例：

store.$subribe((mutation: {storeId, type, events}, state) => {
  console.log(storeId)
  console.log(type)
  console.log(state)
}, { detached: true, flush: 'sync' })

$subscribe源码：

function $subscribe(callback, options = {}) {
  // 将callback添加到subscriptions中，以便使用$patch更新状态时，触发回调
  // 当使用removeSubscription移除callback时，停止对pinia.state.value[$id]监听
  const removeSubscription = addSubscription(
    subscriptions,
    callback,
    options.detached,
    () => stopWatcher()
  )
  const stopWatcher = scope.run(() =>
    // 监听pinia.state.value[$id]，以触发callback，当使用$patch更新state时，不会进入触发这里的callback
    watch(
      () => pinia.state.value[$id] as UnwrapRef<S>,
      (state) => {
        if (options.flush === 'sync' ? isSyncListening : isListening) {
          callback(
            {
              storeId: $id,
              type: MutationType.direct,
              events: debuggerEvents as DebuggerEvent,
            },
            state
          )
        }
      },
      assign({}, $subscribeOptions, options)
    )
  )!

  return removeSubscription
}

在callback中的第一个参数中有个type属性，表示是通过什么方式更新的state，它有三个值：

1. MutationType.direct：通过state.name='xxx'/store.$state.name='xxx'等方式修改
2. MutationType.patchObject：通过store.$patch({ name: 'xxx' })方式修改
3. MutationType.patchFunction：通过store.$patch((state) => state.name='xxx')方式修改

$dispose

销毁store。

function $dispose() {
  // 停止监听
  scope.stop()
  // 清空subscriptions及actionSubscriptions
  subscriptions = []
  actionSubscriptions = []
  // 从pinia._s中删除store
  pinia._s.delete($id)
}

createOptionsStore

createOptionsStore可接收参数如下：

function createOptionsStore<
  Id extends string,
  S extends StateTree,
  G extends _GettersTree<S>,
  A extends _ActionsTree
>(
  id: Id,
  options: DefineStoreOptions<Id, S, G, A>,
  pinia: Pinia,
  hot?: boolean
): Store<Id, S, G, A> {
  const { state, actions, getters } = options

  const initialState: StateTree | undefined = pinia.state.value[id]

  let store: Store<Id, S, G, A>

  function setup() {
    // 如果pinia.state.value[id]不存在，进行初始化
    if (!initialState && (!__DEV__ || !hot)) {
      if (isVue2) {
        set(pinia.state.value, id, state ? state() : {})
      } else {
        pinia.state.value[id] = state ? state() : {}
      }
    }

    // 将pinia.state.value[id]各属性值转为响应式对象
    const localState =
      __DEV__ && hot
        ? // use ref() to unwrap refs inside state TODO: check if this is still necessary
          toRefs(ref(state ? state() : {}).value)
        : toRefs(pinia.state.value[id])

    // 处理getters，并将处理后的getters和actions合并到localState中
    return assign(
      localState,
      actions,
      Object.keys(getters || {}).reduce((computedGetters, name) => {
        computedGetters[name] = markRaw(
          computed(() => {
            setActivePinia(pinia)
            const store = pinia._s.get(id)!
            
            if (isVue2 && !store._r) return

            return getters![name].call(store, store)
          })
        )
        return computedGetters
      }, {} as Record<string, ComputedRef>)
    )
  }

  // 利用createSetupStore创建store
  store = createSetupStore(id, setup, options, pinia, hot, true)

  // 重写store.$reset
  store.$reset = function $reset() {
    const newState = state ? state() : {}
    this.$patch(($state) => {
      assign($state, newState)
    })
  }

  return store as any
}

在createOptionsStore中会根据传入参数构造一个setup函数，然后通过createSetupStore创建一个store，并重写store.$reset方法，最后返回store。

这个setup函数中会将state()的返回值赋值给pinia.state.value[id]，然后将pinia.state.value[id]进行toRefs，得到localState，最后将处理后的getters和actions都合并到localState中，将其返回。对于getters的处理：将每个getter函数都转成一个计算属性。

总结

defineStore返回一个useStore函数，通过执行useStore可以获取对应的store。调用useStore时我们并没有传入id，为什么能准确获取store呢？这是因为useStore是个闭包，在执行useStore执行过程中会自动获取id。

获取store的过程：

1. 首先获取组件实例
2. 使用inject(piniaSymbol)获取pinia实例
3. 判断pinia._s中是否有对应id的键，如果有直接取对应的值作为store，如果没有则创建store

store创建流程分两种：setup方式与options方式

setup方式：

1. 首先在pinia.state.value中添加键为$id的空对象，以便后续赋值
2. 使用reactive声明一个响应式对象store
3. 将store存至pinia._s中
4. 执行setup获取返回值setupStore
5. 遍历setupStore的键值，如果值是ref（不是computed）或reactive，将键值添加到pinia.state.value[$id]中；如果值时function，首先将值使用wrapAction包装，然后用包装后的function替换setupStore中对应的值
6. 将setupStore合并到store中
7. 拦截store.$state，使get操作可以正确获取pinia.state.value[$id]，set操作使用this.$patch更新
8. 调用pinia._p中的扩展函数，扩展store

options方式：

1. 从options中提取state、getter、actions
2. 构建setup函数，在setup函数中会将getter处理成计算属性
3. 使用setup方式创建store
4. 重写store.$reset