toy-valtio.ts

import { useSyncExternalStore } from 'use-sync-external-store/shim'
import {
  affectedToPathList,
  createProxy as createProxyToCompare,
  isChanged,
} from 'proxy-compare';
import {useCallback, useDebugValue, useEffect, useMemo, useRef} from "react";

type ProxyObject = {
  listeners: Set<Function>;
  snapshot: Function;
};

const snapshotCache = new WeakMap();
let globalVersion = 1;

export function proxy(initialObject: any) {
  let version = globalVersion;
  const listeners = new Set<Function>();
  const notifyUpdate = (op: any, nextVersion = ++globalVersion) => {
    version = nextVersion;
    listeners.forEach((listener) => listener(op, nextVersion));
  };
  const handler = {
    get(target: any, prop: any, receiver: any) {
      if (prop === 'listeners') {
        return listeners;
      }
      if (prop === 'snapshot') {
        const cache = snapshotCache.get(receiver);
        if (cache?.[0] === version) {
          return cache[1];
        }
        const snapshot = { ...target };
        snapshotCache.set(receiver, [version, snapshot])
        return snapshot;
      }
      return Reflect.get(target, prop, receiver);
    },
    deleteProperty(target: any, prop: any) {
      const deleted = Reflect.deleteProperty(target, prop);
      if (deleted) {
        notifyUpdate(['delete', [prop]]);
      }
      return deleted;
    },
    set(target: any, prop: any, value: any, receiver: any) {
      version += 1;
      Reflect.set(target, prop, value, receiver);
      notifyUpdate(['set', [prop], value]);
      return true;
    },
  };
  const baseObject = Array.isArray(initialObject)
    ? []
    : Object.create(Object.getPrototypeOf(initialObject));
  const proxyObject = new Proxy(baseObject, handler);
  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];
    }
  })
  return proxyObject;
}

export function useSnapshot(proxyObject: ProxyObject) {
  const lastSnapshot = useRef();
  const lastAffected = useRef<WeakMap<any, any>>();
  const currSnapshot = useSyncExternalStore(
    useCallback((callback: Function) => {
      proxyObject.listeners.add(callback);
      return () => proxyObject.listeners.delete(callback);
    }, [proxyObject]),
    () => {
      const nextSnapshot = proxyObject.snapshot;
      if (
        lastSnapshot.current &&
        lastAffected.current &&
        !isChanged(
          lastSnapshot.current,
          nextSnapshot,
          lastAffected.current,
          new WeakMap()
        )
      ) {
        return lastSnapshot.current;
      }
      return nextSnapshot;
    },
  );
  const currAffected = new WeakMap();
  useEffect(() => {
    lastSnapshot.current = currSnapshot;
    lastAffected.current = currAffected;
  });
  const pathList = useRef<(string | number | symbol)[][]>();
  useEffect(() => {
    pathList.current = affectedToPathList(currSnapshot, currAffected);
  });
  useDebugValue(pathList.current);
  const proxyCache = useMemo(() => new WeakMap(), []);
  return createProxyToCompare(currSnapshot, currAffected, proxyCache);
}