From 42fdc5e476ebb205b7a36f4183aa6a3ba39395c5 Mon Sep 17 00:00:00 2001
From: Marija Najdova <mnajdova@gmail.com>
Date: Wed, 15 Jun 2022 17:58:55 +0200
Subject: [PATCH] [Masonry] Fix flickering when used with React 18

---
 packages/mui-lab/src/Masonry/Masonry.js | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/packages/mui-lab/src/Masonry/Masonry.js b/packages/mui-lab/src/Masonry/Masonry.js
index 48ffe778a8a004..cb84179e7b6836 100644
--- a/packages/mui-lab/src/Masonry/Masonry.js
+++ b/packages/mui-lab/src/Masonry/Masonry.js
@@ -1,4 +1,5 @@
 import { unstable_composeClasses as composeClasses } from '@mui/base';
+import { flushSync } from 'react-dom';
 import { styled, useThemeProps } from '@mui/material/styles';
 import {
   createUnarySpacing,
@@ -245,9 +246,13 @@ const Masonry = React.forwardRef(function Masonry(inProps, ref) {
       }
     });
     if (!skip) {
-      setMaxColumnHeight(Math.max(...columnHeights));
-      const numOfLineBreaks = currentNumberOfColumns > 0 ? currentNumberOfColumns - 1 : 0;
-      setNumberOfLineBreaks(numOfLineBreaks);
+      // In React 18, state updates in a ResizeObserver's callback are happening after the paint which causes flickering
+      // when doing some visual updates in it. Using flushSync ensures that the dom will be painted after the states updates happen
+      // Related issue - https://github.com/facebook/react/issues/24331
+      flushSync(() => {
+        setMaxColumnHeight(Math.max(...columnHeights));
+        setNumberOfLineBreaks(currentNumberOfColumns > 0 ? currentNumberOfColumns - 1 : 0);
+      });
     }
   };