前言

看 2022 年 react 生态，大家都用啥中提到 valtio, 考察了一番, 在最近的 side project B 站首页推荐脚本中也有使用, API 比较简单, 也比较轻量.

这里做一个简单的源码解析.

环境

valtiio

`proxy(object: T): T`

生成一个 proxy, proxy 可以简化成下面这样

function proxy(initialObject) {
  const baseObject = Array.isArray(initialObject)
    ? []
    : Object.create(Object.getPrototypeOf(initialObject))
  const proxyObject = new Proxy(baseObject, proxyHandlers)

  Reflect.ownKeys(initialObject).forEach((key) => {
    const desc = Object.getOwnPropertyDescriptor(initialObject, key) as PropertyDescriptor
    if (desc.get || desc.set) {
      Object.defineProperty(baseObject, key, desc)
    } else {
      proxyObject[key] = initialObject[key as keyof T]
    }
  })

  return proxyObject
}

即

新建一个 initialObject 的副本 baseObject, 作为 new Proxy 的第一个参数
new Proxy
将 initialObject 中的值, 通过 proxy 拷贝到 baseObject 中, Object.defineProperty(baseObject, key, desc) / proxyObjectp[key] = initialObject[key]
- proxyObject[key] = 会走到 proxyHandlers 的 set trap 中去

`proxyHandlers`

target === baseObject , new Proxy() 第一个参数
reciver === proxyObject new Proxy() 返回值

const handler = {
  get(target: T, prop: string | symbol, receiver: any) {
    if (prop === VERSION) {
      return version
    }
    if (prop === LISTENERS) {
      return listeners
    }
    if (prop === SNAPSHOT) {
      return createSnapshot(target, receiver)
    }
    if (prop === HANDLER) {
      return handler
    }
    return Reflect.get(target, prop, receiver)
  },
  deleteProperty(target: T, prop: string | symbol) {
    const prevValue = Reflect.get(target, prop)
    const childListeners = prevValue?.[LISTENERS]
    if (childListeners) {
      childListeners.delete(popPropListener(prop))
    }
    const deleted = Reflect.deleteProperty(target, prop)
    if (deleted) {
      notifyUpdate(['delete', [prop], prevValue])
    }
    return deleted
  },
  is: Object.is,
  canProxy,
  set(target: T, prop: string | symbol, value: any, receiver: any) {
    const hasPrevValue = Reflect.has(target, prop)
    const prevValue = Reflect.get(target, prop, receiver)
    if (hasPrevValue && this.is(prevValue, value)) {
      return true
    }
    const childListeners = prevValue?.[LISTENERS]
    if (childListeners) {
      childListeners.delete(popPropListener(prop))
    }
    if (isObject(value)) {
      value = getUntracked(value) || value
    }
    let nextValue: any

    // 1
    if (Object.getOwnPropertyDescriptor(target, prop)?.set) {
      nextValue = value
    }

    // 2 promise
    else if (value instanceof Promise) {
      nextValue = value
        .then((v) => {
          nextValue[PROMISE_RESULT] = v
          notifyUpdate(['resolve', [prop], v])
          return v
        })
        .catch((e) => {
          nextValue[PROMISE_ERROR] = e
          notifyUpdate(['reject', [prop], e])
        })
    }

    // 3 proxy object
    else if (value?.[LISTENERS]) {
      nextValue = value
      nextValue[LISTENERS].add(getPropListener(prop))
    }

    // 4 nested object
    else if (this.canProxy(value)) {
      nextValue = proxy(value)
      nextValue[LISTENERS].add(getPropListener(prop))
    }

    // others: like plain value / can not proxied value
    else {
      nextValue = value
    }
    Reflect.set(target, prop, nextValue, receiver)
    notifyUpdate(['set', [prop], value, prevValue])
    return true
  },
}

`set`

set trap 比较重要, 给 proxyObject[key] = val 赋值会走到这里, proxy() 中拷贝 initialObject 中的值也会走到这里

