From acc96c12946e92f9284941f6132140dd8cc70f7a Mon Sep 17 00:00:00 2001
From: noa cohen <noaco@wix.com>
Date: Sun, 14 Jul 2024 16:01:03 +0300
Subject: [PATCH] wrapWithShouldUpdateExperimental

---
 src/connectWithShell.tsx | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/connectWithShell.tsx b/src/connectWithShell.tsx
index cc0067f..106cdfc 100644
--- a/src/connectWithShell.tsx
+++ b/src/connectWithShell.tsx
@@ -39,6 +39,31 @@ function wrapWithShouldUpdate<Props extends unknown, F extends (next: Props, pre
     return ((...args: Parameters<F>) => (shouldUpdate && !shouldUpdate(shell, getOwnProps()) ? true : func(args[0], args[1]))) as F
 }
 
+function wrapWithShouldUpdateExperimental<Props extends unknown, F extends (next: Props, prev: Props) => boolean>(
+    shouldUpdate: Maybe<(shell: Shell, ownProps?: Props) => boolean>,
+    arePropsEqual: F,
+    getOwnProps: () => Props,
+    shell: Shell
+): F {
+    if (!shouldUpdate) {
+        return arePropsEqual
+    }
+    let changeInPropsDetected = false
+    return ((...args: Parameters<F>) => {
+        if (!shouldUpdate(shell, getOwnProps())) {
+            if (!changeInPropsDetected) {
+                changeInPropsDetected = !arePropsEqual(args[0], args[1])
+            }
+            return true
+        }
+        if (changeInPropsDetected) {
+            changeInPropsDetected = false
+            return false
+        }
+        return arePropsEqual(args[0], args[1])
+    }) as F
+}
+
 function wrapWithShellContext<State, OwnProps, StateProps, DispatchProps>(
     component: React.ComponentType<OwnProps & StateProps & DispatchProps>,
     mapStateToProps: MapStateToProps<State, OwnProps, StateProps>,
@@ -98,7 +123,7 @@ function wrapWithShellContext<State, OwnProps, StateProps, DispatchProps>(
                               this.getOwnProps,
                               boundShell
                           ),
-                          areOwnPropsEqual: wrapWithShouldUpdate(
+                          areOwnPropsEqual: wrapWithShouldUpdateExperimental(
                               shouldComponentUpdate,
                               reduxConnectOptions.areOwnPropsEqual,
                               this.getOwnProps,