Factor out navigation stuff

src/Navigation.tsx

@@ -26,7 +26,7 @@ import {BottomBar} from './view/shell/bottom-bar/BottomBar'
 import {buildStateObject} from 'lib/routes/helpers'
 import {State, RouteParams} from 'lib/routes/types'
 import {colors} from 'lib/styles'
-import {isNative, isWeb} from 'platform/detection'
+import {isNative} from 'platform/detection'
 import {useColorSchemeStyle} from 'lib/hooks/useColorSchemeStyle'
 import {router} from './routes'
 import {usePalette} from 'lib/hooks/usePalette'
@@ -41,6 +41,7 @@ import {
   setEmailConfirmationRequested,
 } from './state/shell/reminders'
 import {init as initAnalytics} from './lib/analytics/analytics'
+import {useWebScrollRestoration} from './lib/hooks/useWebScrollRestoration'
 
 import {HomeScreen} from './view/screens/Home'
 import {SearchScreen} from './view/screens/Search'
@@ -385,9 +386,6 @@ function MyProfileTabNavigator() {
   )
 }
 
-const webScrollPositions = new Map()
-let webFocusedScreen = null
-
 /**
  * The FlatNavigator is used by Web to represent the routes
  * in a single ("flat") stack.
@@ -395,24 +393,11 @@ let webFocusedScreen = null
 const FlatNavigator = () => {
   const pal = usePalette('default')
   const numUnread = useUnreadNotifications()
-
   const title = (page: string) => bskyTitle(page, numUnread)
+  const screenListeners = useWebScrollRestoration()
   return (
     <Flat.Navigator
-      screenListeners={{
+      screenListeners={screenListeners}
-        beforeRemove(e) {
-          if (isWeb) {
-            webScrollPositions.delete(e.target)
-          }
-        },
-        focus(e) {
-          if (isWeb) {
-            const scrollY = webScrollPositions.get(e.target) ?? 0
-            window.scrollTo(0, scrollY)
-            webFocusedScreen = e.target
-          }
-        },
-      }}
       screenOptions={{
         gestureEnabled: true,
         fullScreenGestureEnabled: true,
@@ -515,16 +500,6 @@ function RoutesContainer({children}: React.PropsWithChildren<{}>) {
       onReady={() => {
         logModuleInitTime()
         onReady()
-
-        if (isWeb) {
-          history.scrollRestoration = 'manual'
-          // TODO: Clean up?
-          navigationRef.current?.addListener('__unsafe_action__', e => {
-            if (webFocusedScreen) {
-              webScrollPositions.set(webFocusedScreen, window.scrollY)
-            }
-          })
-        }
       }}>
       {children}
     </NavigationContainer>

src/lib/hooks/useWebScrollRestoration.native.ts

@@ -0,0 +1,3 @@
+export function useWebScrollRestoration() {
+  return undefined
+}

src/lib/hooks/useWebScrollRestoration.ts

@@ -0,0 +1,51 @@
+import {useMemo, useState, useEffect} from 'react'
+import {EventArg, useNavigation} from '@react-navigation/core'
+
+if ('scrollRestoration' in history) {
+  // Tell the brower not to mess with the scroll.
+  // We're doing that manuall below.
+  history.scrollRestoration = 'manual'
+}
+
+function createInitialScrollState() {
+  return {
+    scrollYs: new Map(),
+    focusedKey: null as string | null,
+  }
+}
+
+export function useWebScrollRestoration() {
+  const [state] = useState(createInitialScrollState)
+  const navigation = useNavigation()
+
+  useEffect(() => {
+    function onDispatch() {
+      if (state.focusedKey) {
+        // Remember where we were for later.
+        state.scrollYs.set(state.focusedKey, window.scrollY)
+      }
+    }
+    // We want to intercept any push/pop/replace *before* the re-render.
+    // There is no official way to do this yet, but this works okay for now.
+    // https://twitter.com/satya164/status/1737301243519725803
+    navigation.addListener('__unsafe_action__' as any, onDispatch)
+    return () => {
+      navigation.removeListener('__unsafe_action__' as any, onDispatch)
+    }
+  }, [state, navigation])
+
+  const screenListeners = useMemo(
+    () => ({
+      beforeRemove(e: EventArg<'beforeRemove', boolean | undefined, unknown>) {
+        state.scrollYs.delete(e.target)
+      },
+      focus(e: EventArg<'focus', boolean | undefined, unknown>) {
+        const scrollY = state.scrollYs.get(e.target) ?? 0
+        window.scrollTo(0, scrollY)
+        state.focusedKey = e.target ?? null
+      },
+    }),
+    [state],
+  )
+  return screenListeners
+}