fix: reduce useState and useEffect (#2223)

packages/hooks/src/__tests__/index.spec.ts

@@ -3,6 +3,7 @@
  */
 import 'fake-indexeddb/auto';
 
+import { UNTITLED_WORKSPACE_NAME } from '@affine/env';
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import type { Page } from '@blocksuite/store';
 import { assertExists } from '@blocksuite/store';
@@ -10,7 +11,7 @@ import { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
 import { renderHook } from '@testing-library/react';
 import { useBlockSuiteWorkspacePageIsPublic } from '@toeverything/hooks/use-block-suite-workspace-page-is-public';
 import { useBlockSuiteWorkspacePageTitle } from '@toeverything/hooks/use-block-suite-workspace-page-title';
-import { describe, expect, test } from 'vitest';
+import { describe, expect, test, vitest } from 'vitest';
 import { beforeEach } from 'vitest';
 
 import { useBlockSuiteWorkspaceName } from '../use-block-suite-workspace-name';
@@ -30,9 +31,9 @@ beforeEach(async () => {
     const frameId = page.addBlock('affine:frame', {}, pageBlockId);
     page.addBlock('affine:paragraph', {}, frameId);
   };
-  initPage(blockSuiteWorkspace.createPage('page0'));
-  initPage(blockSuiteWorkspace.createPage('page1'));
-  initPage(blockSuiteWorkspace.createPage('page2'));
+  initPage(blockSuiteWorkspace.createPage({ id: 'page0' }));
+  initPage(blockSuiteWorkspace.createPage({ id: 'page1' }));
+  initPage(blockSuiteWorkspace.createPage({ id: 'page2' }));
 });
 
 describe('useBlockSuiteWorkspaceName', () => {
@@ -48,6 +49,17 @@ describe('useBlockSuiteWorkspaceName', () => {
     workspaceNameHook.result.current[1]('test 3');
     expect(blockSuiteWorkspace.meta.name).toBe('test 3');
   });
+
+  test('null', () => {
+    const workspaceNameHook = renderHook(() =>
+      useBlockSuiteWorkspaceName(null)
+    );
+    vitest.spyOn(globalThis.console, 'warn');
+    expect(workspaceNameHook.result.current[0]).toBe(UNTITLED_WORKSPACE_NAME);
+    workspaceNameHook.result.current[1]('test');
+    expect(globalThis.console.warn).toHaveBeenCalledTimes(2);
+    expect(workspaceNameHook.result.current[0]).toBe(UNTITLED_WORKSPACE_NAME);
+  });
 });
 
 describe('useBlockSuiteWorkspacePageTitle', () => {

packages/hooks/src/use-block-suite-workspace-name.ts

@@ -1,32 +1,51 @@
 import { UNTITLED_WORKSPACE_NAME } from '@affine/env';
 import type { Workspace } from '@blocksuite/store';
 import { assertExists } from '@blocksuite/store';
-import { useCallback, useEffect, useState } from 'react';
+import type { Atom, WritableAtom } from 'jotai';
+import { atom, useAtom } from 'jotai';
+
+const weakMap = new WeakMap<
+  Workspace,
+  WritableAtom<string, [string], void> & Atom<string>
+>();
+
+const emptyWorkspaceNameAtom = atom(UNTITLED_WORKSPACE_NAME, () => {
+  console.warn('you cannot set the name of an null workspace.');
+  console.warn('this is a bug in the code.');
+});
 
 export function useBlockSuiteWorkspaceName(
   blockSuiteWorkspace: Workspace | null
 ) {
-  const [name, set] = useState(
-    () => blockSuiteWorkspace?.meta.name ?? UNTITLED_WORKSPACE_NAME
-  );
-  useEffect(() => {
-    if (blockSuiteWorkspace) {
-      set(blockSuiteWorkspace.meta.name ?? '');
+  let nameAtom:
+    | (WritableAtom<string, [string], void> & Atom<string>)
+    | undefined;
+  if (!blockSuiteWorkspace) {
+    nameAtom = emptyWorkspaceNameAtom;
+  } else if (!weakMap.has(blockSuiteWorkspace)) {
+    const baseAtom = atom<string>(
+      blockSuiteWorkspace.meta.name ?? UNTITLED_WORKSPACE_NAME
+    );
+    const writableAtom = atom(
+      get => get(baseAtom),
+      (get, set, name: string) => {
+        blockSuiteWorkspace.meta.setName(name);
+        set(baseAtom, name);
+      }
+    );
+    baseAtom.onMount = set => {
       const dispose = blockSuiteWorkspace.meta.commonFieldsUpdated.on(() => {
         set(blockSuiteWorkspace.meta.name ?? '');
       });
       return () => {
         dispose.dispose();
       };
-    }
-  }, [blockSuiteWorkspace]);
-  const setName = useCallback(
-    (name: string) => {
-      assertExists(blockSuiteWorkspace);
-      blockSuiteWorkspace.meta.setName(name);
-      set(name);
-    },
-    [blockSuiteWorkspace]
-  );
-  return [name, setName] as const;
+    };
+    weakMap.set(blockSuiteWorkspace, writableAtom);
+    nameAtom = writableAtom;
+  } else {
+    nameAtom = weakMap.get(blockSuiteWorkspace);
+    assertExists(nameAtom);
+  }
+  return useAtom(nameAtom);
 }

packages/hooks/src/use-block-suite-workspace-page-is-public.ts

@@ -1,24 +1,36 @@
 import type { Page } from '@blocksuite/store';
-import { useCallback, useEffect, useState } from 'react';
+import { assertExists } from '@blocksuite/store';
+import type { Atom, WritableAtom } from 'jotai';
+import { atom, useAtom } from 'jotai';
+
+const weakMap = new WeakMap<
+  Page,
+  WritableAtom<boolean, [boolean], void> & Atom<boolean>
+>();
 
 export function useBlockSuiteWorkspacePageIsPublic(page: Page) {
-  const [isPublic, set] = useState<boolean>(() => page.meta.isPublic ?? false);
-  useEffect(() => {
-    const disposable = page.workspace.meta.pageMetasUpdated.on(() => {
-      set(page.meta.isPublic ?? false);
-    });
-    return () => {
-      disposable.dispose();
-    };
-  }, [page]);
-  const setIsPublic = useCallback(
-    (isPublic: boolean) => {
-      set(isPublic);
-      page.workspace.setPageMeta(page.id, {
-        isPublic,
+  if (!weakMap.has(page)) {
+    const baseAtom = atom<boolean>(page.meta.isPublic ?? false);
+    const writableAtom = atom(
+      get => get(baseAtom),
+      (get, set, isPublic: boolean) => {
+        page.workspace.setPageMeta(page.id, {
+          isPublic,
+        });
+        set(baseAtom, isPublic);
+      }
+    );
+    baseAtom.onMount = set => {
+      const disposable = page.workspace.meta.pageMetasUpdated.on(() => {
+        set(page.meta.isPublic ?? false);
       });
-    },
-    [page.id, page.workspace]
-  );
-  return [isPublic, setIsPublic] as const;
+      return () => {
+        disposable.dispose();
+      };
+    };
+    weakMap.set(page, writableAtom);
+  }
+  const isPublicAtom = weakMap.get(page);
+  assertExists(isPublicAtom);
+  return useAtom(isPublicAtom);
 }

packages/hooks/src/use-block-suite-workspace-page-title.ts

@@ -1,25 +1,40 @@
 import type { Workspace } from '@blocksuite/store';
 import { assertExists } from '@blocksuite/store';
-import { useEffect, useState } from 'react';
+import type { Atom } from 'jotai';
+import { atom, useAtomValue } from 'jotai';
+
+const weakMap = new WeakMap<Workspace, Map<string, Atom<string>>>();
+
+function getAtom(w: Workspace, pageId: string): Atom<string> {
+  if (!weakMap.has(w)) {
+    weakMap.set(w, new Map());
+  }
+  const map = weakMap.get(w);
+  assertExists(map);
+  if (!map.has(pageId)) {
+    const baseAtom = atom<string>(w.getPage(pageId)?.meta.title || 'Untitled');
+    baseAtom.onMount = set => {
+      const disposable = w.meta.pageMetasUpdated.on(() => {
+        const page = w.getPage(pageId);
+        assertExists(page);
+        set(page?.meta.title || 'Untitled');
+      });
+      return () => {
+        disposable.dispose();
+      };
+    };
+    map.set(pageId, baseAtom);
+    return baseAtom;
+  } else {
+    return map.get(pageId) as Atom<string>;
+  }
+}
 
 export function useBlockSuiteWorkspacePageTitle(
   blockSuiteWorkspace: Workspace,
   pageId: string
 ) {
-  const page = blockSuiteWorkspace.getPage(pageId);
-  const [title, setTitle] = useState(() => page?.meta.title || 'AFFiNE');
-  useEffect(() => {
-    const page = blockSuiteWorkspace.getPage(pageId);
-    setTitle(page?.meta.title || 'Untitled');
-    const dispose = blockSuiteWorkspace.meta.pageMetasUpdated.on(() => {
-      const page = blockSuiteWorkspace.getPage(pageId);
-      assertExists(page);
-      setTitle(page?.meta.title || 'Untitled');
-    });
-    return () => {
-      dispose.dispose();
-    };
-  }, [blockSuiteWorkspace, pageId]);
-
-  return title;
+  const titleAtom = getAtom(blockSuiteWorkspace, pageId);
+  assertExists(titleAtom);
+  return useAtomValue(titleAtom);
 }

packages/hooks/src/use-block-suite-workspace-page.ts

@@ -1,30 +1,55 @@
 import type { Page, Workspace } from '@blocksuite/store';
-import { useEffect, useState } from 'react';
+import { assertExists, DisposableGroup } from '@blocksuite/store';
+import type { Atom } from 'jotai';
+import { atom, useAtomValue } from 'jotai';
+
+const weakMap = new WeakMap<Workspace, Map<string, Atom<Page | null>>>();
+
+const emptyAtom = atom<Page | null>(null);
+
+function getAtom(w: Workspace, pageId: string | null): Atom<Page | null> {
+  if (!pageId) {
+    return emptyAtom;
+  }
+  if (!weakMap.has(w)) {
+    weakMap.set(w, new Map());
+  }
+  const map = weakMap.get(w);
+  assertExists(map);
+  if (!map.has(pageId)) {
+    const baseAtom = atom(w.getPage(pageId));
+    baseAtom.onMount = set => {
+      const group = new DisposableGroup();
+      group.add(
+        w.slots.pageAdded.on(id => {
+          if (pageId === id) {
+            set(w.getPage(id));
+          }
+        })
+      );
+      group.add(
+        w.slots.pageRemoved.on(id => {
+          if (pageId === id) {
+            set(null);
+          }
+        })
+      );
+      return () => {
+        group.dispose();
+      };
+    };
+    map.set(pageId, baseAtom);
+    return baseAtom;
+  } else {
+    return map.get(pageId) as Atom<Page | null>;
+  }
+}
 
 export function useBlockSuiteWorkspacePage(
   blockSuiteWorkspace: Workspace,
   pageId: string | null
 ): Page | null {
-  const [page, setPage] = useState(() => {
-    if (pageId === null) {
-      return null;
-    }
-    return blockSuiteWorkspace.getPage(pageId);
-  });
-  useEffect(() => {
-    if (pageId) {
-      setPage(blockSuiteWorkspace.getPage(pageId));
-    }
-  }, [blockSuiteWorkspace, pageId]);
-  useEffect(() => {
-    const disposable = blockSuiteWorkspace.slots.pageAdded.on(id => {
-      if (pageId === id) {
-        setPage(blockSuiteWorkspace.getPage(id));
-      }
-    });
-    return () => {
-      disposable.dispose();
-    };
-  }, [blockSuiteWorkspace, pageId]);
-  return page;
+  const pageAtom = getAtom(blockSuiteWorkspace, pageId);
+  assertExists(pageAtom);
+  return useAtomValue(pageAtom);
 }