参照源码编写简易版 Piniavue3 项目一般都会使用 pinia 作为状态管理库，本文就参照源码，编写一个简易版的

前言

vue3 项目一般都会使用 pinia 作为状态管理库，本文就参照源码，编写一个简易版的 pinia，目的在于了解 pinia 背后的实现原理，做到知其所以然。

回顾 pinia

先来回顾下 pinia 本身是如何使用的。安装后，在 src\main.ts，从 pinia 导入 createPinia 方法调用，然后将返回结果 store 传给 app.use() 进行安装：

// 代码片段 1
import { createApp } from 'vue'
import { createPinia } from 'pinia'
import App from './App.vue'

const app = createApp(App)
const store = createPinia()
app.use(store)
console.log(store)

既然是被 app.use() 调用的，说明 createPinia() 返回的 store 或是个包含 install 函数的对象，或本身就是函数。查看 pinia 源码，可知 createPinia() 返回的是类型为 Pinia 的对象：

Pinia 是个接口，定义如下：

// Every application must own its own pinia to be able to create stores
export interface Pinia {
  install: (app: App) => void

  // root state
  state: Ref<Record<string, StateTree>>
  
  /**
   * Adds a store plugin to extend every store
   * @param plugin - store plugin to add
   */
  use(plugin: PiniaPlugin): Pinia

  // Installed store plugins
  _p: PiniaPlugin[]

  // App linked to this Pinia instance
  _a: App

  // Effect scope the pinia is attached to
  _e: EffectScope

  // Registry of stores used by this pinia.
  _s: Map<string, StoreGeneric>

  // Added by `createTestingPinia()` to bypass `useStore(pinia)`.
  _testing?: boolean
}

其中确实有个 install 方法用于将 pinia 实例挂载到 vue 应用中。另外还有些属性，解释如下：

state 为响应式的根状态容器，其中存储了所有 store 的根状态，也就是我们用 defineStore 创建的一个个 store 对象中的 state 所返回的对象，以 store 的 id 为键；
use 用于安装全局的 store 插件，比如想把 store 存储到 localStorage，可以通过 use 添加状态持久化插件 pinia-plugin-persistedstate：

import { createPinia } from 'pinia'
// 引入插件
import piniaPluginPersistedstate from 'pinia-plugin-persistedstate'
const store = createPinia()
// 使用插件
store.use(piniaPluginPersistedstate)

_p 为已安装的插件数组；
_a 是关联的 vue 应用实例，可用于获取全局配置；
_e 的值 EffectScope 是使用 vue3 提供的 api effectscope 创建的响应式作用域，以管理 pinia 实例的所有响应式副作用；
_s 是个 Map 对象，里面管理着所有 store 对象。

我定义了个 store 如下：

// src\stores\counter.ts 代码片段 2
import { ref, computed } from 'vue'
import { defineStore } from 'pinia'

export const useCounterStore = defineStore('counter', () => {
  const count = ref(0)
  const doubleCount = computed(() => count.value * 2)
  function increment() {
    count.value++
  }
  return { count, doubleCount, increment }
})

在 src\App.vue 中使用：

<!-- src\App.vue 代码片段 2.1 -->
<script setup lang="ts">
  import { storeToRefs } from 'pinia'
  import { useCounterStore } from '@/stores/counter'
  
  const counterStore = useCounterStore()
  const { count, doubleCount } = storeToRefs(counterStore)
  const add = () => {
    counterStore.increment()
  }
</script>

<template>
  <div>
    <div>count: {{ count }}</div>
    <div>doubleCount: {{ doubleCount }}</div>
    <button @click="add">add</button>
  </div>
</template>

现在查看代码片段 1 中的 console.log(store) 的结果，如下：

实现 `createPinia`

创建文件 src\myPinia\createPinia.ts 用于定义我们自己的 createPinia：

// src\myPinia\createPinia.ts
import { effectScope, ref } from 'vue'
import { piniaSymbol } from './rootStore'
import type { App } from 'vue'

export function createPinia() {
  const scope = effectScope()
  const state = scope.run(() => ref({}))
  const pinia = {
    install(app: App) {
      app.provide(piniaSymbol, pinia)
    },
    _e: scope,
    _s: new Map(),
    state
   }
  return pinia
}

作为简易版本的 createPinia，其返回的对象 pinia 只包含 install、_e、_s 和 state：

在 install 方法中，通过 app 的 provide 方法向所有组件提供 pinia，这样之后在组件中使用某个 store 时，才可以将对应的 store 存放进 pinia 的 _s 中。piniaSymbol 定义如下：

// src\myPinia\rootStore.ts
export const piniaSymbol = Symbol('pinia')

_e 的值 scope 通过 effectScope() 得到，state 的值就是 scope.run() 中的回调函数所返回的值。

effectScope