// 3 proxy object
else if (value?.[LISTENERS]) {
  nextValue = value
  nextValue[LISTENERS].add(getPropListener(prop))
}

// 4 nested object
else if (this.canProxy(value)) {
  nextValue = proxy(value)
  nextValue[LISTENERS].add(getPropListener(prop))
}

这里可以看到, 对于 nested object 对 proxy() 来说, 其实是一样的. 即

// form 1: match 3 proxy object
const state = proxy({nested: proxy({key: 'val'})})

// form 2: match 4 nested object
const state = proxy({nested: {key: 'val'}})

form 2, set trap 中, canProxy({ key: 'val' }) 为 true, 则会自动调用 nextValue = proxy(value), 余下和 form 1 一致

`listeners`

const LISTENERS = __DEV__ ? Symbol('LISTENERS') : Symbol()

type Path = (string | symbol)[]

type Op =
  | [op: 'set', path: Path, value: unknown, prevValue: unknown]
  | [op: 'delete', path: Path, prevValue: unknown]
  | [op: 'resolve', path: Path, value: unknown]
  | [op: 'reject', path: Path, error: unknown]

type Listener = (op: Op, nextVersion: number) => void

// via `get` trap
proxyObject[LISTENERS]: Set<Listener>

getPropListener(prop: string | symbol) 返回一个 Listener
- 如果 prop 对应的 child proxy object 发生任何变更操作, 调用当前 proxyObject 的 notifyUpdate
- op[1] = [prop, ...op[1]] path 中 prepend 当前 prop
- 调用时机在上面说到的 set trap
popPropListener 删除一个 prop listener
- 调用时机为 set trap 剔除掉旧值的 prop listener, 或者 deleteProperty trap

const propListeners = new Map<string | symbol, Listener>()
const getPropListener = (prop: string | symbol) => {
  let propListener = propListeners.get(prop)
  if (!propListener) {
    propListener = (op, nextVersion) => {
      const newOp: Op = [...op]
      newOp[1] = [prop, ...(newOp[1] as Path)]
      notifyUpdate(newOp, nextVersion)
    }
    propListeners.set(prop, propListener)
  }
  return propListener
}

const popPropListener = (prop: string | symbol) => {
  const propListener = propListeners.get(prop)
  propListeners.delete(prop)
  return propListener
}

`subscribe()`

export function subscribe<T extends object>(
  proxyObject: T,
  callback: (ops: Op[]) => void,
  notifyInSync?: boolean
) {
  if (__DEV__ && !(proxyObject as any)?.[LISTENERS]) {
    console.warn('Please use proxy object')
  }
  let promise: Promise<void> | undefined
  const ops: Op[] = []
  const listener: Listener = (op) => {
    ops.push(op)
    if (notifyInSync) {
      callback(ops.splice(0))
      return
    }
    if (!promise) {
      promise = Promise.resolve().then(() => {
        promise = undefined
        callback(ops.splice(0))
      })
    }
  }
  ;(proxyObject as any)[LISTENERS].add(listener)
  return () => {
    ;(proxyObject as any)[LISTENERS].delete(listener)
  }
}

即是

state[LISTENERS].add(listener)
callback 参数是 ops: Op[] 需要注意, 用于批处理 op, 文档未标注, 一般无需使用

`snapshot()`

export function snapshot<T extends object>(proxyObject: T): Snapshot<T> {
  return (proxyObject as any)[SNAPSHOT]
}

