From 1b89ec7221dc87792fc56886b170ff6f1d06e230 Mon Sep 17 00:00:00 2001
From: Andrei Alecu <andreialecu@users.noreply.github.com>
Date: Wed, 3 Apr 2024 19:06:44 +0300
Subject: [PATCH] fix: lazy FlashList scroll sync issue

---
 .../src/Shared/ExampleComponentFlashList.tsx  |  2 +-
 .../ExampleComponentMasonryFlashList.tsx      |  2 +-
 src/FlashList.tsx                             | 30 +++++++++++++------
 src/MasonryFlashList.tsx                      | 30 +++++++++++++------
 src/hooks.tsx                                 | 19 ++++++++++--
 5 files changed, 60 insertions(+), 23 deletions(-)

diff --git a/example/src/Shared/ExampleComponentFlashList.tsx b/example/src/Shared/ExampleComponentFlashList.tsx
index 38b68ad..2ca75c6 100644
--- a/example/src/Shared/ExampleComponentFlashList.tsx
+++ b/example/src/Shared/ExampleComponentFlashList.tsx
@@ -19,7 +19,7 @@ type Props = {
 const Example = React.forwardRef<CollapsibleRef, Props>(
   ({ emptyContacts, ...props }, ref) => {
     return (
-      <Tabs.Container ref={ref} headerHeight={HEADER_HEIGHT} {...props}>
+      <Tabs.Container ref={ref} headerHeight={HEADER_HEIGHT} lazy {...props}>
         {props.hideArticleTab ? (
           <Tabs.Tab name="article" label="Article">
             <Article />
diff --git a/example/src/Shared/ExampleComponentMasonryFlashList.tsx b/example/src/Shared/ExampleComponentMasonryFlashList.tsx
index f15eb76..e3a00ad 100644
--- a/example/src/Shared/ExampleComponentMasonryFlashList.tsx
+++ b/example/src/Shared/ExampleComponentMasonryFlashList.tsx
@@ -19,7 +19,7 @@ type Props = {
 const Example = React.forwardRef<CollapsibleRef, Props>(
   ({ emptyContacts, ...props }, ref) => {
     return (
-      <Tabs.Container ref={ref} headerHeight={HEADER_HEIGHT} {...props}>
+      <Tabs.Container ref={ref} headerHeight={HEADER_HEIGHT} lazy {...props}>
         {props.hideArticleTab ? (
           <Tabs.Tab name="article" label="Article">
             <Article />
diff --git a/src/FlashList.tsx b/src/FlashList.tsx
index b18949f..ba8ee6e 100644
--- a/src/FlashList.tsx
+++ b/src/FlashList.tsx
@@ -3,10 +3,12 @@ import type {
   FlashList as SPFlashList,
 } from '@shopify/flash-list'
 import React, { useCallback } from 'react'
-import Animated from 'react-native-reanimated'
+import Animated, {
+  useSharedValue,
+  useAnimatedReaction,
+} from 'react-native-reanimated'
 
 import {
-  useAfterMountEffect,
   useChainCallback,
   useCollapsibleStyle,
   useConvertAnimatedToValue,
@@ -72,12 +74,22 @@ function FlashListImpl<R>(
 
   const { scrollHandler, enable } = useScrollHandlerY(name)
 
-  const onLayout = useAfterMountEffect(rest.onLayout, () => {
-    'worklet'
-    // we enable the scroll event after mounting
-    // otherwise we get an `onScroll` call with the initial scroll position which can break things
-    enable(true)
-  })
+  const hadLoad = useSharedValue(false)
+
+  const onLoad = useCallback(() => {
+    hadLoad.value = true
+  }, [hadLoad])
+
+  useAnimatedReaction(
+    () => {
+      return hadLoad.value
+    },
+    (ready) => {
+      if (ready) {
+        enable(true)
+      }
+    }
+  )
 
   const { progressViewOffset, contentContainerStyle } = useCollapsibleStyle()
 
@@ -141,7 +153,7 @@ function FlashListImpl<R>(
     // @ts-expect-error typescript complains about `unknown` in the memo, it should be T
     <FlashListMemo
       {...rest}
-      onLayout={onLayout}
+      onLoad={onLoad}
       ref={refWorkaround}
       contentContainerStyle={memoContentContainerStyle}
       bouncesZoom={false}
diff --git a/src/MasonryFlashList.tsx b/src/MasonryFlashList.tsx
index 7db0cc9..bc140e2 100644
--- a/src/MasonryFlashList.tsx
+++ b/src/MasonryFlashList.tsx
@@ -1,9 +1,11 @@
 import { MasonryFlashListProps, MasonryFlashListRef } from '@shopify/flash-list'
 import React, { useCallback } from 'react'
-import Animated from 'react-native-reanimated'
+import Animated, {
+  useAnimatedReaction,
+  useSharedValue,
+} from 'react-native-reanimated'
 
 import {
-  useAfterMountEffect,
   useChainCallback,
   useCollapsibleStyle,
   useConvertAnimatedToValue,
@@ -76,12 +78,22 @@ function MasonryFlashListImpl<R>(
 
   const { scrollHandler, enable } = useScrollHandlerY(name)
 
-  const onLayout = useAfterMountEffect(rest.onLayout, () => {
-    'worklet'
-    // we enable the scroll event after mounting
-    // otherwise we get an `onScroll` call with the initial scroll position which can break things
-    enable(true)
-  })
+  const hadLoad = useSharedValue(false)
+
+  const onLoad = useCallback(() => {
+    hadLoad.value = true
+  }, [hadLoad])
+
+  useAnimatedReaction(
+    () => {
+      return hadLoad.value
+    },
+    (ready) => {
+      if (ready) {
+        enable(true)
+      }
+    }
+  )
 
   const { progressViewOffset, contentContainerStyle } = useCollapsibleStyle()
 
@@ -145,7 +157,7 @@ function MasonryFlashListImpl<R>(
     // @ts-expect-error typescript complains about `unknown` in the memo, it should be T
     <MasonryFlashListMemo
       {...rest}
-      onLayout={onLayout}
+      onLoad={onLoad}
       contentContainerStyle={memoContentContainerStyle}
       ref={refWorkaround}
       bouncesZoom={false}
diff --git a/src/hooks.tsx b/src/hooks.tsx
index 75e9c9f..061759f 100644
--- a/src/hooks.tsx
+++ b/src/hooks.tsx
@@ -272,12 +272,27 @@ export const useScrollHandlerY = (name: TabName) => {
 
   const enabled = useSharedValue(false)
 
+  const scrollTo = useScroller()
+
   const enable = useCallback(
     (toggle: boolean) => {
       'worklet'
       enabled.value = toggle
+
+      if (toggle) {
+        const tabIndex = tabNames.value.findIndex((n) => n === name)
+
+        const ref = refMap[name]
+        scrollTo(
+          ref,
+          0,
+          scrollY.value[tabIndex],
+          false,
+          `[${name}] restore scroll position - enable`
+        )
+      }
     },
-    [enabled]
+    [enabled, name, refMap, scrollTo, scrollY.value, tabNames.value]
   )
 
   /**
@@ -293,8 +308,6 @@ export const useScrollHandlerY = (name: TabName) => {
     name,
   ])
 
-  const scrollTo = useScroller()
-
   const scrollAnimation = useSharedValue<number | undefined>(undefined)
 
   useAnimatedReaction(