此处趁便介绍 effectScope 的相关知识，它用于创建一个 effect 作用域，可以捕获到其中所创建的响应式副作用，即 watch 和 watchEffect（还包括 render，在 vue3.5 之前，还包括 computed，具体可查阅 pinia 仓库的这条 issue，里面有尤雨溪本人的解释）：

<script setup lang="ts">
  import { effectScope, ref, watch, watchEffect } from 'vue'
  const count = ref(0)
  const scope = effectScope()
  scope.run(() => {
    watch(count, (newVal: number) => {
      console.log('watch', newVal)
    })
    watchEffect(() => {
      console.log('watchEffect', count.value)
    })
  })
  const add = () => {
    count.value++
  }
  const stop = () => {
    scope.stop()
  }
</script>

<template>
  <div>
    <div>count: {{ count }}</div>
    <button @click="add">add</button>
    <button @click="stop">stop</button>
  </div>
</template>

点击 add 按钮，会触发定义于传给 scope.run() 的回调中的 watch 和 watchEffect，但是当点击 stop 按钮后，由于调用了 scope.stop()，停止了 effect 作用域，在其中创建的 watch 和 watchEffect 不再生效：

scope.run() 的返回值为传给它的回调函数的的返回值，这从其类型定义就可以看出：

export declare class EffectScope {
    // ...省略
    run<T>(fn: () => T): T | undefined;
    stop(fromParent?: boolean): void;
}

实现 defineStore

defineStore() 用于定义一个个的 store，它有 2 种语法风格 —— Option Store 和 Setup Store，代码片段 2 使用的就是 Setup 的语法。无论使用哪种语法，defineStore() 的第 1 个参数，都需要是个独一无二的 id。注意，从 v3.0 版本开始，id 不再可以像下面这样定义在选项对象内部：

// v3.0 开始废弃
export const useCounterStore = defineStore({
  id: 'counter',
  // ...
})

Option 语法的第 2 个参数为对象，而 Setup 语法的第 2 个参数为函数，我们使用函数的重载来定义：

// src\myPinia\store.ts 代码片段 3
import { inject } from "vue"
import { piniaSymbol } from "./rootStore"
import type { createPinia } from "./createPinia"

type Pinia = ReturnType<typeof createPinia>

export function defineStore(id: string, options: Record<string, any>): () => any
// setup 风格的写法其实还能接收第 3个参数作为插件选项，此处忽略
export function defineStore<SS>(id: string, storeSetup: () => SS): () => any
export function defineStore(id: string, setup: any) {
  const isSetupStore = typeof setup === 'function'
  const options = isSetupStore ? {} : setup
  function useStore() {
    const pinia = inject(piniaSymbol) as Pinia
    if (!pinia._s.has(id)) {
      if (isSetupStore) {
        createSetupStore(id, setup, pinia)
      } else {
        createOptionsStore(id, options, pinia)
      }
    }
    return pinia._s.get(id)
  }
  return useStore
}

以代码片段 2 定义的 useCounterStore 为例，在组件中使用时，是需要调用一下去获取 store 的：

<!-- src\App.vue 代码片段 4 -->
<script setup lang="ts">
  import { useCounterStore } from '@/stores/counter'
  const counterStore = useCounterStore()
</script>

由此可见，defineStore() 返回的应该是个函数，且该函数执行后返回的是 store。所以在代码片段 3 中，在 defineStore() 的实现里，最后 return 的为函数 useStore。代码片段 2 或 4 中的 useCounterStore 其实就是 useStore，它的执行是位于组件内的，所以可以使用 inject 去注入之前在 createPinia 时通过 app.provide 提供的 pinia 对象，返回结果为 pinia._s 这个 Map 对象中，键为 id（即 counter）所对应的 store。

处理 Setup Store

useCounterStore 在被首次调用前，pinia._s 中是没有键为 counter 的元素的，即 pinia._s.has(id) 为 false（这也解释了为何说 pinia 中的 store 都是懒加载的，只有第一次被使用时才会创建对应的实例）。如果传入 defineStore 的第 2 个参数为函数，说明使用的是 Setup 语法，调用 createSetupStore 处理：

// src\myPinia\store.ts
function createSetupStore<SS>(id: string, setup: () => SS, pinia: Pinia) {
  const store: Record<string, any> = reactive({})
  const setupStore = setup()
  let storeScope: EffectScope
  const result = pinia._e.run(() => {
    storeScope = effectScope()
    return storeScope.run(() => processSetup<SS>(id, setupStore, pinia))
  })
  Object.assign(store, result)
  pinia._s.set(id, store)
  store.$id = id
}

createSetupStore 的目的在于往 pinia._s 中存储键为 id，值为 store 的映射关系，store 是一个 reactive 对象，它的值由 Object.assign(store, result) 获取，result 其实就是 setup 执行的返回结果 setupStore，也就是代码片段 2 中作为传给 defineStore() 的第 2 个参数的那个函数的返回值。

至于 processSetup 的作用，则是将 store 中的 state，以 id 为键存入 pinia.state 中：