// proxyHandlers.get trap
get() {
  if (prop === SNAPSHOT) {
    return createSnapshot(target, receiver)
  }
}

 const createSnapshot = (target: T, receiver: any) => {
    const cache = snapshotCache.get(receiver)
    if (cache?.[0] === version) {
      return cache[1]
    }
    const snapshot: any = Array.isArray(target)
      ? []
      : Object.create(Object.getPrototypeOf(target))
    markToTrack(snapshot, true) // mark to track
    snapshotCache.set(receiver, [version, snapshot])
    Reflect.ownKeys(target).forEach((key) => {
      const value = Reflect.get(target, key, receiver)
      if (refSet.has(value)) {
        markToTrack(value, false) // mark not to track
        snapshot[key] = value
      } else if (value instanceof Promise) {
        if (PROMISE_RESULT in value) {
          snapshot[key] = (value as any)[PROMISE_RESULT]
        } else {
          const errorOrPromise = (value as any)[PROMISE_ERROR] || value
          Object.defineProperty(snapshot, key, {
            get() {
              if (PROMISE_RESULT in value) {
                return (value as any)[PROMISE_RESULT]
              }
              throw errorOrPromise
            },
          })
        }
      } else if (value?.[LISTENERS]) {
        snapshot[key] = value[SNAPSHOT]
      } else {
        snapshot[key] = value
      }
    })
    Object.freeze(snapshot)
    return snapshot
  }

const snapshotCache = new WeakMap<
  object,
  [version: number, snapshot: unknown]
>()

该方法用于根据 proxyObject 创建一个快照
useSyncExternalStore 在 render 里连续调用两次 getSnapshot, 以提示 getSnapshot 在没有变化时应该返回缓存的版本. snapshot()自带根据版本的缓存, 所以这个 getSnapshot 可以通过 useSyncExternalStore 的检查.

react

`react.useSyncExternalStore`

https://github.com/magicdawn/magicdawn/issues/8#issuecomment-1135633629

可以简写成以下代码

useSyncExternalStore(subscribe, getSnapshot) {
  const [snapshot, setSnapshot] = useState(getSnapshot())
  useEffect(() => {
    return subscribe(() => {
      const nextSnapshot = getSnapshot()
      if(Object.is(snapshot, nextSnapshot)) {
        setSnapshot(nextSnapshot())
      }
    })
  }, [snapshot, subscribe, getSnapshot])
  return snapshot
}

通过 subscribe 感知外部 store 的变化
通过 getSnapshot 获取外部 store 的最新快照
感知到变化之后, 通过对比之前的 snapshot 与新获得的 nextSnapshot 是否是同一个对象, 不是就通过 setState 触发 re-render
getSnapshot 会在 render 中 or subscribe 回调中 调用, 要求外部没有变化时返回同一个 snapshot

`valtio.useSnapshot`

export function useSnapshot<T extends object>(proxyObject: T, options?: Options): Snapshot<T> {
  const notifyInSync = options?.sync
  const lastSnapshot = useRef<Snapshot<T>>()
  const lastAffected = useRef<WeakMap<object, unknown>>()
  let inRender = true
  const currSnapshot = useSyncExternalStore(
    useCallback(
      (callback) => {
        const unsub = subscribe(proxyObject, callback, notifyInSync)
        callback() // Note: do we really need this?
        return unsub
      },
      [proxyObject, notifyInSync]
    ),
    () => {
      const nextSnapshot = snapshot(proxyObject)
      try {
        if (
          !inRender &&
          lastSnapshot.current &&
          lastAffected.current &&
          !isChanged(lastSnapshot.current, nextSnapshot, lastAffected.current, new WeakMap())
        ) {
          // not changed
          return lastSnapshot.current
        }
      } catch (e) {
        // ignore if a promise or something is thrown
      }
      return nextSnapshot
    },
    () => snapshot(proxyObject)
  )
  inRender = false
  const currAffected = new WeakMap()
  useEffect(() => {
    lastSnapshot.current = currSnapshot
    lastAffected.current = currAffected
  })
  const proxyCache = useMemo(() => new WeakMap(), []) // per-hook proxyCache
  return createProxyToCompare(currSnapshot, currAffected, proxyCache)
}

