refactor(editor): improve the implementation of flat proxy (#10872)

blocksuite/framework/store/src/reactive/flat-native-y/proxy.ts

@@ -1,18 +1,34 @@
 import { signal } from '@preact/signals-core';
 
-import { Boxed } from '../boxed';
 import { isPureObject } from '../is-pure-object';
 import { native2Y } from '../native-y';
-import { Text } from '../text';
 import type { UnRecord } from '../types';
+import { signalUpdater } from './signal-updater';
 import type { CreateProxyOptions } from './types';
 import {
+  bindOnChangeIfNeed,
   getFirstKey,
   isProxy,
-  keyWithoutPrefix,
   keyWithPrefix,
   markProxy,
 } from './utils';
+import { yMapUpdater } from './y-map-updater';
+
+export function createProxy(options: CreateProxyOptions): UnRecord {
+  const { base } = options;
+
+  if (isProxy(base)) {
+    return base;
+  }
+
+  initializeProxy(options);
+
+  const proxyHandler = createProxyHandler(options);
+  const proxy = new Proxy(base, proxyHandler);
+  markProxy(proxy);
+
+  return proxy;
+}
 
 function initializeProxy(options: CreateProxyOptions) {
   const { basePath, yMap, base, root } = options;
@@ -31,30 +47,71 @@ function initializeProxy(options: CreateProxyOptions) {
   });
 }
 
