From 586aa95001c36a3210622da3e422dc7e3a786b00 Mon Sep 17 00:00:00 2001
From: Petyo Ivanov <underlog@gmail.com>
Date: Wed, 9 Aug 2023 11:50:54 +0300
Subject: [PATCH] fix: state restore works with window scrolling

Fixes  #957
---
 src/stateLoadSystem.ts | 82 ++++++++++++++++++++++++++----------------
 1 file changed, 52 insertions(+), 30 deletions(-)

diff --git a/src/stateLoadSystem.ts b/src/stateLoadSystem.ts
index ec96aadad..b49ef34a9 100644
--- a/src/stateLoadSystem.ts
+++ b/src/stateLoadSystem.ts
@@ -1,39 +1,61 @@
 import { domIOSystem } from './domIOSystem'
 import { initialTopMostItemIndexSystem } from './initialTopMostItemIndexSystem'
-import { StateSnapshot, StateCallback } from './interfaces'
+import { StateSnapshot, StateCallback, ScrollContainerState, WindowViewportInfo } from './interfaces'
 import { propsReadySystem } from './propsReadySystem'
 import { sizeSystem, sizeTreeToRanges } from './sizeSystem'
 import * as u from './urx'
+import { windowScrollerSystem } from './windowScrollerSystem'
 
-export const stateLoadSystem = u.system(([{ sizes, sizeRanges }, { scrollTop }, { initialTopMostItemIndex }, { didMount }]) => {
-  const getState = u.stream<StateCallback>()
-  const restoreStateFrom = u.statefulStream<StateSnapshot | undefined>(undefined)
-
-  u.subscribe(u.pipe(getState, u.withLatestFrom(sizes, scrollTop)), ([callback, sizes, scrollTop]) => {
-    const ranges = sizeTreeToRanges(sizes.sizeTree)
-    callback({ ranges, scrollTop })
-  })
-
-  u.connect(u.pipe(restoreStateFrom, u.filter(u.isDefined), u.map(locationFromSnapshot)), initialTopMostItemIndex)
-
-  u.connect(
-    u.pipe(
-      didMount,
-      u.withLatestFrom(restoreStateFrom),
-      u.filter(([, state]) => state !== undefined),
-      u.distinctUntilChanged(),
-      u.map(([, snapshot]) => {
-        return snapshot!.ranges
-      })
-    ),
-    sizeRanges
-  )
-
-  return {
-    getState,
-    restoreStateFrom,
-  }
-}, u.tup(sizeSystem, domIOSystem, initialTopMostItemIndexSystem, propsReadySystem))
+export const stateLoadSystem = u.system(
+  ([
+    { sizes, sizeRanges },
+    { scrollTop },
+    { initialTopMostItemIndex },
+    { didMount },
+    { useWindowScroll, windowScrollContainerState, windowViewportRect },
+  ]) => {
+    const getState = u.stream<StateCallback>()
+    const restoreStateFrom = u.statefulStream<StateSnapshot | undefined>(undefined)
+
+    const statefulWindowScrollContainerState = u.statefulStream<ScrollContainerState | null>(null)
+    const statefulWindowViewportRect = u.statefulStream<WindowViewportInfo | null>(null)
+
+    u.connect(windowScrollContainerState, statefulWindowScrollContainerState)
+    u.connect(windowViewportRect, statefulWindowViewportRect)
+
+    u.subscribe(
+      u.pipe(getState, u.withLatestFrom(sizes, scrollTop, useWindowScroll, statefulWindowScrollContainerState, statefulWindowViewportRect)),
+      ([callback, sizes, scrollTop, useWindowScroll, windowScrollContainerState, windowViewportRect]) => {
+        const ranges = sizeTreeToRanges(sizes.sizeTree)
+        if (useWindowScroll && windowScrollContainerState !== null && windowViewportRect !== null) {
+          scrollTop = windowScrollContainerState.scrollTop - windowViewportRect.offsetTop
+        }
+        callback({ ranges, scrollTop })
+      }
+    )
+
+    u.connect(u.pipe(restoreStateFrom, u.filter(u.isDefined), u.map(locationFromSnapshot)), initialTopMostItemIndex)
+
+    u.connect(
+      u.pipe(
+        didMount,
+        u.withLatestFrom(restoreStateFrom),
+        u.filter(([, state]) => state !== undefined),
+        u.distinctUntilChanged(),
+        u.map(([, snapshot]) => {
+          return snapshot!.ranges
+        })
+      ),
+      sizeRanges
+    )
+
+    return {
+      getState,
+      restoreStateFrom,
+    }
+  },
+  u.tup(sizeSystem, domIOSystem, initialTopMostItemIndexSystem, propsReadySystem, windowScrollerSystem)
+)
 
 function locationFromSnapshot(snapshot: StateSnapshot | undefined) {
   return { offset: snapshot!.scrollTop, index: 0, align: 'start' }