refactor: extract useBlockSuite from useAppState (#1001)

apps/web/src/components/edgeless-toolbar/index.tsx

@@ -17,9 +17,9 @@ import {
 import { MuiSlide } from '@affine/component';
 import { Tooltip } from '@affine/component';
 import useCurrentPageMeta from '@/hooks/use-current-page-meta';
-import { useAppState } from '@/providers/app-state-provider';
 import useHistoryUpdated from '@/hooks/use-history-update';
 import { useTranslation } from '@affine/i18n';
+import { useBlockSuite } from '@/store/workspace';
 
 const useToolbarList1 = () => {
   const { t } = useTranslation();
@@ -87,7 +87,7 @@ const useToolbarList1 = () => {
 const UndoRedo = () => {
   const [canUndo, setCanUndo] = useState(false);
   const [canRedo, setCanRedo] = useState(false);
-  const { currentPage } = useAppState();
+  const currentPage = useBlockSuite(store => store.currentPage);
   const onHistoryUpdated = useHistoryUpdated();
   const { t } = useTranslation();
   useEffect(() => {

apps/web/src/components/header/EditorHeader.tsx

@@ -6,17 +6,17 @@ import {
   StyledTitleWrapper,
 } from './styles';
 import { Content } from '@affine/component';
-import { useAppState } from '@/providers/app-state-provider';
 import EditorModeSwitch from '@/components/editor-mode-switch';
 import QuickSearchButton from './QuickSearchButton';
 import Header from './Header';
 import usePropsUpdated from '@/hooks/use-props-updated';
 import useCurrentPageMeta from '@/hooks/use-current-page-meta';
+import { useBlockSuite } from '@/store/workspace';
 
 export const EditorHeader = () => {
   const [title, setTitle] = useState('');
   const [isHover, setIsHover] = useState(false);
-  const { editor } = useAppState();
+  const editor = useBlockSuite(store => store.editor);
   const { trash: isTrash = false } = useCurrentPageMeta() || {};
   const onPropsUpdated = usePropsUpdated();
 

apps/web/src/components/header/header-right-items/EditorOptionMenu.tsx

@@ -11,14 +11,14 @@ import {
   PaperIcon,
   TrashIcon,
 } from '@blocksuite/icons';
-import { useAppState } from '@/providers/app-state-provider';
 import { usePageHelper } from '@/hooks/use-page-helper';
 import { useConfirm } from '@/providers/ConfirmProvider';
 import useCurrentPageMeta from '@/hooks/use-current-page-meta';
 import { toast } from '@affine/component';
 import { useTranslation } from '@affine/i18n';
+import { useBlockSuite } from '@/store/workspace';
 const PopoverContent = () => {
-  const { editor } = useAppState();
+  const editor = useBlockSuite(store => store.editor);
   const { toggleFavoritePage, toggleDeletePage } = usePageHelper();
   const { changePageMode } = usePageHelper();
   const confirm = useConfirm(store => store.confirm);

apps/web/src/components/help-island/index.tsx

@@ -11,7 +11,7 @@ import { Tooltip } from '@affine/component';
 import { useTranslation } from '@affine/i18n';
 import { useModal } from '@/store/globalModal';
 import { MuiFade } from '@affine/component';
-import { useAppState } from '@/providers/app-state-provider';
+import { useBlockSuite } from '@/store/workspace';
 export type IslandItemNames = 'contact' | 'shortcuts';
 export const HelpIsland = ({
   showList = ['contact', 'shortcuts'],
@@ -20,7 +20,7 @@ export const HelpIsland = ({
 }) => {
   const [spread, setShowSpread] = useState(false);
   const { triggerShortcutsModal, triggerContactModal } = useModal();
-  const { blockHub } = useAppState();
+  const blockHub = useBlockSuite(store => store.blockHub);
   const { t } = useTranslation();
 
   useEffect(() => {

apps/web/src/hooks/use-current-page-meta.ts

@@ -1,19 +1,23 @@
 import { useCallback, useEffect, useState } from 'react';
-import { useAppState, PageMeta } from '@/providers/app-state-provider';
+import { PageMeta } from '@/providers/app-state-provider';
+import { useBlockSuite } from '@/store/workspace';
 
 export const useCurrentPageMeta = (): PageMeta | null => {
-  const { currentPage, currentWorkspace } = useAppState();
+  const currentPage = useBlockSuite(store => store.currentPage);
+  const currentBlockSuiteWorkspace = useBlockSuite(
+    store => store.currentWorkspace
+  );
 
   const pageMetaHandler = useCallback((): PageMeta | null => {
-    if (!currentPage || !currentWorkspace) {
+    if (!currentPage || !currentBlockSuiteWorkspace) {
       return null;
     }
     return (
-      (currentWorkspace.blocksuiteWorkspace?.meta.pageMetas.find(
+      (currentBlockSuiteWorkspace.meta.pageMetas.find(
         p => p.id === currentPage.id
       ) as PageMeta) ?? null
     );
-  }, [currentPage, currentWorkspace]);
+  }, [currentPage, currentBlockSuiteWorkspace]);
 
   const [currentPageMeta, setCurrentPageMeta] = useState<PageMeta | null>(
     pageMetaHandler
@@ -22,16 +26,14 @@ export const useCurrentPageMeta = (): PageMeta | null => {
   useEffect(() => {
     setCurrentPageMeta(pageMetaHandler);
 
-    const dispose = currentWorkspace?.blocksuiteWorkspace?.meta.pagesUpdated.on(
-      () => {
-        setCurrentPageMeta(pageMetaHandler);
-      }
-    ).dispose;
+    const dispose = currentBlockSuiteWorkspace?.meta.pagesUpdated.on(() => {
+      setCurrentPageMeta(pageMetaHandler);
+    }).dispose;
 
     return () => {
       dispose?.();
     };
-  }, [currentPage, currentWorkspace, pageMetaHandler]);
+  }, [currentPage, currentBlockSuiteWorkspace, pageMetaHandler]);
 
   return currentPageMeta;
 };

apps/web/src/hooks/use-history-update.ts

@@ -1,13 +1,12 @@
 import { Page } from '@blocksuite/store';
-import { useAppState } from '@/providers/app-state-provider';
 import { useEffect, useRef } from 'react';
+import { useBlockSuite } from '@/store/workspace';
 
 export type EventCallBack<T> = (callback: (props: T) => void) => void;
 export type UseHistoryUpdated = (page?: Page) => EventCallBack<Page>;
 
 export const useHistoryUpdate: UseHistoryUpdated = () => {
-  const { currentPage } = useAppState();
-
+  const currentPage = useBlockSuite(store => store.currentPage);
   const callbackQueue = useRef<((page: Page) => void)[]>([]);
 
   useEffect(() => {

apps/web/src/hooks/use-page-helper.ts

@@ -5,6 +5,7 @@ import { EditorContainer } from '@blocksuite/editor';
 import { useChangePageMeta } from '@/hooks/use-change-page-meta';
 import { useRouter } from 'next/router';
 import { WorkspaceUnit } from '@affine/datacenter';
+import { useBlockSuite } from '@/store/workspace';
 
 export type EditorHandlers = {
   createPage: (params?: {
@@ -39,7 +40,8 @@ const getPageMeta = (workspace: WorkspaceUnit | null, pageId: string) => {
 export const usePageHelper = (): EditorHandlers => {
   const router = useRouter();
   const changePageMeta = useChangePageMeta();
-  const { currentWorkspace, editor } = useAppState();
+  const editor = useBlockSuite(store => store.editor);
+  const { currentWorkspace } = useAppState();
 
   return {
     createPage: ({

apps/web/src/hooks/use-props-updated.ts

@@ -1,6 +1,6 @@
 import { useEffect, useRef } from 'react';
 import { EditorContainer } from '@blocksuite/editor';
-import { useAppState } from '@/providers/app-state-provider';
+import { useBlockSuite } from '@/store/workspace';
 export type EventCallBack<T> = (callback: (props: T) => void) => void;
 
 export type UsePropsUpdated = (
@@ -8,7 +8,7 @@ export type UsePropsUpdated = (
 ) => EventCallBack<EditorContainer>;
 
 export const usePropsUpdated: UsePropsUpdated = () => {
-  const { editor } = useAppState();
+  const editor = useBlockSuite(store => store.editor);
 
   const callbackQueue = useRef<((editor: EditorContainer) => void)[]>([]);
 

apps/web/src/pages/_app.tsx

@@ -23,6 +23,7 @@ import Head from 'next/head';
 import '@affine/i18n';
 import { useTranslation } from '@affine/i18n';
 import React from 'react';
+import { BlockSuiteProvider } from '@/store/workspace';
 
 const ThemeProvider = dynamic(() => import('@/providers/ThemeProvider'), {
   ssr: false,
@@ -67,20 +68,22 @@ const App = ({ Component, pageProps }: AppPropsWithLayout) => {
         <title>AFFiNE</title>
       </Head>
       <Logger />
-      <ProviderComposer
-        contexts={[
-          <ThemeProvider key="ThemeProvider" />,
-          <AppStateProvider key="appStateProvider" />,
-          <ModalProvider key="ModalProvider" />,
-          <ConfirmProvider key="ConfirmProvider" />,
-        ]}
-      >
-        {NoNeedAppStatePageList.includes(router.route) ? (
-          getLayout(<Component {...pageProps} />)
-        ) : (
-          <AppDefender>{getLayout(<Component {...pageProps} />)}</AppDefender>
-        )}
-      </ProviderComposer>
+      <BlockSuiteProvider key="BlockSuiteProvider">
+        <ProviderComposer
+          contexts={[
+            <ThemeProvider key="ThemeProvider" />,
+            <AppStateProvider key="appStateProvider" />,
+            <ModalProvider key="ModalProvider" />,
+            <ConfirmProvider key="ConfirmProvider" />,
+          ]}
+        >
+          {NoNeedAppStatePageList.includes(router.route) ? (
+            getLayout(<Component {...pageProps} />)
+          ) : (
+            <AppDefender>{getLayout(<Component {...pageProps} />)}</AppDefender>
+          )}
+        </ProviderComposer>
+      </BlockSuiteProvider>
     </>
   );
 };

apps/web/src/pages/workspace/[workspaceId]/[pageId].tsx

@@ -1,10 +1,4 @@
-import {
-  PropsWithChildren,
-  ReactElement,
-  useCallback,
-  useEffect,
-  useState,
-} from 'react';
+import { PropsWithChildren, ReactElement, useEffect, useState } from 'react';
 import { EditorHeader } from '@/components/header';
 import MobileModal from '@/components/mobile-modal';
 import { useAppState } from '@/providers/app-state-provider';
@@ -13,20 +7,16 @@ import WorkspaceLayout from '@/components/workspace-layout';
 import { useRouter } from 'next/router';
 import { usePageHelper } from '@/hooks/use-page-helper';
 import dynamic from 'next/dynamic';
-import { EditorContainer } from '@blocksuite/editor';
 import Head from 'next/head';
 import { useTranslation } from '@affine/i18n';
-import { BlockHub } from '@blocksuite/blocks';
+import { useBlockSuite } from '@/store/workspace';
 const DynamicBlocksuite = dynamic(() => import('@/components/editor'), {
   ssr: false,
 });
 
 const BlockHubAppender = () => {
-  const { setBlockHub, editor } = useAppState();
-  const setBlockHubHandler = useCallback(
-    (blockHub: BlockHub) => setBlockHub.current(blockHub),
-    [setBlockHub]
-  );
+  const setBlockHub = useBlockSuite(store => store.setBlockHub);
+  const editor = useBlockSuite(store => store.editor);
   useEffect(() => {
     let blockHubElement: HTMLElement | null = null;
 
@@ -37,23 +27,20 @@ const BlockHubAppender = () => {
         return;
       }
       blockHubElement = blockHub;
-      // setBlockHubHandler(blockHub);
+      setBlockHub(blockHub);
       toolWrapper.appendChild(blockHub);
     });
     return () => {
       blockHubElement?.remove();
     };
-  }, [editor, setBlockHubHandler]);
+  }, [editor, setBlockHub]);
   return null;
 };
 
 const Page: NextPageWithLayout = () => {
-  const { currentPage, currentWorkspace, setEditor } = useAppState();
-
-  const setEditorHandler = useCallback(
-    (editor: EditorContainer) => setEditor.current(editor),
-    [setEditor]
-  );
+  const currentPage = useBlockSuite(store => store.currentPage);
+  const setEditor = useBlockSuite(store => store.setEditor);
+  const { currentWorkspace } = useAppState();
 
   const { t } = useTranslation();
 
@@ -70,7 +57,7 @@ const Page: NextPageWithLayout = () => {
           <DynamicBlocksuite
             page={currentPage}
             workspace={currentWorkspace.blocksuiteWorkspace}
-            setEditor={setEditorHandler}
+            setEditor={setEditor}
           />
           <BlockHubAppender />
         </>
@@ -82,7 +69,8 @@ const Page: NextPageWithLayout = () => {
 const PageDefender = ({ children }: PropsWithChildren) => {
   const router = useRouter();
   const [pageLoaded, setPageLoaded] = useState(false);
-  const { currentWorkspace, loadPage } = useAppState();
+  const loadPage = useBlockSuite(store => store.loadPage);
+  const { currentWorkspace } = useAppState();
   const { createPage } = usePageHelper();
 
   useEffect(() => {

apps/web/src/providers/app-state-provider/Provider.tsx

@@ -9,6 +9,7 @@ import {
 } from './interface';
 import { createDefaultWorkspace } from './utils';
 import { User } from '@affine/datacenter';
+import { useBlockSuiteApi } from '@/store/workspace';
 
 export interface Disposable {
   dispose(): void;
@@ -22,6 +23,7 @@ export const useAppState = () => useContext(AppState);
 export const AppStateProvider = ({
   children,
 }: PropsWithChildren<AppStateContextProps>) => {
+  const blocksuiteApi = useBlockSuiteApi();
   const [appState, setAppState] = useState<AppStateValue>({} as AppStateValue);
   const { dataCenter } = appState;
   const [blobState, setBlobState] = useState(false);
@@ -43,9 +45,6 @@ export const AppStateProvider = ({
         workspaceList: dataCenter.workspaces,
         currentWorkspace: null,
         pageList: [],
-        currentPage: null,
-        editor: null,
-        blockHub: null,
         synced: true,
         isOwner: false,
       });
@@ -86,19 +85,6 @@ export const AppStateProvider = ({
     );
   }, [dataCenter]);
 
-  const loadPage = useRef<AppStateFunction['loadPage']>();
-  loadPage.current = pageId => {
-    const { currentWorkspace, currentPage } = appState;
-    if (pageId === currentPage?.id) {
-      return;
-    }
-    const page = currentWorkspace?.blocksuiteWorkspace?.getPage(pageId) || null;
-    setAppState({
-      ...appState,
-      currentPage: page,
-    });
-  };
-
   const loadWorkspace: AppStateFunction['loadWorkspace'] =
     useRef() as AppStateFunction['loadWorkspace'];
   loadWorkspace.current = async (workspaceId, abort) => {
@@ -136,13 +122,13 @@ export const AppStateProvider = ({
 
     const pageList =
       (workspace?.blocksuiteWorkspace?.meta.pageMetas as PageMeta[]) ?? [];
+    if (workspace?.blocksuiteWorkspace) {
+      blocksuiteApi.getState().setWorkspace(workspace.blocksuiteWorkspace);
+    }
     setAppState({
       ...appState,
       currentWorkspace: workspace,
       pageList: pageList,
-      currentPage: null,
-      editor: null,
-      blockHub: null,
       isOwner,
     });
 
@@ -171,23 +157,6 @@ export const AppStateProvider = ({
     };
   }, [appState.currentWorkspace]);
 
-  const setEditor: AppStateFunction['setEditor'] =
-    useRef() as AppStateFunction['setEditor'];
-  setEditor.current = editor => {
-    setAppState({
-      ...appState,
-      editor,
-    });
-  };
-  const setBlockHub: AppStateFunction['setBlockHub'] =
-    useRef() as AppStateFunction['setBlockHub'];
-  setBlockHub.current = blockHub => {
-    setAppState({
-      ...appState,
-      blockHub,
-    });
-  };
-
   const login = async () => {
     const { dataCenter } = appState;
     try {
@@ -213,9 +182,6 @@ export const AppStateProvider = ({
     <AppState.Provider
       value={{
         ...appState,
-        setEditor,
-        setBlockHub,
-        loadPage: loadPage.current,
         loadWorkspace: loadWorkspace,
         login,
         logout,

apps/web/src/providers/app-state-provider/interface.ts

@@ -1,6 +1,5 @@
 import { DataCenter, User, WorkspaceUnit } from '@affine/datacenter';
 import type { EditorContainer } from '@blocksuite/editor';
-import { BlockHub } from '@blocksuite/blocks';
 
 import type {
   Page as StorePage,
@@ -21,22 +20,15 @@ export type AppStateValue = {
   workspaceList: WorkspaceUnit[];
   currentWorkspace: WorkspaceUnit | null;
   pageList: PageMeta[];
-  currentPage: StorePage | null;
-  editor?: EditorContainer | null;
-  blockHub?: BlockHub | null;
   synced: boolean;
   isOwner?: boolean;
   blobDataSynced?: boolean;
 };
 
 export type AppStateFunction = {
-  setEditor: MutableRefObject<(page: EditorContainer) => void>;
-  setBlockHub: MutableRefObject<(BlockHub: BlockHub) => void>;
-
   loadWorkspace: MutableRefObject<
     (workspaceId: string, abort?: AbortSignal) => Promise<WorkspaceUnit | null>
   >;
-  loadPage: (pageId: string) => void;
 
   login: () => Promise<User | null>;
   logout: () => Promise<void>;

apps/web/src/store/workspace/index.tsx

@@ -0,0 +1,98 @@
+import type React from 'react';
+import { createContext, useContext, useMemo } from 'react';
+import { createStore, useStore } from 'zustand';
+import { combine, subscribeWithSelector } from 'zustand/middleware';
+import type { UseBoundStore } from 'zustand/react';
+import type { Page } from '@blocksuite/store';
+import type { BlockHub } from '@blocksuite/blocks';
+import type { Workspace } from '@blocksuite/store';
+import type { EditorContainer } from '@blocksuite/editor';
+
+export type BlockSuiteState = {
+  currentWorkspace: Workspace | null;
+  editor: EditorContainer | null;
+  currentPage: Page | null;
+  blockHub: BlockHub | null;
+};
+
+export type BlockSuiteActions = {
+  loadPage: (pageId: string) => void;
+  setEditor: (editor: EditorContainer) => void;
+  setWorkspace: (workspace: Workspace) => void;
+  setBlockHub: (blockHub: BlockHub) => void;
+};
+
+const create = () =>
+  createStore(
+    subscribeWithSelector(
+      combine<BlockSuiteState, BlockSuiteActions>(
+        {
+          currentWorkspace: null,
+          currentPage: null,
+          blockHub: null,
+          editor: null,
+        },
+        (set, get) => ({
+          setWorkspace: workspace => {
+            set({
+              currentWorkspace: workspace,
+            });
+          },
+          setEditor: editor => {
+            set({
+              editor,
+            });
+          },
+          loadPage: pageId => {
+            const { currentWorkspace } = get();
+            if (currentWorkspace === null) {
+              console.warn('currentWorkspace is null');
+              return;
+            }
+            const page = currentWorkspace.getPage(pageId);
+            if (page === null) {
+              console.warn('cannot find page ', pageId);
+              return;
+            }
+            set({
+              currentPage: page,
+            });
+          },
+          setBlockHub: blockHub => {
+            set({
+              blockHub,
+            });
+          },
+        })
+      )
+    )
+  );
+type Store = ReturnType<typeof create>;
+
+const BlockSuiteContext = createContext<Store | null>(null);
+
+export const useBlockSuiteApi = () => {
+  const api = useContext(BlockSuiteContext);
+  if (!api) {
+    throw new Error('cannot find modal context');
+  }
+  return api;
+};
+
+export const useBlockSuite: UseBoundStore<Store> = ((
+  selector: Parameters<UseBoundStore<Store>>[0],
+  equals: Parameters<UseBoundStore<Store>>[1]
+) => {
+  const api = useBlockSuiteApi();
+  return useStore(api, selector, equals);
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+}) as any;
+
+export const BlockSuiteProvider: React.FC<React.PropsWithChildren> =
+  function ModelProvider({ children }) {
+    return (
+      <BlockSuiteContext.Provider value={useMemo(() => create(), [])}>
+        {children}
+      </BlockSuiteContext.Provider>
+    );
+  };