inRender 前面提到 getSnapshot 会在 render 函数或者 subscribe 回调中被调用, inRender 用于区分这两种情况
- inRender 直接返回 snapshot() 的结果
- 在 subscribe 回调里使用时, 根据 isChanged 结果, 如果没有改变, 就会返回 lastSnapshot 缓存的版本, 也就不会 re-render.
createProxy & isChanged 是 valtio 相比 zustand / redux 等库从机制上优越的地方
- createProxy / isChanged 需要传入一个 affected 参数, 表示在 render 函数中哪些 prop 被访问了
- 比如 const {a, b} = useSnapshot(state), 等 render 函数执行完, affected 会记录 Set(a,b) 已被访问
- 在 isChanged 里需要传入 affected, 只会判断 a,b 是否变化, 如果是另外的字段如 c 变化了, 此处还是可以使用缓存的版本.
- 而传统的 store(redux / zustand) 等需要先 select(a,b), 再使用 shallowEqual 判断, 以避免多余的 render, 而 valtio 可以直接使用解构

proxy-compare

Types

type Affected = WeakMap<object, Set<string | symbol>>

type ProxyCache<T extends object> = WeakMap<object, ProxyHandler<T>>

type Affected = WeakMap<object, Set<string | symbol>>
type ProxyCache<T extends object> = WeakMap<object, ProxyHandler<T>>
type ProxyHandler<T extends object> = {
  // is frozen
  [FROZEN_PROPERTY]: boolean

  // Proxy
  [PROXY_PROPERTY]?: T

  // ProxyCache
  [PROXY_CACHE_PROPERTY]?: ProxyCache<object>

  // affected
  [AFFECTED_PROPERTY]?: Affected

  // proxy traps
  get(target: T, key: string | symbol): unknown
  has(target: T, key: string | symbol): boolean
  getOwnPropertyDescriptor(target: T, key: string | symbol): PropertyDescriptor | undefined
  ownKeys(target: T): (string | symbol)[]
  set?(target: T, key: string | symbol, value: unknown): boolean
  deleteProperty?(target: T, key: string | symbol): boolean
}

Affected 表示一个 object有哪些 key 被访问过了
ProxyCache 用于 proxy() 的 cache, key 是 originalObject
ProxyHandler 是 new Proxy(target, handlers) 第二个参数, 再加上一些自定义属性

`proxy-coampre.createProxy`

`createProxy`

export const createProxy = <T>(
  obj: T,
  affected: WeakMap<object, unknown>,
  proxyCache?: WeakMap<object, unknown>
): T => {
  if (!isObjectToTrack(obj)) return obj
  const origObj = (obj as {[GET_ORIGINAL_SYMBOL]?: typeof obj})[GET_ORIGINAL_SYMBOL] // unwrap proxy

  // new Proxy 第一个参数, 确保是 plain object
  const target = origObj || obj
  const frozen = isFrozen(target)

  // read cache
  let proxyHandler: ProxyHandler<typeof target> | undefined =
    proxyCache && (proxyCache as ProxyCache<typeof target>).get(target)

  if (!proxyHandler || proxyHandler[FROZEN_PROPERTY] !== frozen) {
    proxyHandler = createProxyHandler<T extends object ? T : never>(target, frozen)
    proxyHandler[PROXY_PROPERTY] = new Proxy(
      frozen ? unfreeze(target) : target,
      proxyHandler
    ) as typeof target
    if (proxyCache) {
      proxyCache.set(target, proxyHandler)
    }
  }
  proxyHandler[AFFECTED_PROPERTY] = affected as Affected
  proxyHandler[PROXY_CACHE_PROPERTY] = proxyCache as ProxyCache<object> | undefined
  return proxyHandler[PROXY_PROPERTY] as typeof target
}

最后返回了 Proxy
该方法比较直观, 查 cache, 没有 fallback 到 new Proxy, 逻辑在下面的 createProxyHandler
useSnapshot 在最后使用了该方法

`createProxyHandler`

