[Lab][Masonry] Fix layout flicker and single column issue

packages/mui-lab/src/Masonry/Masonry.js

@@ -208,118 +208,149 @@ const Masonry = React.forwardRef(function Masonry(inProps, ref) {
 
   const classes = useUtilityClasses(ownerState);
 
-  const handleResize = React.useCallback(
-    (masonryChildren) => {
-      if (!masonryRef.current || !masonryChildren || masonryChildren.length === 0) {
-        return;
-      }
+  const handleResize = React.useCallback(() => {
+    if (!masonryRef.current) {
+      return;
+    }
 
-      const masonry = masonryRef.current;
-      const masonryFirstChild = masonryRef.current.firstChild;
-      const parentWidth = masonry.clientWidth;
-      const firstChildWidth = masonryFirstChild.clientWidth;
+    const masonry = masonryRef.current;
+    const firstVisibleChild = Array.from(masonry.childNodes).find(
+      (child) =>
+        child.nodeType === Node.ELEMENT_NODE &&
+        child.dataset.class !== 'line-break' &&
+        window.getComputedStyle(child).display !== 'none',
+    );
 
-      if (parentWidth === 0 || firstChildWidth === 0) {
-        return;
-      }
+    if (!firstVisibleChild) {
+      return;
+    }
 
-      const firstChildComputedStyle = window.getComputedStyle(masonryFirstChild);
-      const firstChildMarginLeft = parseToNumber(firstChildComputedStyle.marginLeft);
-      const firstChildMarginRight = parseToNumber(firstChildComputedStyle.marginRight);
+    const parentWidth = masonry.clientWidth;
+    const firstChildWidth = firstVisibleChild.clientWidth;
 
-      const currentNumberOfColumns = Math.round(
-        parentWidth / (firstChildWidth + firstChildMarginLeft + firstChildMarginRight),
-      );
+    if (parentWidth === 0 || firstChildWidth === 0) {
+      return;
+    }
 
-      const columnHeights = new Array(currentNumberOfColumns).fill(0);
-      let skip = false;
-      let nextOrder = 1;
-      masonry.childNodes.forEach((child) => {
-        if (child.nodeType !== Node.ELEMENT_NODE || child.dataset.class === 'line-break' || skip) {
-          return;
-        }
-        const childComputedStyle = window.getComputedStyle(child);
-        const childMarginTop = parseToNumber(childComputedStyle.marginTop);
-        const childMarginBottom = parseToNumber(childComputedStyle.marginBottom);
-        // if any one of children isn't rendered yet, masonry's height shouldn't be computed yet
-        const childHeight = parseToNumber(childComputedStyle.height)
-          ? Math.ceil(parseToNumber(childComputedStyle.height)) + childMarginTop + childMarginBottom
-          : 0;
-        if (childHeight === 0) {
+    const firstChildComputedStyle = window.getComputedStyle(firstVisibleChild);
+    const firstChildMarginLeft = parseToNumber(firstChildComputedStyle.marginLeft);
+    const firstChildMarginRight = parseToNumber(firstChildComputedStyle.marginRight);
+
+    const currentNumberOfColumns = Math.round(
+      parentWidth / (firstChildWidth + firstChildMarginLeft + firstChildMarginRight),
+    );
+
+    const columnHeights = new Array(currentNumberOfColumns).fill(0);
+    let skip = false;
+    let nextOrder = 1;
+    masonry.childNodes.forEach((child) => {
+      if (child.nodeType !== Node.ELEMENT_NODE || child.dataset.class === 'line-break' || skip) {
+        return;
+      }
+      const childComputedStyle = window.getComputedStyle(child);
+      if (childComputedStyle.display === 'none') {
+        return;
+      }
+      const childMarginTop = parseToNumber(childComputedStyle.marginTop);
+      const childMarginBottom = parseToNumber(childComputedStyle.marginBottom);
+      const childHeight = parseToNumber(childComputedStyle.height)
+        ? Math.ceil(parseToNumber(childComputedStyle.height)) + childMarginTop + childMarginBottom
+        : 0;
+      if (childHeight === 0) {
+        skip = true;
+        return;
+      }
+      for (let i = 0; i < child.childNodes.length; i += 1) {
+        const nestedChild = child.childNodes[i];
+        if (nestedChild.tagName === 'IMG' && nestedChild.clientHeight === 0) {
           skip = true;
-          return;
+          break;
         }
-        // if there is a nested image that isn't rendered yet, masonry's height shouldn't be computed yet
-        for (let i = 0; i < child.childNodes.length; i += 1) {
-          const nestedChild = child.childNodes[i];
-          if (nestedChild.tagName === 'IMG' && nestedChild.clientHeight === 0) {
-            skip = true;
-            break;
+      }
+      if (!skip) {
+        if (sequential) {
+          columnHeights[nextOrder - 1] += childHeight;
+          child.style.order = nextOrder;
+          nextOrder += 1;
+          if (nextOrder > currentNumberOfColumns) {
+            nextOrder = 1;
           }
+        } else {
+          const currentMinColumnIndex = columnHeights.indexOf(Math.min(...columnHeights));
+          columnHeights[currentMinColumnIndex] += childHeight;
+          const order = currentMinColumnIndex + 1;
+          child.style.order = order;
         }
-        if (!skip) {
-          if (sequential) {
-            columnHeights[nextOrder - 1] += childHeight;
-            child.style.order = nextOrder;
-            nextOrder += 1;
-            if (nextOrder > currentNumberOfColumns) {
-              nextOrder = 1;
-            }
-          } else {
-            // find the current shortest column (where the current item will be placed)
-            const currentMinColumnIndex = columnHeights.indexOf(Math.min(...columnHeights));
-            columnHeights[currentMinColumnIndex] += childHeight;
-            const order = currentMinColumnIndex + 1;
-            child.style.order = order;
-          }
+      }
+    });
+    if (!skip) {
+      queueMicrotask(() => {
+        if (masonryRef.current) {
+          ReactDOM.flushSync(() => {
+            setMaxColumnHeight(Math.max(...columnHeights));
+            setNumberOfLineBreaks(currentNumberOfColumns > 0 ? currentNumberOfColumns - 1 : 0);
+          });
         }
       });
-      if (!skip) {
-        // 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
-        ReactDOM.flushSync(() => {
-          setMaxColumnHeight(Math.max(...columnHeights));
-          setNumberOfLineBreaks(currentNumberOfColumns > 0 ? currentNumberOfColumns - 1 : 0);
-        });
-      }
-    },
-    [sequential],
-  );
+    }
+  }, [sequential]);
 
   useEnhancedEffect(() => {
-    // IE and old browsers are not supported
-    if (typeof ResizeObserver === 'undefined') {
+    if (typeof ResizeObserver === 'undefined' || typeof MutationObserver === 'undefined') {
+      return undefined;
+    }
+
+    const masonry = masonryRef.current;
+    if (!masonry) {
       return undefined;
     }
 
-    let animationFrame;
+    let resizeTimeout;
+    const debouncedHandleResize = () => {
+      clearTimeout(resizeTimeout);
+      resizeTimeout = setTimeout(handleResize, 16); // ~60fps
+    };
 
-    const resizeObserver = new ResizeObserver(() => {
-      // see https://github.com/mui/material-ui/issues/36909
-      animationFrame = requestAnimationFrame(handleResize);
+    const resizeObserver = new ResizeObserver(debouncedHandleResize);
+    // Observes for child additions or removals to update the layout.
+    const mutationObserver = new MutationObserver((mutations) => {
+      mutations.forEach((mutation) => {
+        mutation.addedNodes.forEach((node) => {
+          if (node instanceof HTMLElement && node.dataset.class !== 'line-break') {
+            resizeObserver.observe(node);
+          }
+        });
+        mutation.removedNodes.forEach((node) => {
+          if (node instanceof HTMLElement && node.dataset.class !== 'line-break') {
+            resizeObserver.unobserve(node);
+          }
+        });
+      });
+      handleResize();
     });
 
-    if (masonryRef.current) {
-      masonryRef.current.childNodes.forEach((childNode) => {
+    Array.from(masonry.childNodes).forEach((childNode) => {
+      if (childNode instanceof HTMLElement && childNode.dataset.class !== 'line-break') {
         resizeObserver.observe(childNode);
-      });
-    }
+      }
+    });
+
+    mutationObserver.observe(masonry, {
+      childList: true,
+    });
+
+    handleResize();
 
     return () => {
-      if (animationFrame) {
-        cancelAnimationFrame(animationFrame);
-      }
-      if (resizeObserver) {
-        resizeObserver.disconnect();
-      }
+      clearTimeout(resizeTimeout);
+      resizeObserver.disconnect();
+      mutationObserver.disconnect();
     };
-  }, [columns, spacing, children, handleResize]);
+  }, [handleResize, columns, spacing, children]);
 
   const handleRef = useForkRef(ref, masonryRef);
 
-  //  columns are likely to have different heights and hence can start to merge;
-  //  a line break at the end of each column prevents columns from merging
+  // A line break is added to the end of each column to prevent columns from merging.
   const lineBreaks = new Array(numberOfLineBreaks)
     .fill('')
     .map((_, index) => (