From 3e5c6868a049aa0ab32fc1caaea83b5ddfbf4d8a Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@me.com>
Date: Fri, 1 Oct 2021 13:37:42 +0100
Subject: [PATCH] Flush styles synchronously

---
 package.json |  7 ++++---
 src/index.ts | 23 +++++++++++++----------
 yarn.lock    | 24 +++++++++++++++++++-----
 3 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/package.json b/package.json
index ecf78bf..88d28f2 100644
--- a/package.json
+++ b/package.json
@@ -22,7 +22,8 @@
     "release": "np --no-2fa"
   },
   "peerDependencies": {
-    "react": ">=16.8"
+    "react": ">=16.8",
+    "react-dom": ">=16.8"
   },
   "husky": {
     "hooks": {
@@ -52,8 +53,8 @@
     "@testing-library/react": "^10.0.2",
     "@types/jest": "^25.1.2",
     "@types/raf": "^3.4.0",
-    "@types/react": "^16.9.19",
-    "@types/react-dom": "^16.9.5",
+    "@types/react": "^17.0.3",
+    "@types/react-dom": "^17.0.3",
     "@types/styled-components": "^5.0.1",
     "babel-loader": "^8.0.6",
     "cypress": "^4.2.0",
diff --git a/src/index.ts b/src/index.ts
index 78f0c36..538d5d2 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,10 +1,5 @@
-import {
-  useState,
-  useRef,
-  useCallback,
-  TransitionEvent,
-  CSSProperties,
-} from 'react';
+import { useState, useRef, TransitionEvent, CSSProperties } from 'react';
+import { flushSync } from 'react-dom';
 import {
   noop,
   callAll,
@@ -54,12 +49,20 @@ export default function useCollapse({
     height: collapsedHeight,
     overflow: 'hidden',
   };
-  const [styles, setStyles] = useState<CSSProperties>(
+  const [styles, setStylesRaw] = useState<CSSProperties>(
     isExpanded ? {} : collapsedStyles
   );
-  const mergeStyles = useCallback((newStyles: {}): void => {
+  const setStyles = (newStyles: {} | ((oldStyles: {}) => {})): void => {
+    // We rely on reading information from layout
+    // at arbitrary times, so ensure all style changes
+    // happen before we might attempt to read them.
+    flushSync(() => {
+      setStylesRaw(newStyles);
+    });
+  };
+  const mergeStyles = (newStyles: {}): void => {
     setStyles((oldStyles) => ({ ...oldStyles, ...newStyles }));
-  }, []);
+  };
 
   function getTransitionStyles(
     height: number | string
diff --git a/yarn.lock b/yarn.lock
index 0150ba2..11f33e3 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2595,10 +2595,10 @@
     "@types/history" "*"
     "@types/react" "*"
 
-"@types/react-dom@^16.9.5":
-  version "16.9.8"
-  resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-16.9.8.tgz#fe4c1e11dfc67155733dfa6aa65108b4971cb423"
-  integrity sha512-ykkPQ+5nFknnlU6lDd947WbQ6TE3NNzbQAkInC2EKY1qeYdTKp7onFusmYZb+ityzx2YviqT6BXSu+LyWWJwcA==
+"@types/react-dom@^17.0.3":
+  version "17.0.9"
+  resolved "https://registry.yarnpkg.com/@types/react-dom/-/react-dom-17.0.9.tgz#441a981da9d7be117042e1a6fd3dac4b30f55add"
+  integrity sha512-wIvGxLfgpVDSAMH5utdL9Ngm5Owu0VsGmldro3ORLXV8CShrL8awVj06NuEXFQ5xyaYfdca7Sgbk/50Ri1GdPg==
   dependencies:
     "@types/react" "*"
 
@@ -2623,7 +2623,7 @@
   dependencies:
     "@types/react" "*"
 
-"@types/react@*", "@types/react@^16.9.19":
+"@types/react@*":
   version "16.9.46"
   resolved "https://registry.yarnpkg.com/@types/react/-/react-16.9.46.tgz#f0326cd7adceda74148baa9bff6e918632f5069e"
   integrity sha512-dbHzO3aAq1lB3jRQuNpuZ/mnu+CdD3H0WVaaBQA8LTT3S33xhVBUj232T8M3tAhSWJs/D/UqORYUlJNl/8VQZg==
@@ -2631,6 +2631,15 @@
     "@types/prop-types" "*"
     csstype "^3.0.2"
 
+"@types/react@^17.0.3":
+  version "17.0.26"
+  resolved "https://registry.yarnpkg.com/@types/react/-/react-17.0.26.tgz#960ea4b3518cc154ed7df3b35656c508df653331"
+  integrity sha512-MXxuXrH2xOcv5cp/su4oz69dNQnSA90JjFw5HBd5wifw6Ihi94j7dRJm7qNsB30tnruXSCPc9qmlhGop4nh9Hw==
+  dependencies:
+    "@types/prop-types" "*"
+    "@types/scheduler" "*"
+    csstype "^3.0.2"
+
 "@types/resolve@0.0.8":
   version "0.0.8"
   resolved "https://registry.yarnpkg.com/@types/resolve/-/resolve-0.0.8.tgz#f26074d238e02659e323ce1a13d041eee280e194"
@@ -2653,6 +2662,11 @@
     "@types/glob" "*"
     "@types/node" "*"
 
+"@types/scheduler@*":
+  version "0.16.2"
+  resolved "https://registry.yarnpkg.com/@types/scheduler/-/scheduler-0.16.2.tgz#1a62f89525723dde24ba1b01b092bf5df8ad4d39"
+  integrity sha512-hppQEBDmlwhFAXKJX2KnWLYu5yMfi91yazPb2l+lbJiwW+wdo1gNeRA+3RgNSO39WYX2euey41KEwnqesU2Jew==
+
 "@types/shelljs@^0.8.5":
   version "0.8.8"
   resolved "https://registry.yarnpkg.com/@types/shelljs/-/shelljs-0.8.8.tgz#e439c69929b88a2c8123c1a55e09eb708315addf"