const createProxyHandler = <T extends object>(origObj: T, frozen: boolean) => {
  let trackObject = false // for trackMemo

  const recordUsage = (h: ProxyHandler<T>, key: string | symbol, skipWithOwnKeys?: boolean) => {
    if (!trackObject) {
      let used = (h[AFFECTED_PROPERTY] as Affected).get(origObj)
      if (!used) {
        used = new Set()
        ;(h[AFFECTED_PROPERTY] as Affected).set(origObj, used)
      }

      // 默认不传 skipWithOwnKeys = false, 总是 used.add
      // 传 skipWithOwnKeys = true,
      //   不包含 OWN_KEYS_SYMBOL, 则 used.add
      //   包含则跳过
      if (!skipWithOwnKeys || !used.has(OWN_KEYS_SYMBOL)) {
        used.add(key)
      }
    }
  }

  const recordObjectAsUsed = (h: ProxyHandler<T>) => {
    trackObject = true
    ;(h[AFFECTED_PROPERTY] as Affected).delete(origObj)
  }

  const handler: ProxyHandler<T> = {
    [FROZEN_PROPERTY]: frozen,
    get(target, key) {
      if (key === GET_ORIGINAL_SYMBOL) {
        return origObj
      }
      recordUsage(this, key)
      return createProxy(
        (target as any)[key],
        this[AFFECTED_PROPERTY] as Affected,
        this[PROXY_CACHE_PROPERTY]
      )
    },
    has(target, key) {
      if (key === TRACK_MEMO_SYMBOL) {
        recordObjectAsUsed(this)
        return true
      }
      // LIMITATION: We simply record the same as `get`.
      // This means { a: {} } and { a: {} } is detected as changed,
      // if `'a' in obj` is handled.
      recordUsage(this, key)
      return key in target
    },
    getOwnPropertyDescriptor(target, key) {
      // LIMITATION: We simply record the same as `get`.
      // This means { a: {} } and { a: {} } is detected as changed,
      // if `obj.getOwnPropertyDescriptor('a'))` is handled.
      recordUsage(this, key, true)
      return Object.getOwnPropertyDescriptor(target, key)
    },
    ownKeys(target) {
      recordUsage(this, OWN_KEYS_SYMBOL)
      return Reflect.ownKeys(target)
    },
  }
  if (frozen) {
    handler.set = handler.deleteProperty = () => false
  }
  return handler
}

关于 `trackMemo` / `TRACK_MEMO_SYMBOL`

在单测看到用法

it('should fail without trackMemo', () => {
  const proxyCache = new WeakMap()
  const s1 = {a: {b: 1, c: 2}}
  const a1 = new WeakMap()
  const p1 = createProxy(s1, a1, proxyCache)
  noop(p1.a.b)
  expect(isChanged(s1, {a: s1.a}, a1)).toBe(false)
  expect(isChanged(s1, {a: {b: 3, c: 2}}, a1)).toBe(true)
  expect(isChanged(s1, {a: {b: 1, c: 3}}, a1)).not.toBe(true) // false
})

it('should work with trackMemo', () => {
  const proxyCache = new WeakMap()
  const s1 = {a: {b: 1, c: 2}}
  const a1 = new WeakMap()
  const p1 = createProxy(s1, a1, proxyCache)
  noop(p1.a.b)
  trackMemo(p1.a) // 不同1, 加了 trackMemo 调用
  expect(isChanged(s1, {a: s1.a}, a1)).toBe(false)
  expect(isChanged(s1, {a: {b: 3, c: 2}}, a1)).toBe(true)
  expect(isChanged(s1, {a: {b: 1, c: 3}}, a1)).toBe(true) // 不同2
})

注意看最后一行, 使用了 trackMemo(p1.a) 后, 就算 s1.a.b 和 s2.a.b 相同, 结果 isChanged 还是 true, 也就是改变了
trackMemo()调用链
- trackMemo(obj)
- 实现中有 return TRACK_MEMO_SYMBOL in obj
- 走到 Proxy has trap, 即上面 createProxyHandler 第 40 行
- recordObjectAsUsed(this)
- 会调用 affected.delete(originalObject), 也就是没有 used 了
文档未标注, 不太确定用途, 表现就是不会存 affected, 没有 used set, isChanged 返回总是返回 true

`get` trap

可以看到这里有递归调用, 对于嵌套 object = { nested: { key: 'val' } }

