From b4e80b602d5692f4212db26873d800671825636c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arno=C5=A1t=20Pleskot?= <arnostpleskot@gmail.com>
Date: Wed, 13 Sep 2023 15:57:50 +0200
Subject: [PATCH] feat: add memoization

---
 src/scene/scrollConstraints.ts | 104 ++++++++++++++++++++++-----------
 1 file changed, 71 insertions(+), 33 deletions(-)

diff --git a/src/scene/scrollConstraints.ts b/src/scene/scrollConstraints.ts
index 973d2b2f06..b88db8fe17 100644
--- a/src/scene/scrollConstraints.ts
+++ b/src/scene/scrollConstraints.ts
@@ -1,4 +1,5 @@
 import { AppState, ScrollConstraints } from "../types";
+import { isShallowEqual } from "../utils";
 import { getNormalizedZoom } from "./zoom";
 
 /**
@@ -301,6 +302,15 @@ const isViewportOutsideOfConstrainedArea = (state: AppState) => {
   );
 };
 
+let memoizedValues: {
+  previousState: Pick<
+    AppState,
+    "zoom" | "width" | "height" | "scrollConstraints"
+  >;
+  constraints: ReturnType<typeof calculateConstraints>;
+  constraintsWithoutOverscroll: ReturnType<typeof calculateConstraints>;
+} | null = null;
+
 /**
  * Constrains the AppState scroll values within the defined scroll constraints.
  *
@@ -323,47 +333,75 @@ export const constrainScrollState = (
   }
   const { scrollX, scrollY, width, height, scrollConstraints, zoom } = state;
 
+  const canUseMemoizedValues =
+    memoizedValues?.previousState.scrollConstraints && // can't use memoized values if there were no scrollConstraints in memoizedValues
+    memoizedValues && // there are memoized values
+    isShallowEqual(
+      // current scrollConstraints are the same as in memoizedValues
+      state.scrollConstraints,
+      memoizedValues.previousState.scrollConstraints!,
+    ) &&
+    isShallowEqual(
+      // current zoom and window dimensions are equal to those in memoizedValues
+      { zoom: zoom.value, width, height },
+      {
+        zoom: memoizedValues.previousState.zoom.value,
+        width: memoizedValues.previousState.width,
+        height: memoizedValues.previousState.height,
+      },
+    );
+
+  const constraints = canUseMemoizedValues
+    ? memoizedValues!.constraints
+    : calculateConstraints({
+        scrollConstraints,
+        width,
+        height,
+        zoom,
+        cursorButton: "down",
+      });
+
+  const constraintsWithoutOverscroll = canUseMemoizedValues
+    ? memoizedValues!.constraintsWithoutOverscroll
+    : calculateConstraints({
+        scrollConstraints,
+        width,
+        height,
+        zoom,
+        cursorButton: "up",
+      });
+
   const constrainedValues = constrainScrollValues({
-    ...calculateConstraints({
-      scrollConstraints,
-      width,
-      height,
-      zoom,
-      cursorButton: "down",
-    }),
+    ...constraints,
     scrollX,
     scrollY,
   });
 
-  const shouldAnimate = isViewportOutsideOfConstrainedArea(state);
+  const animateTo = constrainScrollValues({
+    ...constraintsWithoutOverscroll,
+    scrollX,
+    scrollY,
+  });
 
-  const animateTo = shouldAnimate
-    ? constrainScrollValues({
-        ...calculateConstraints({
-          scrollConstraints,
-          width,
-          height,
-          zoom,
-          cursorButton: "up",
-        }),
-        scrollX,
-        scrollY,
-      })
-    : null;
-
-  if (
-    constrainedValues.scrollX !== scrollX ||
-    constrainedValues.scrollY !== scrollY
-  ) {
-    return {
-      state: {
-        ...state,
-        ...constrainedValues,
+  if (!canUseMemoizedValues) {
+    memoizedValues = {
+      previousState: {
+        zoom: state.zoom,
+        width: state.width,
+        height: state.height,
+        scrollConstraints: state.scrollConstraints,
       },
-      shouldAnimate,
-      animateTo,
+      constraints,
+      constraintsWithoutOverscroll,
     };
   }
 
-  return { state, shouldAnimate, animateTo };
+  return {
+    state: {
+      ...state,
+      ...constrainedValues,
+    },
+    shouldAnimate: isViewportOutsideOfConstrainedArea(state),
+    animateTo,
+  };
 };