Mini Reactive

mini-reactive.ts

interface ProxyState {
  initialState: object
  getNotifyVersion: (nextVersion?: number) => number
  createSnapshot: <State extends object>(proxyState: State, version: number) => State
  addListener: (listener: (change: Change, nextVersion: number) => void) => () => void
}

const state2proxyCache = new WeakMap<object, object>()
const proxy2proxyStateCache = new WeakMap<object, ProxyState>()
const proxy2snapshotCache = new WeakMap<object, [version: number, snapshot: object]>()
const refObjectSet = new WeakSet<object>()

type Change = ({ type: 'delete' } | { type: 'set'; value: unknown }) & {
  props: (string | symbol)[]
  previous: unknown
}

const VERSION = { proxyState: 0, emptyListenerProp: 0 }

export function proxy<State extends object>(state: State): State {
  if (state2proxyCache.has(state)) {
    return state2proxyCache.get(state) as State
  }
  let [stateVersion, elpVersion] = [VERSION.proxyState, VERSION.emptyListenerProp]
  const baseState: State = Array.isArray(state) ? [] : Object.create(Object.getPrototypeOf(state))
  const listeners = new Set<(change: Change, version: number) => void>()
  const propStateMap = new Map<
    string | symbol,
    { proxyState: ProxyState; removeListener?: () => void }
  >()

  function addPropListener(prop: string | symbol, propProxyState: ProxyState) {
    if (listeners.size) {
      const removeListener = propProxyState.addListener(createPropListener(prop))
      propStateMap.set(prop, { proxyState: propProxyState, removeListener })
    } else {
      propStateMap.set(prop, { proxyState: propProxyState })
    }
  }

  function removePropListener(prop: string | symbol) {
    const propProxyState = propStateMap.get(prop)
    if (propProxyState) {
      propProxyState.removeListener?.()
      propStateMap.delete(prop)
    }
  }

  function addListener(listener: (change: Change, version: number) => void) {
    if (!listeners.size) {
      for (const [prop, { proxyState }] of propStateMap) {
        const removeListener = proxyState.addListener(createPropListener(prop))
        propStateMap.set(prop, { proxyState, removeListener })
      }
    }
    listeners.add(listener)
    return () => {
      listeners.delete(listener)
      if (!listeners.size) {
        for (const [prop, { proxyState, removeListener }] of propStateMap) {
          if (removeListener) {
            removeListener()
            propStateMap.set(prop, { proxyState })
          }
        }
      }
    }
  }

  function notifyUpdate(change: Change, nextNotifyVersion: number = ++VERSION.proxyState) {
    if (stateVersion !== nextNotifyVersion) {
      stateVersion = nextNotifyVersion
      for (const listener of listeners) listener(change, nextNotifyVersion)
    }
  }

  function createPropListener(
    prop: string | symbol
  ): (change: Change, nextNotifyVersion: number) => void {
    return (change, nextNotifyVersion) => {
      const newChange = { ...change, props: [...change.props, prop] }
      notifyUpdate(newChange, nextNotifyVersion)
    }
  }

  function ensureEmptyPropVersion(nextElpVersion: number = ++VERSION.emptyListenerProp) {
    if (elpVersion !== nextElpVersion) {
      elpVersion = nextElpVersion
      for (const [_, { proxyState }] of propStateMap) {
        const propNotifyVersion = proxyState.getNotifyVersion(nextElpVersion)
        if (propNotifyVersion >= stateVersion) {
          stateVersion = propNotifyVersion
        }
      }
    }
  }

  function getNotifyVersion(nextElpVersion: number = ++VERSION.emptyListenerProp) {
    if (!listeners.size) ensureEmptyPropVersion(nextElpVersion)
    return stateVersion
  }

  const proxyObj = new Proxy(baseState, {
    // get(target, prop, receiver) {
    //   return Reflect.get(target, prop, receiver)
    // },
    set(target, prop, value, receiver: object) {
      const hasPreValue = Reflect.has(target, prop)
      const previous = Reflect.get(target, prop, receiver)
      if (
        hasPreValue &&
        (Object.is(previous, value) ||
          (state2proxyCache.has(value) && Object.is(previous, state2proxyCache.get(value))))
      ) {
        return true
      }
      removePropListener(prop)
      let nextValue = value
      if (!proxy2proxyStateCache.has(value) && canProxy(value)) {
        nextValue = proxy(value)
      }
      const childProxyState = proxy2proxyStateCache.get(nextValue)
      if (childProxyState) {
        addPropListener(prop, childProxyState)
      }
      Reflect.set(target, prop, nextValue, receiver)
      notifyUpdate({ type: 'set', props: [prop], value, previous })
      return true
    },
    deleteProperty(target, prop) {
      const prevValue = Reflect.get(target, prop)
      removePropListener(prop)
      const isDeleted = Reflect.deleteProperty(target, prop)
      if (isDeleted) {
        notifyUpdate({ type: 'delete', props: [prop], previous: prevValue })
      }
      return isDeleted
    },
  })
  // add properties to proxyState
  for (const key of Reflect.ownKeys(state)) {
    const descriptor = Object.getOwnPropertyDescriptor(state, key) || {}
    if ('value' in descriptor) {
      proxyObj[key as keyof State] = state[key as keyof State]
      delete descriptor.value
      delete descriptor.writable
    }
    Object.defineProperty(baseState, key, descriptor)
  }
  // set proxied State to cache
  state2proxyCache.set(state, proxyObj)
  // set proxyState to cache
  proxy2proxyStateCache.set(proxyObj, {
    initialState: state,
    getNotifyVersion,
    createSnapshot,
    addListener,
  })
  return proxyObj
}