// src\myPinia\store.ts
function processSetup<SS>(id: string, setupStore: SS, pinia: Pinia) {
  if (!pinia.state?.value[id]) pinia.state.value[id] = {}
  for (const key in setupStore) {
    const prop = setupStore[key]
    if ((isRef(prop) && !isComputed(prop)) || isReactive(prop)) {
      pinia.state.value[id][key] = prop
    }
  }
  return setupStore
}

对于 Setup 语法定义的 store，ref() 或 reactive() 定义的值即为 state 属性，所以判断条件为 (isRef(prop) && !isComputed(prop)) || isReactive(prop)。因为计算属性传给 isRef() 也为 true，而计算属性为 getters，不属于 state，所以需要定义 isComputed 方法来排除，计算属性会有个 effect 属性，而普通的 ref 是没有的：

// src\myPinia\store.ts
function isComputed(value: any) {
  return !!(isRef(value) && (value as any).effect)
}

最后给 store 添加上属性 $id，值为 id，方便之后可以直接通过 store 获取对应 id。

处理 Option Store

当传入 defineStore 的第 2 个参数为对象时，说明使用的是 Option 语法，调用 createOptionsStore 处理。在 pinia 源码中，createOptionsStore 函数内部会创建一个 setup 函数，将 Option 语法中的 state、getters 和 actions 处理为 setup 语法，最后还是会调用处理 Setup Store 的 createSetupStore：

// packages/pinia/src/store.ts 省略部分代码
function createOptionsStore() {
  function setup() {
    // ...
  }
  store = createSetupStore(id, setup, pinia)
  return store as any
}

从这个角度来看，我们在定义 store 时，相较于 Option Store，使用 Setup Store 语法会更高效些。但 Setup 的写法也有弊处，比如不能直接调用 $reset()，需要自己定义。

我们自己实现的 createOptionsStore 定义如下：

// src\myPinia\store.ts
function createOptionsStore(id: string, options: any, pinia: Pinia) {
  const store: Record<string, any> = reactive({})
  let storeScope: EffectScope
  const result = pinia._e.run(() => {
    storeScope = effectScope()
    return storeScope.run(() => processOptions(id, options, pinia, store))
  })
  Object.assign(store, result)
  pinia._s.set(id, store)
  store.$id = id
}

和 createSetupStore 的定义类似，主要区别在于 processOptions 的处理：

// src\myPinia\store.ts
function processOptions(id: string, options: any, pinia: Pinia, store: Reactive<any>) {
  const { state, getters, actions } = options
  
  // 处理 state
  const storeState = toRefs(ref(state ? state() : {}).value)
  pinia.state.value[id] = storeState
  
  // 处理 getters
  const storeGetters = Object.keys(getters || {}).reduce((computedGetters, name) => {
    computedGetters[name] = computed(() => getters[name].call(store, store))
    return computedGetters
  }, {} as Record<string, ComputedRef>)
  
  // 处理 actions
  const storeActions: Record<string, () => any> = {}
  for (const name in actions) {
    const action = actions[name]
    storeActions[name] = function (...rest) {
      return action.apply(store, rest)
    }
  }
  
  return {
    ...storeState,
    ...storeGetters,
    ...storeActions
  }
}

对于 state，就是将 state() 所返回的对象中，各个属性转换为 ref，然后在 pinia.state 中存储相应数据；
对于 getters，所做的是将 getters: { doubleCount: (state) => state.count * 2 } 这样的内容转换成 const doubleCount = computed(() => count.value * 2) 这样的定义。转换时通过 call(store, store) 来绑定 this 以及传参；
对于 actions 的处理则在于绑定 this 为 store，因为参数为数组，所以使用的是 apply。

实现 storeToRefs

代码片段 2.1 中，我使用解构语法从 counterStore 中获取值时，需要先将 counterStore传给 storeToRefs()，否则获取到的 count 等会失去响应式，所以我们还需要实现 storeToRefs 以方便使用：

import { toRaw, isRef, isReactive, toRef } from "vue"

export function storeToRefs(store: any) {
  const rawStore = toRaw(store)
  const refs: Record<string, any>  = {}
  for (const key in rawStore) {
    const value = rawStore[key]
    if (isRef(value) || isReactive(value)) {
      refs[key] = toRef(store, key)
    }
  }
  return refs
}

storeToRefs() 接收 store 作为参数，store 为一个 reactive 对象，使用 toRaw() 获取原始对象 rawStore进行遍历，如果某个属性的值本身为响应式的数据，则通过 toRef() 创建对应的 ref，这样创建的 ref 与其源属性是保持同步的。

验证效果

至此，我们自己实现了 createPinia、defineStore 和 storeToRefs，为了方便大家测试，我把代码打包上传到了 npm 仓库，诸君可以 npm install tiny-pinia 或 pnpm add tiny-pinia 下载到本地项目进行验证。

也可以直接查看完整代码：github / gitee

感谢.gif 点赞.png

参照源码编写简易版 Pinia

前言