refactor(hooks): reduce null types (#3111)

(cherry picked from commit 8b5d997322)

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

@@ -3,14 +3,13 @@
  */
 import 'fake-indexeddb/auto';
 
-import { UNTITLED_WORKSPACE_NAME } from '@affine/env/constant';
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import { assertExists } from '@blocksuite/global/utils';
 import type { Page } from '@blocksuite/store';
 import { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
 import { renderHook } from '@testing-library/react';
 import { useAtomValue } from 'jotai';
-import { describe, expect, test, vitest } from 'vitest';
+import { describe, expect, test } from 'vitest';
 import { beforeEach } from 'vitest';
 
 import { useBlockSuitePagePreview } from '../use-block-suite-page-preview';
@@ -52,17 +51,6 @@ 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-page-references.ts

@@ -1,9 +1,10 @@
+import { assertExists } from '@blocksuite/global/utils';
 import type { Page, Workspace } from '@blocksuite/store';
-import { atom, useAtomValue } from 'jotai';
-import { atomFamily } from 'jotai/utils';
+import { type Atom, atom, useAtomValue } from 'jotai';
 
 import { useBlockSuiteWorkspacePage } from './use-block-suite-workspace-page';
 
+const weakMap = new WeakMap<Page, Atom<string[]>>();
 function getPageReferences(page: Page): string[] {
   // todo: is there a way to use page indexer to get all references?
   return page
@@ -13,26 +14,27 @@ function getPageReferences(page: Page): string[] {
     .filter(Boolean);
 }
 
-const pageReferencesAtomFamily = atomFamily((page: Page | null) => {
-  if (page === null) {
-    return atom([]);
-  }
-  const baseAtom = atom<string[]>(getPageReferences(page));
-  baseAtom.onMount = set => {
-    const dispose = page.slots.yUpdated.on(() => {
-      set(getPageReferences(page));
-    });
-    return () => {
-      dispose.dispose();
+const getPageReferencesAtom = (page: Page) => {
+  if (!weakMap.has(page)) {
+    const baseAtom = atom<string[]>(getPageReferences(page));
+    baseAtom.onMount = set => {
+      const dispose = page.slots.yUpdated.on(() => {
+        set(getPageReferences(page));
+      });
+      return () => {
+        dispose.dispose();
+      };
     };
-  };
-  return baseAtom;
-});
+    weakMap.set(page, baseAtom);
+  }
+  return weakMap.get(page) as Atom<string[]>;
+};
 
 export function useBlockSuitePageReferences(
   blockSuiteWorkspace: Workspace,
   pageId: string
 ): string[] {
   const page = useBlockSuiteWorkspacePage(blockSuiteWorkspace, pageId);
-  return useAtomValue(pageReferencesAtomFamily(page));
+  assertExists(page);
+  return useAtomValue(getPageReferencesAtom(page));
 }

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

@@ -1,4 +1,3 @@
-import { assertExists } from '@blocksuite/global/utils';
 import type { Page, Workspace } from '@blocksuite/store';
 import { useMemo } from 'react';
 
@@ -6,7 +5,6 @@ export function useBlockSuiteWorkspaceHelper(blockSuiteWorkspace: Workspace) {
   return useMemo(
     () => ({
       createPage: (pageId?: string): Page => {
-        assertExists(blockSuiteWorkspace);
         return blockSuiteWorkspace.createPage({ id: pageId });
       },
     }),

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

@@ -1,28 +1,15 @@
 import { UNTITLED_WORKSPACE_NAME } from '@affine/env/constant';
-import { assertExists } from '@blocksuite/global/utils';
 import type { Workspace } from '@blocksuite/store';
 import type { Atom, WritableAtom } from 'jotai';
 import { atom, useAtom } from 'jotai';
 
-const weakMap = new WeakMap<
-  Workspace,
-  WritableAtom<string, [string], void> & Atom<string>
->();
+type StringAtom = 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.');
-});
+const weakMap = new WeakMap<Workspace, StringAtom>();
 
-export function useBlockSuiteWorkspaceName(
-  blockSuiteWorkspace: Workspace | null
-) {
-  let nameAtom:
-    | (WritableAtom<string, [string], void> & Atom<string>)
-    | undefined;
-  if (!blockSuiteWorkspace) {
-    nameAtom = emptyWorkspaceNameAtom;
-  } else if (!weakMap.has(blockSuiteWorkspace)) {
+export function useBlockSuiteWorkspaceName(blockSuiteWorkspace: Workspace) {
+  let nameAtom: StringAtom;
+  if (!weakMap.has(blockSuiteWorkspace)) {
     const baseAtom = atom<string>(
       blockSuiteWorkspace.meta.name ?? UNTITLED_WORKSPACE_NAME
     );
@@ -44,8 +31,7 @@ export function useBlockSuiteWorkspaceName(
     weakMap.set(blockSuiteWorkspace, writableAtom);
     nameAtom = writableAtom;
   } else {
-    nameAtom = weakMap.get(blockSuiteWorkspace);
-    assertExists(nameAtom);
+    nameAtom = weakMap.get(blockSuiteWorkspace) as StringAtom;
   }
   return useAtom(nameAtom);
 }

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

@@ -51,6 +51,5 @@ export function useBlockSuiteWorkspacePage(
 ): Page | null {
   const pageAtom = getAtom(blockSuiteWorkspace, pageId);
   assertExists(pageAtom);
-  const page = useAtomValue(pageAtom);
-  return page;
+  return useAtomValue(pageAtom);
 }