export function subscribe<State extends object>(
  proxyObj: State,
  callback: (changes: Change[]) => void,
  notifyInSync: boolean = false
): () => void {
  const { addListener } = proxy2proxyStateCache.get(proxyObj) as ProxyState
  const changes: Change[] = []
  let promise: Promise<void> | undefined
  let isListenerRemoved = true
  const listener = (change: Change) => {
    changes.push(change)
    if (notifyInSync) {
      callback(changes.splice(0))
      return
    }
    if (!promise) {
      promise = Promise.resolve().then(() => {
        promise = undefined
        if (!isListenerRemoved) callback(changes.splice(0))
      })
    }
  }
  const removeListener = addListener(listener)
  isListenerRemoved = false
  return () => {
    isListenerRemoved = true
    removeListener()
  }
}

function createSnapshot<State extends object>(proxyState: State, version: number): State {
  // judge cache
  const cache = proxy2snapshotCache.get(proxyState)
  if (cache?.[0] === version) return cache[1] as State
  // create empty snapshot
  const snapshot: State = Array.isArray(proxyState)
    ? []
    : Object.create(Object.getPrototypeOf(proxyState))
  // add properties to snapshot
  for (const key of Reflect.ownKeys(proxyState)) {
    // ignore native properties, such as Array#length,
    // it's a getter handled internally
    if (Object.getOwnPropertyDescriptor(snapshot, key)) continue
    const value = Reflect.get(proxyState, key) as any
    const { enumerable = true } = Reflect.getOwnPropertyDescriptor(proxyState, key) || {}
    const descriptor: PropertyDescriptor = { value, enumerable, configurable: true }
    // If the value is in the map (means that it's a proxied object),
    // it should been created recursively.
    if (proxy2proxyStateCache.has(value)) {
      const { getNotifyVersion } = proxy2proxyStateCache.get(value) || {}
      descriptor.value = createSnapshot(value, getNotifyVersion?.() || version)
    }
    Object.defineProperty(snapshot, key, descriptor)
  }
  // prevent snapshot extensions,
  // make sure that the snapshot is immutable to avoid unexpected behaviors
  Object.preventExtensions(snapshot)
  return snapshot
}

export function snapshot<State extends object, StateSlice = State>(
  proxyObj: State,
  selector: (snapshot: State) => StateSlice = (s) => s as unknown as StateSlice
): StateSlice {
  const { getNotifyVersion } = proxy2proxyStateCache.get(proxyObj) as ProxyState
  return selector(createSnapshot(proxyObj, getNotifyVersion()))
}

function isObject(x: unknown): x is object {
  return typeof x === 'object' && x !== null
}

function canProxy(x: unknown) {
  return (
    isObject(x) &&
    !refObjectSet.has(x) &&
    (Array.isArray(x) || !(Symbol.iterator in x)) &&
    !(x instanceof WeakMap) &&
    !(x instanceof WeakSet) &&
    !(x instanceof Error) &&
    !(x instanceof Number) &&
    !(x instanceof Date) &&
    !(x instanceof String) &&
    !(x instanceof RegExp) &&
    !(x instanceof ArrayBuffer)
  )
}

export function ref<State extends object>(state: State): State {
  refObjectSet.add(state)
  return state
}

export function create<State extends object>(
  state: State
): {
  mutate: State
  subscribe: (callback: (changes: Change[]) => void, notifyInSync?: boolean) => () => void
  restore: () => void
  snapshot: <StateSlice = State>(selector?: (snapshot: State) => StateSlice) => StateSlice
} {
  const proxiedState = proxy(state)

  const restore = () => {
    for (const key of Reflect.ownKeys(state)) {
      proxiedState[key as keyof State] = state[key as keyof State]
    }
  }

  return {
    mutate: proxiedState,
    restore,
    subscribe: (callback: (changes: Change[]) => void, notifyInSync: boolean = false) => {
      return subscribe(proxiedState, callback, notifyInSync)
    },
    snapshot: <StateSlice = State>(
      selector: (snapshot: State) => StateSlice = (s) => s as unknown as StateSlice
    ) => {
      return snapshot(proxiedState, selector)
    },
  }
}