访问 proxy.nested, 会返回 createProxy(proxy.nested)

const o = { nested: { key: 'val' } }
const affected = new WeakMap()
const cache = new WeakMap()

// createProxy
const p = createProxy(o, affected, cache)

// access
console.log(p.nested.key)

// Set(1) { 'nested' }
// Set(1) { 'key' }
console.log(affected.get(o))
console.log(affected.get(o.nested))

// affected 会变成
// 伪代码
WeakMap{
  [o]: Set('nested'),
  [o.nested]: Set('key')
}

`proxy-compare.isChanged`

API

export const isChanged = (
  origObj: unknown,
  nextObj: unknown,
  affected: WeakMap<object, unknown>,
  cache?: WeakMap<object, unknown>
): boolean => {
  if (Object.is(origObj, nextObj)) {
    return false
  }
  if (!isObject(origObj) || !isObject(nextObj)) return true
  const used = (affected as Affected).get(origObj)
  if (!used) return true
  if (cache) {
    const hit = (cache as ChangedCache).get(origObj)
    if (hit && hit[NEXT_OBJECT_PROPERTY] === nextObj) {
      return hit[CHANGED_PROPERTY]
    }
    // for object with cycles
    ;(cache as ChangedCache).set(origObj, {
      [NEXT_OBJECT_PROPERTY]: nextObj,
      [CHANGED_PROPERTY]: false,
    })
  }
  let changed: boolean | null = null
  // eslint-disable-next-line no-restricted-syntax
  for (const key of used) {
    const c =
      key === OWN_KEYS_SYMBOL
        ? isOwnKeysChanged(origObj, nextObj)
        : isChanged((origObj as any)[key], (nextObj as any)[key], affected, cache)
    if (c === true || c === false) changed = c
    if (changed) break
  }
  if (changed === null) changed = true
  if (cache) {
    cache.set(origObj, {
      [NEXT_OBJECT_PROPERTY]: nextObj,
      [CHANGED_PROPERTY]: changed,
    })
  }
  return changed
}

调用流程

抛开 cache 不看, 就是
根据 origObj 从 affected 拿到 used set, 没有表示 changed.
- 这里需要注意 useSnapshot 的用法, 如果没有使用 useSnapshot()中的值, 只是用这个 hook, 那么 affected 中会是空, 最后每次 isChanged 都是 true, 每次都 re-render
遍历 used keys
- 如果是 OWN_KEYS_SYMBOL, 则比较两个 object 所拿到的 ownKeys 是否 deep equal
- 不是, 则继续递归 isChanged(origObj[key], nextObj[key], affected, cache)

Tips

这里 cache 可以给个默认值, cache = new WeakMap(), 方便外部使用

END

差不多了.

Magicdawn

valtio / proxy-compare 源码解析

前言

环境

valtiio

`proxy(object: T): T`

`proxyHandlers`

`set`

`listeners`

`snapshot()`

react

`react.useSyncExternalStore`

`valtio.useSnapshot`

proxy-compare

Types

`proxy-coampre.createProxy`

`createProxy`

`createProxyHandler`

关于 `trackMemo` / `TRACK_MEMO_SYMBOL`

`get` trap

`proxy-compare.isChanged`

调用流程

Tips

END

前言

环境

valtiio

proxy(object: T): T

proxyHandlers

set

listeners

subscribe()

snapshot()

react

react.useSyncExternalStore

valtio.useSnapshot

proxy-compare

Types

proxy-coampre.createProxy

createProxy

createProxyHandler

关于 trackMemo / TRACK_MEMO_SYMBOL

get trap

proxy-compare.isChanged

调用流程

Tips

END

`proxy(object: T): T`

`proxyHandlers`

`set`

`listeners`

`subscribe()`

`snapshot()`

`react.useSyncExternalStore`

`valtio.useSnapshot`

`proxy-coampre.createProxy`

`createProxy`

`createProxyHandler`

关于 `trackMemo` / `TRACK_MEMO_SYMBOL`

`get` trap

`proxy-compare.isChanged`