refactor(frontend): improve HMR experience

Use a HMR acceptor as in
https://github.com/vitejs/vite/issues/10227#issuecomment-1256969751

Also updates frontend tooling to the latest version
(yarn now support typescript 4.8.4)

12 files changed, 145 insertions, 80 deletions

diff --git a/subprojects/frontend/src/App.tsx b/subprojects/frontend/src/App.tsx
index 90514044..a471690b 100644
--- a/subprojects/frontend/src/App.tsx
+++ b/subprojects/frontend/src/App.tsx
@@ -1,21 +1,33 @@
-import Grow from '@mui/material/Grow';
-import Stack from '@mui/material/Stack';
-import { SnackbarProvider } from 'notistack';
-import React from 'react';
+import Box from '@mui/material/Box';
+import CssBaseline from '@mui/material/CssBaseline';
+import React, { StrictMode, Suspense, lazy } from 'react';
 
-import TopBar from './TopBar';
-import UpdateNotification from './UpdateNotification';
-import EditorPane from './editor/EditorPane';
+import Loading from './Loading';
+import type RootStore from './RootStore';
+import RootStoreProvider from './RootStoreProvider';
+import WindowControlsOverlayColor from './WindowControlsOverlayColor';
+import ThemeProvider from './theme/ThemeProvider';
 