-export function createProxy(options: CreateProxyOptions): UnRecord {
+function updateSignal(
+  value: unknown,
+  prop: string,
+  receiver: any,
+  options: CreateProxyOptions
+) {
+  const {
+    root,
+    shouldByPassSignal,
+    byPassSignalUpdate,
+    onChange,
+    basePath,
+    initialized,
+    onDispose,
+  } = options;
+
+  const fullPath = basePath ? `${basePath}.${prop}` : prop;
+  const firstKey = getFirstKey(fullPath);
+  signalUpdater({
+    root,
+    firstKey,
+    shouldByPassSignal,
+    byPassSignalUpdate,
+    onChange,
+    basePath,
+    value,
+    handleNestedUpdate: (signalKey: string) => {
+      if (value === undefined) {
+        delete root[signalKey];
+        return;
+      }
+
+      const signalData = signal(value);
+      root[signalKey] = signalData;
+      const unsubscribe = signalData.subscribe(next => {
+        if (!initialized()) {
+          return;
+        }
+        byPassSignalUpdate(() => {
+          receiver[prop] = next;
+          onChange?.(firstKey, true);
+        });
+      });
+      const subscription = onDispose.subscribe(() => {
+        subscription.unsubscribe();
+        unsubscribe();
+      });
+    },
+  });
+}
+
+function createProxyHandler(
+  options: CreateProxyOptions
+): ProxyHandler<UnRecord> {
   const {
     yMap,
-    base,
-    root,
-    onDispose,
-    shouldByPassSignal,
     shouldByPassYjs,
-    byPassSignalUpdate,
     basePath,
     onChange,
     initialized,
     transform,
     stashed,
   } = options;
-  const isRoot = !basePath;
 
-  if (isProxy(base)) {
-    return base;
-  }
-
-  initializeProxy(options);
-
-  const proxy = new Proxy(base, {
+  return {
     has: (target, p) => {
       return Reflect.has(target, p);
     },
@@ -63,90 +120,20 @@ export function createProxy(options: CreateProxyOptions): UnRecord {
     },
     set: (target, p, value, receiver) => {
       if (typeof p === 'string') {
-        const list: Array<() => void> = [];
         const fullPath = basePath ? `${basePath}.${p}` : p;
         const firstKey = getFirstKey(fullPath);
         const isStashed = stashed.has(firstKey);
 
-        const updateSignal = (value: unknown) => {
-          if (shouldByPassSignal()) {
-            return;
-          }
-
-          const signalKey = `${firstKey}$`;
-          if (!(signalKey in root)) {
-            if (!isRoot) {
-              return;
-            }
-            const signalData = signal(value);
-            root[signalKey] = signalData;
-            const unsubscribe = signalData.subscribe(next => {
-              if (!initialized()) {
-                return;
-              }
-              byPassSignalUpdate(() => {
-                proxy[p] = next;
-                onChange?.(firstKey, true);
-              });
-            });
-            const subscription = onDispose.subscribe(() => {
-              subscription.unsubscribe();
-              unsubscribe();
-            });
-            return;
-          }
-          byPassSignalUpdate(() => {
-            const prev = root[firstKey];
-            const next = isRoot
-              ? value
-              : isPureObject(prev)
-                ? { ...prev }
-                : Array.isArray(prev)
-                  ? [...prev]
-                  : prev;
-            // @ts-expect-error allow magic props
-            root[signalKey].value = next;
-            onChange?.(firstKey, true);
-          });
-        };
-
         if (isPureObject(value)) {
-          const syncYMap = () => {
-            if (shouldByPassYjs()) {
-              return;
-            }
-            yMap.forEach((_, key) => {
-              if (initialized() && keyWithoutPrefix(key).startsWith(fullPath)) {
-                yMap.delete(key);
-              }
+          const syncYMap = () =>
+            yMapUpdater({
+              shouldByPassYjs,
+              yMap,
+              initialized,
+              onChange,
+              fullPath,
+              value,
             });
-            const run = (obj: object, basePath: string) => {
-              Object.entries(obj).forEach(([key, value]) => {
-                const fullPath = basePath ? `${basePath}.${key}` : key;
-                if (isPureObject(value)) {
-                  run(value, fullPath);
-                } else {
-                  list.push(() => {
-                    if (value instanceof Text || Boxed.is(value)) {
-                      value.bind(() => {
-                        onChange?.(firstKey, true);
-                      });
-                    }
-                    yMap.set(keyWithPrefix(fullPath), native2Y(value));
-                  });
-                }
-              });
-            };
-            run(value, fullPath);
-            if (list.length && initialized()) {
-              yMap.doc?.transact(
-                () => {
-                  list.forEach(fn => fn());
-                },
-                { proxy: true }
-              );
-            }
-          };
 
           if (!isStashed) {
             syncYMap();
@@ -155,23 +142,20 @@ export function createProxy(options: CreateProxyOptions): UnRecord {
           const next = createProxy({
             ...options,
             basePath: fullPath,
-            yMap,
             base: value as UnRecord,
-            root,
           });
 
           const result = Reflect.set(target, p, next, receiver);
-          updateSignal(next);
+          updateSignal(next, p, receiver, options);
           return result;
         }
 
-        if (value instanceof Text || Boxed.is(value)) {
-          value.bind(() => {
-            onChange?.(firstKey, true);
-          });
-        }
+        bindOnChangeIfNeed(value, () => {
+          onChange?.(firstKey, true);
+        });
         const yValue = native2Y(value);
         const next = transform(firstKey, value, yValue);
+
         if (!isStashed && initialized() && !shouldByPassYjs()) {
           yMap.doc?.transact(
             () => {
@@ -182,7 +166,7 @@ export function createProxy(options: CreateProxyOptions): UnRecord {
         }
 
         const result = Reflect.set(target, p, next, receiver);
-        updateSignal(next);
+        updateSignal(next, p, receiver, options);
         return result;
       }
       return Reflect.set(target, p, value, receiver);
@@ -193,34 +177,6 @@ export function createProxy(options: CreateProxyOptions): UnRecord {
         const firstKey = getFirstKey(fullPath);
         const isStashed = stashed.has(firstKey);
 
-        const updateSignal = () => {
-          if (shouldByPassSignal()) {
-            return;
-          }
-
-          const signalKey = `${firstKey}$`;
-          if (!(signalKey in root)) {
-            if (!isRoot) {
-              return;
-            }
-            delete root[signalKey];
-            return;
-          }
-          byPassSignalUpdate(() => {
-            const prev = root[firstKey];
-            const next = isRoot
-              ? prev
-              : isPureObject(prev)
-                ? { ...prev }
-                : Array.isArray(prev)
-                  ? [...prev]
-                  : prev;
-            // @ts-expect-error allow magic props
-            root[signalKey].value = next;
-            onChange?.(firstKey, true);
-          });
-        };
-
         if (!isStashed && initialized() && !shouldByPassYjs()) {
           yMap.doc?.transact(
             () => {
@@ -236,14 +192,10 @@ export function createProxy(options: CreateProxyOptions): UnRecord {
         }
 
         const result = Reflect.deleteProperty(target, p);
-        updateSignal();
+        updateSignal(undefined, p, undefined, options);
         return result;
       }
       return Reflect.deleteProperty(target, p);
     },
-  });
-
-  markProxy(proxy);
-
-  return proxy;
+  };
 }

blocksuite/framework/store/src/reactive/flat-native-y/signal-updater.ts

@@ -0,0 +1,49 @@
+import { isPureObject } from '../is-pure-object';
+import type { CreateProxyOptions } from './types';
+
+type UpdateSignalOptions = Pick<
+  CreateProxyOptions,
+  'shouldByPassSignal' | 'root' | 'onChange' | 'byPassSignalUpdate' | 'basePath'
+> & {
+  firstKey: string;
+  value: unknown;
+  handleNestedUpdate: (signalKey: string) => void;
+};
+
+export function signalUpdater({
+  root,
+  firstKey,
+  shouldByPassSignal,
+  byPassSignalUpdate,
+  onChange,
+  basePath,
+  value,
+  handleNestedUpdate,
+}: UpdateSignalOptions): void {
+  const isRoot = !basePath;
+  if (shouldByPassSignal()) {
+    return;
+  }
+
+  const signalKey = `${firstKey}$`;
+  if (!(signalKey in root)) {
+    if (!isRoot) {
+      return;
+    }
+    handleNestedUpdate(signalKey);
+    return;
+  }
+  byPassSignalUpdate(() => {
+    const prev = root[firstKey];
+    const next = isRoot
+      ? value
+      : isPureObject(prev)
+        ? { ...prev }
+        : Array.isArray(prev)
+          ? [...prev]
+          : prev;
+    // @ts-expect-error allow magic props
+    root[signalKey].value = next;
+    onChange?.(firstKey, true);
+  });
+}

blocksuite/framework/store/src/reactive/flat-native-y/utils.ts

@@ -1,4 +1,6 @@
 import { SYS_KEYS } from '../../consts';
+import { Boxed } from '../boxed';
+import { Text } from '../text';
 import type { UnRecord } from '../types';
 
 export const keyWithoutPrefix = (key: string) => key.replace(/(prop|sys):/, '');
@@ -33,10 +35,16 @@ export function deleteEmptyObject(
   }
 }
 
-export function getFirstKey(key: string) {
+export function getFirstKey(key: string): string {
   const result = key.split('.').at(0);
   if (!result) {
     throw new Error(`Invalid key for: ${key}`);
   }
   return result;
 }
+
+export function bindOnChangeIfNeed(value: unknown, onChange: () => void): void {
+  if (value instanceof Text || Boxed.is(value)) {
+    value.bind(onChange);
+  }
+}

blocksuite/framework/store/src/reactive/flat-native-y/y-map-updater.ts

@@ -0,0 +1,61 @@
+import { isPureObject } from '../is-pure-object';
+import { native2Y } from '../native-y';
+import type { CreateProxyOptions } from './types';
+import {
+  bindOnChangeIfNeed,
+  getFirstKey,
+  keyWithoutPrefix,
+  keyWithPrefix,
+} from './utils';
+
+type YMapOptions = Pick<
+  CreateProxyOptions,
+  'shouldByPassYjs' | 'yMap' | 'initialized' | 'onChange'
+> & {
+  fullPath: string;
+  value: unknown;
+};
+
+export function yMapUpdater({
+  shouldByPassYjs,
+  yMap,
+  initialized,
+  onChange,
+  fullPath,
+  value,
+}: YMapOptions) {
+  const firstKey = getFirstKey(fullPath);
+  if (shouldByPassYjs()) {
+    return;
+  }
+  const list: Array<() => void> = [];
+  yMap.forEach((_, key) => {
+    if (initialized() && keyWithoutPrefix(key).startsWith(fullPath)) {
+      yMap.delete(key);
+    }
+  });
+  const run = (obj: object, basePath: string) => {
+    Object.entries(obj).forEach(([key, value]) => {
+      const fullPath = basePath ? `${basePath}.${key}` : key;
+      if (isPureObject(value)) {
+        run(value, fullPath);
+      } else {
+        list.push(() => {
+          bindOnChangeIfNeed(value, () => {
+            onChange?.(firstKey, true);
+          });
+          yMap.set(keyWithPrefix(fullPath), native2Y(value));
+        });
+      }
+    });
+  };
+  run(value as object, fullPath);
+  if (list.length && initialized()) {
+    yMap.doc?.transact(
+      () => {
+        list.forEach(fn => fn());
+      },
+      { proxy: true }
+    );
+  }
+}