-export default function App(): JSX.Element {
+const Refinery = lazy(() => import('./Refinery.js'));
+
+export default function App({
+  rootStore,
+}: {
+  rootStore: RootStore;
+}): JSX.Element {
   return (
-    // @ts-expect-error -- notistack has problems with `exactOptionalPropertyTypes
-    <SnackbarProvider TransitionComponent={Grow}>
-      <UpdateNotification />
-      <Stack direction="column" height="100vh" overflow="auto">
-        <TopBar />
-        <EditorPane />
-      </Stack>
-    </SnackbarProvider>
+    <StrictMode>
+      <RootStoreProvider rootStore={rootStore}>
+        <ThemeProvider>
+          <CssBaseline enableColorScheme />
+          <WindowControlsOverlayColor />
+          <Box height="100vh" overflow="auto">
+            <Suspense fallback={<Loading />}>
+              <Refinery />
+            </Suspense>
+          </Box>
+        </ThemeProvider>
+      </RootStoreProvider>
+    </StrictMode>
   );
 }
diff --git a/subprojects/frontend/src/Refinery.tsx b/subprojects/frontend/src/Refinery.tsx
new file mode 100644
index 00000000..41f835b3
--- /dev/null
+++ b/subprojects/frontend/src/Refinery.tsx
@@ -0,0 +1,21 @@
+import Grow from '@mui/material/Grow';
+import Stack from '@mui/material/Stack';
+import { SnackbarProvider } from 'notistack';
+import React from 'react';
+
+import TopBar from './TopBar';
+import UpdateNotification from './UpdateNotification';
+import EditorPane from './editor/EditorPane';
+
+export default function Refinery(): JSX.Element {
+  return (
+    // @ts-expect-error -- notistack has problems with `exactOptionalPropertyTypes
+    <SnackbarProvider TransitionComponent={Grow}>
+      <UpdateNotification />
+      <Stack direction="column" height="100vh" overflow="auto">
+        <TopBar />
+        <EditorPane />
+      </Stack>
+    </SnackbarProvider>
+  );
+}
diff --git a/subprojects/frontend/src/RootStore.tsx b/subprojects/frontend/src/RootStore.ts
index c1674a3c..54a80501 100644
--- a/subprojects/frontend/src/RootStore.tsx
+++ b/subprojects/frontend/src/RootStore.ts
@@ -1,6 +1,5 @@
 import { getLogger } from 'loglevel';
 import { makeAutoObservable, runInAction } from 'mobx';
-import React, { createContext, useContext } from 'react';
 
 import PWAStore from './PWAStore';
 import type EditorStore from './editor/EditorStore';
@@ -15,6 +14,8 @@ export default class RootStore {
 
   readonly themeStore: ThemeStore;
 
+  disposed = false;
+
   constructor(initialValue: string) {
     this.pwaStore = new PWAStore();
     this.themeStore = new ThemeStore();
@@ -25,6 +26,9 @@ export default class RootStore {
     import('./editor/EditorStore')
       .then(({ default: EditorStore }) => {
         runInAction(() => {
+          if (this.disposed) {
+            return;
+          }
           this.editorStore = new EditorStore(initialValue, this.pwaStore);
         });
       })
@@ -32,29 +36,12 @@ export default class RootStore {
         log.error('Failed to load EditorStore', error);
       });
   }
-}
-
-const StoreContext = createContext<RootStore | undefined>(undefined);
 
-export interface RootStoreProviderProps {
-  children: JSX.Element;
-
-  rootStore: RootStore;
-}
-
-export function RootStoreProvider({
-  children,
-  rootStore,
-}: RootStoreProviderProps): JSX.Element {
-  return (
-    <StoreContext.Provider value={rootStore}>{children}</StoreContext.Provider>
-  );
-}
-
-export const useRootStore = (): RootStore => {
-  const rootStore = useContext(StoreContext);
-  if (!rootStore) {
-    throw new Error('useRootStore must be used within RootStoreProvider');
+  dispose(): void {
+    if (this.disposed) {
+      return;
+    }
+    this.editorStore?.dispose();
+    this.disposed = true;
   }
-  return rootStore;
-};
+}
diff --git a/subprojects/frontend/src/RootStoreProvider.tsx b/subprojects/frontend/src/RootStoreProvider.tsx
new file mode 100644
index 00000000..70ac7776
--- /dev/null
+++ b/subprojects/frontend/src/RootStoreProvider.tsx
@@ -0,0 +1,29 @@
+import React, { type ReactNode, createContext, useContext } from 'react';
+
+import type RootStore from './RootStore';
+
+const StoreContext = createContext<RootStore | undefined>(undefined);
+
+export function useRootStore(): RootStore {
+  const rootStore = useContext(StoreContext);
+  if (!rootStore) {
+    throw new Error('useRootStore must be used within RootStoreProvider');
+  }
+  return rootStore;
+}
+
+export default function RootStoreProvider({
+  children,
+  rootStore,
+}: {
+  children?: ReactNode;
+  rootStore: RootStore;
+}): JSX.Element {
+  return (
+    <StoreContext.Provider value={rootStore}>{children}</StoreContext.Provider>
+  );
+}
+
+RootStoreProvider.defaultProps = {
+  children: undefined,
+};
diff --git a/subprojects/frontend/src/ToggleDarkModeButton.tsx b/subprojects/frontend/src/ToggleDarkModeButton.tsx
index 5602f3bc..52606242 100644
--- a/subprojects/frontend/src/ToggleDarkModeButton.tsx
+++ b/subprojects/frontend/src/ToggleDarkModeButton.tsx
@@ -4,7 +4,7 @@ import IconButton from '@mui/material/IconButton';
 import { observer } from 'mobx-react-lite';
 import React from 'react';
 
-import { useRootStore } from './RootStore';
+import { useRootStore } from './RootStoreProvider';
 
 export default observer(function ToggleDarkModeButton(): JSX.Element {
   const { themeStore } = useRootStore();
diff --git a/subprojects/frontend/src/TopBar.tsx b/subprojects/frontend/src/TopBar.tsx
index a060ff9e..79ea0b65 100644
--- a/subprojects/frontend/src/TopBar.tsx
+++ b/subprojects/frontend/src/TopBar.tsx
@@ -7,7 +7,7 @@ import { throttle } from 'lodash-es';
 import { observer } from 'mobx-react-lite';
 import React, { useEffect, useMemo, useState } from 'react';
 
-import { useRootStore } from './RootStore';
+import { useRootStore } from './RootStoreProvider';
 import ToggleDarkModeButton from './ToggleDarkModeButton';
 import GenerateButton from './editor/GenerateButton';
 
diff --git a/subprojects/frontend/src/UpdateNotification.tsx b/subprojects/frontend/src/UpdateNotification.tsx
index d260e3b7..8c46186a 100644
--- a/subprojects/frontend/src/UpdateNotification.tsx
+++ b/subprojects/frontend/src/UpdateNotification.tsx
@@ -2,7 +2,7 @@ import Button from '@mui/material/Button';
 import { observer } from 'mobx-react-lite';
 import React, { useEffect } from 'react';
 
-import { useRootStore } from './RootStore';
+import { useRootStore } from './RootStoreProvider';
 import { ContrastThemeProvider } from './theme/ThemeProvider';
 import useDelayedSnackbar from './utils/useDelayedSnackbar';
 
diff --git a/subprojects/frontend/src/editor/EditorPane.tsx b/subprojects/frontend/src/editor/EditorPane.tsx
index 199a8e1d..0db4b892 100644
--- a/subprojects/frontend/src/editor/EditorPane.tsx
+++ b/subprojects/frontend/src/editor/EditorPane.tsx
@@ -7,7 +7,7 @@ import useMediaQuery from '@mui/material/useMediaQuery';
 import { observer } from 'mobx-react-lite';
 import React, { useState } from 'react';
 
-import { useRootStore } from '../RootStore';
+import { useRootStore } from '../RootStoreProvider';
 
 import ConnectionStatusNotification from './ConnectionStatusNotification';
 import EditorArea from './EditorArea';
diff --git a/subprojects/frontend/src/editor/EditorStore.ts b/subprojects/frontend/src/editor/EditorStore.ts
index d13f644c..4ee24779 100644
--- a/subprojects/frontend/src/editor/EditorStore.ts
+++ b/subprojects/frontend/src/editor/EditorStore.ts
@@ -263,4 +263,8 @@ export default class EditorStore {
     this.client.formatText();
     return true;
   }
+
+  dispose(): void {
+    this.client.dispose();
+  }
 }
diff --git a/subprojects/frontend/src/index.tsx b/subprojects/frontend/src/index.tsx
index f7fe61b9..2c0259bf 100644
--- a/subprojects/frontend/src/index.tsx
+++ b/subprojects/frontend/src/index.tsx
@@ -1,16 +1,9 @@
-import Box from '@mui/material/Box';
-import CssBaseline from '@mui/material/CssBaseline';
 import { configure } from 'mobx';
-import React, { Suspense, lazy } from 'react';
-import { createRoot } from 'react-dom/client';
+import React from 'react';
+import { type Root, createRoot } from 'react-dom/client';
 
-import Loading from './Loading';
-import RootStore, { RootStoreProvider } from './RootStore';
-import WindowControlsOverlayColor from './WindowControlsOverlayColor';
-import ThemeProvider from './theme/ThemeProvider';
-import getLogger from './utils/getLogger';
-
-const log = getLogger('index');
+import App from './App';
+import RootStore from './RootStore';
 
 const initialValue = `class Family {
     contains Person[] members
@@ -57,30 +50,45 @@ configure({
   enforceActions: 'always',
 });
 
-const rootStore = new RootStore(initialValue);
+let HotRootStore = RootStore;
+let HotApp = App;
+
+function createStore(): RootStore {
+  return new HotRootStore(initialValue);
+}
 
-const App = lazy(() => import('./App.js'));
+let rootStore = createStore();
 
-const app = (
-  <React.StrictMode>
-    <RootStoreProvider rootStore={rootStore}>
-      <ThemeProvider>
-        <CssBaseline enableColorScheme />
-        <WindowControlsOverlayColor />
-        <Box height="100vh" overflow="auto">
-          <Suspense fallback={<Loading />}>
-            <App />
-          </Suspense>
-        </Box>
-      </ThemeProvider>
-    </RootStoreProvider>
-  </React.StrictMode>
-);
+let root: Root | undefined;
 
-const rootElement = document.getElementById('app');
-if (rootElement === null) {
-  log.error('Root element not found');
-} else {
-  const root = createRoot(rootElement);
-  root.render(app);
+function render(): void {
+  root?.render(<HotApp rootStore={rootStore} />);
 }
+
+if (import.meta.hot) {
+  import.meta.hot.accept('./App', (module) => {
+    if (module === undefined) {
+      return;
+    }
+    ({ default: HotApp } = module as unknown as typeof import('./App'));
+    render();
+  });
+  import.meta.hot.accept('./RootStore', (module) => {
+    if (module === undefined) {
+      return;
+    }
+    ({ default: HotRootStore } =
+      module as unknown as typeof import('./RootStore'));
+    rootStore.dispose();
+    rootStore = createStore();
+    render();
+  });
+}
+
+document.addEventListener('DOMContentLoaded', () => {
+  const rootElement = document.getElementById('app');
+  if (rootElement !== null) {
+    root = createRoot(rootElement);
+    render();
+  }
+});
diff --git a/subprojects/frontend/src/theme/ThemeProvider.tsx b/subprojects/frontend/src/theme/ThemeProvider.tsx
index d990fd5d..060a5e9a 100644
--- a/subprojects/frontend/src/theme/ThemeProvider.tsx
+++ b/subprojects/frontend/src/theme/ThemeProvider.tsx
@@ -14,7 +14,7 @@ import {
 import { observer } from 'mobx-react-lite';
 import React, { type ReactNode, createContext, useContext } from 'react';
 
-import { useRootStore } from '../RootStore';
+import { useRootStore } from '../RootStoreProvider';
 
 interface OuterPalette {
   background: string;
diff --git a/subprojects/frontend/src/xtext/XtextClient.ts b/subprojects/frontend/src/xtext/XtextClient.ts
index cd5d280d..e7d26ae6 100644
--- a/subprojects/frontend/src/xtext/XtextClient.ts
+++ b/subprojects/frontend/src/xtext/XtextClient.ts
@@ -120,4 +120,8 @@ export default class XtextClient {
       log.error('Error while formatting text', e);
     });
   }
+
+  dispose(): void {
+    this.webSocketClient.disconnect();
+  }
 }