Make layout reactivity less brittle

Follow-up to ea2d98179c

This took a couple of iterations to find something that works without creating update loops, but I think that by automatically informing Grid whenever a layout component is re-rendered, we'll have a much easier time ensuring that our layouts are fully reactive.

src/grid/Grid.tsx

@@ -25,10 +25,15 @@ import {
   CSSProperties,
   ComponentProps,
   ComponentType,
+  Dispatch,
   FC,
   LegacyRef,
   ReactNode,
-  useCallback,
+  SetStateAction,
+  createContext,
+  forwardRef,
+  memo,
+  useContext,
   useEffect,
   useMemo,
   useRef,
@@ -113,6 +118,27 @@ function offset(element: HTMLElement, relativeTo: Element): Offset {
   }
 }
 
+interface LayoutContext {
+  setGeneration: Dispatch<SetStateAction<number | null>>;
+}
+
+const LayoutContext = createContext<LayoutContext | null>(null);
+
+/**
+ * Enables Grid to react to layout changes. You must call this in your Layout
+ * component or else Grid will not be reactive.
+ */
+export function useLayout(): void {
+  const context = useContext(LayoutContext);
+  if (context === null)
+    throw new Error("useLayout called outside of a Grid layout component");
+
+  // On every render, tell Grid that the layout may have changed
+  useEffect(() =>
+    context.setGeneration((prev) => (prev === null ? 0 : prev + 1)),
+  );
+}
+
 export interface LayoutProps<LayoutModel, TileModel, R extends HTMLElement> {
   ref: LegacyRef<R>;
   model: LayoutModel;
@@ -158,6 +184,11 @@ interface Drag {
 
 export type DragCallback = (drag: Drag) => void;
 
+interface LayoutMemoProps<LayoutModel, TileModel, R extends HTMLElement>
+  extends LayoutProps<LayoutModel, TileModel, R> {
+  Layout: ComponentType<LayoutProps<LayoutModel, TileModel, R>>;
+}
+
 interface Props<
   LayoutModel,
   TileModel,
@@ -209,7 +240,7 @@ export function Grid<
   const [gridRoot, gridRef2] = useState<HTMLElement | null>(null);
   const gridRef = useMergedRefs<HTMLElement>(gridRef1, gridRef2);
 
-  const [layoutRoot, setLayoutRoot] = useState<HTMLElement | null>(null);
+  const [layoutRoot, layoutRef] = useState<HTMLElement | null>(null);
   const [generation, setGeneration] = useState<number | null>(null);
   const tiles = useInitial(() => new Map<string, Tile<TileModel>>());
   const prefersReducedMotion = usePrefersReducedMotion();
@@ -236,27 +267,22 @@ export function Grid<
     [tiles],
   );
 
-  const layoutRef = useCallback(
-    (e: HTMLElement | null) => {
-      setLayoutRoot(e);
-      if (e !== null)
-        setGeneration(parseInt(e.getAttribute("data-generation")!));
-    },
-    [setLayoutRoot, setGeneration],
+  // We must memoize the Layout component to break the update loop where a
+  // render of Grid causes a re-render of Layout, which in turn re-renders Grid
+  const LayoutMemo = useMemo(
+    () =>
+      memo(
+        forwardRef<
+          LayoutRef,
+          LayoutMemoProps<LayoutModel, TileModel, LayoutRef>
+        >(function LayoutMemo({ Layout, ...props }, ref): ReactNode {
+          return <Layout {...props} ref={ref} />;
+        }),
+      ),
+    [],
   );
 
-  useEffect(() => {
-    if (layoutRoot !== null) {
-      const observer = new MutationObserver((mutations) => {
-        if (mutations.some((m) => m.type === "attributes")) {
-          setGeneration(parseInt(layoutRoot.getAttribute("data-generation")!));
-        }
-      });
-
-      observer.observe(layoutRoot, { attributes: true });
-      return (): void => observer.disconnect();
-    }
-  }, [layoutRoot, setGeneration]);
+  const context: LayoutContext = useMemo(() => ({ setGeneration }), []);
 
   // Combine the tile definitions and slots together to create placed tiles
   const placedTiles = useMemo(() => {
@@ -279,10 +305,10 @@ export function Grid<
     }
 
     return result;
-    // The rects may change due to the grid updating to a new generation, but
-    // eslint can't statically verify this
+    // The rects may change due to the grid resizing or updating to a new
+    // generation, but eslint can't statically verify this
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [gridRoot, layoutRoot, tiles, generation]);
+  }, [gridRoot, layoutRoot, tiles, gridBounds, generation]);
 
   // Drag state is stored in a ref rather than component state, because we use
   // react-spring's imperative API during gestures to improve responsiveness
@@ -463,7 +489,9 @@ export function Grid<
       className={classNames(className, styles.grid)}
       style={style}
     >
-      <Layout ref={layoutRef} model={model} Slot={Slot} />
+      <LayoutContext.Provider value={context}>
+        <LayoutMemo ref={layoutRef} Layout={Layout} model={model} Slot={Slot} />
+      </LayoutContext.Provider>
       {tileTransitions((spring, { id, model, onDrag, width, height }) => (
         <TileWrapper
           key={id}

src/grid/GridLayout.tsx

@@ -20,7 +20,6 @@ import { useObservableEagerState } from "observable-hooks";
 
 import { GridLayout as GridLayoutModel } from "../state/CallViewModel";
 import styles from "./GridLayout.module.css";
-import { useReactiveState } from "../useReactiveState";
 import { useInitial } from "../useInitial";
 import {
   CallLayout,
@@ -28,7 +27,7 @@ import {
   TileModel,
   arrangeTiles,
 } from "./CallLayout";
-import { DragCallback } from "./Grid";
+import { DragCallback, useLayout } from "./Grid";
 
 interface GridCSSProperties extends CSSProperties {
   "--gap": string;
@@ -49,7 +48,7 @@ export const makeGridLayout: CallLayout<GridLayoutModel> = ({
   // The "fixed" (non-scrolling) part of the layout is where the spotlight tile
   // lives
   fixed: forwardRef(function GridLayoutFixed({ model, Slot }, ref) {
-    const { width, height } = useObservableEagerState(minBounds);
+    useLayout();
     const alignment = useObservableEagerState(
       useInitial(() =>
         spotlightAlignment.pipe(
@@ -68,10 +67,6 @@ export const makeGridLayout: CallLayout<GridLayoutModel> = ({
         },
       [model.spotlight],
     );
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [model.spotlight === undefined, width, height, alignment],
-    );
 
     const onDragSpotlight: DragCallback = useCallback(
       ({ xRatio, yRatio }) =>
@@ -83,7 +78,7 @@ export const makeGridLayout: CallLayout<GridLayoutModel> = ({
     );
 
     return (
-      <div ref={ref} className={styles.fixed} data-generation={generation}>
+      <div ref={ref} className={styles.fixed}>
         {tileModel && (
           <Slot
             className={styles.slot}
@@ -100,17 +95,13 @@ export const makeGridLayout: CallLayout<GridLayoutModel> = ({
 
   // The scrolling part of the layout is where all the grid tiles live
   scrolling: forwardRef(function GridLayout({ model, Slot }, ref) {
+    useLayout();
     const { width, height: minHeight } = useObservableEagerState(minBounds);
     const { gap, tileWidth, tileHeight } = useMemo(
       () => arrangeTiles(width, minHeight, model.grid.length),
       [width, minHeight, model.grid.length],
     );
 
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [model.grid, width, minHeight],
-    );
-
     const tileModels: GridTileModel[] = useMemo(
       () => model.grid.map((vm) => ({ type: "grid", vm })),
       [model.grid],
@@ -119,7 +110,6 @@ export const makeGridLayout: CallLayout<GridLayoutModel> = ({
     return (
       <div
         ref={ref}
-        data-generation={generation}
         className={styles.scrolling}
         style={
           {

src/grid/OneOnOneLayout.tsx

@@ -20,9 +20,8 @@ import classNames from "classnames";
 
 import { OneOnOneLayout as OneOnOneLayoutModel } from "../state/CallViewModel";
 import { CallLayout, GridTileModel, arrangeTiles } from "./CallLayout";
-import { useReactiveState } from "../useReactiveState";
 import styles from "./OneOnOneLayout.module.css";
-import { DragCallback } from "./Grid";
+import { DragCallback, useLayout } from "./Grid";
 
 /**
  * An implementation of the "one-on-one" layout, in which the remote participant
@@ -35,10 +34,12 @@ export const makeOneOnOneLayout: CallLayout<OneOnOneLayoutModel> = ({
   scrollingOnTop: false,
 
   fixed: forwardRef(function OneOnOneLayoutFixed(_props, ref) {
-    return <div ref={ref} data-generation={0} />;
+    useLayout();
+    return <div ref={ref} />;
   }),
 
   scrolling: forwardRef(function OneOnOneLayoutScrolling({ model, Slot }, ref) {
+    useLayout();
     const { width, height } = useObservableEagerState(minBounds);
     const pipAlignmentValue = useObservableEagerState(pipAlignment);
     const { tileWidth, tileHeight } = useMemo(
@@ -46,11 +47,6 @@ export const makeOneOnOneLayout: CallLayout<OneOnOneLayoutModel> = ({
       [width, height],
     );
 
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [width, height, pipAlignmentValue],
-    );
-
     const remoteTileModel: GridTileModel = useMemo(
       () => ({ type: "grid", vm: model.remote }),
       [model.remote],
@@ -70,7 +66,7 @@ export const makeOneOnOneLayout: CallLayout<OneOnOneLayoutModel> = ({
     );
 
     return (
-      <div ref={ref} data-generation={generation} className={styles.layer}>
+      <div ref={ref} className={styles.layer}>
         <Slot
           id={remoteTileModel.vm.id}
           model={remoteTileModel}

src/grid/SpotlightExpandedLayout.tsx

@@ -19,9 +19,8 @@ import { useObservableEagerState } from "observable-hooks";
 
 import { SpotlightExpandedLayout as SpotlightExpandedLayoutModel } from "../state/CallViewModel";
 import { CallLayout, GridTileModel, SpotlightTileModel } from "./CallLayout";
-import { DragCallback } from "./Grid";
+import { DragCallback, useLayout } from "./Grid";
 import styles from "./SpotlightExpandedLayout.module.css";
-import { useReactiveState } from "../useReactiveState";
 
 /**
  * An implementation of the "expanded spotlight" layout, in which the spotlight
@@ -29,27 +28,21 @@ import { useReactiveState } from "../useReactiveState";
  */
 export const makeSpotlightExpandedLayout: CallLayout<
   SpotlightExpandedLayoutModel
-> = ({ minBounds, pipAlignment }) => ({
+> = ({ pipAlignment }) => ({
   scrollingOnTop: true,
 
   fixed: forwardRef(function SpotlightExpandedLayoutFixed(
     { model, Slot },
     ref,
   ) {
-    const { width, height } = useObservableEagerState(minBounds);
-
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [width, height, model.spotlight],
-    );
-
+    useLayout();
     const spotlightTileModel: SpotlightTileModel = useMemo(
       () => ({ type: "spotlight", vms: model.spotlight, maximised: true }),
       [model.spotlight],
     );
 
     return (
-      <div ref={ref} data-generation={generation} className={styles.layer}>
+      <div ref={ref} className={styles.layer}>
         <Slot
           className={styles.spotlight}
           id="spotlight"
@@ -63,14 +56,9 @@ export const makeSpotlightExpandedLayout: CallLayout<
     { model, Slot },
     ref,
   ) {
-    const { width, height } = useObservableEagerState(minBounds);
+    useLayout();
     const pipAlignmentValue = useObservableEagerState(pipAlignment);
 
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [width, height, model.pip === undefined, pipAlignmentValue],
-    );
-
     const pipTileModel: GridTileModel | undefined = useMemo(
       () => model.pip && { type: "grid", vm: model.pip },
       [model.pip],
@@ -86,7 +74,7 @@ export const makeSpotlightExpandedLayout: CallLayout<
     );
 
     return (
-      <div ref={ref} data-generation={generation} className={styles.layer}>
+      <div ref={ref} className={styles.layer}>
         {pipTileModel && (
           <Slot
             className={styles.pip}

src/grid/SpotlightLandscapeLayout.tsx

@@ -21,7 +21,7 @@ import classNames from "classnames";
 import { CallLayout, GridTileModel, TileModel } from "./CallLayout";
 import { SpotlightLandscapeLayout as SpotlightLandscapeLayoutModel } from "../state/CallViewModel";
 import styles from "./SpotlightLandscapeLayout.module.css";
-import { useReactiveState } from "../useReactiveState";
+import { useLayout } from "./Grid";
 
 /**
  * An implementation of the "spotlight landscape" layout, in which the spotlight
@@ -37,7 +37,8 @@ export const makeSpotlightLandscapeLayout: CallLayout<
     { model, Slot },
     ref,
   ) {
-    const { width, height } = useObservableEagerState(minBounds);
+    useLayout();
+    useObservableEagerState(minBounds);
     const tileModel: TileModel = useMemo(
       () => ({
         type: "spotlight",
@@ -46,13 +47,9 @@ export const makeSpotlightLandscapeLayout: CallLayout<
       }),
       [model.spotlight],
     );
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [model.grid.length, width, height, model.spotlight],
-    );
 
     return (
-      <div ref={ref} data-generation={generation} className={styles.layer}>
+      <div ref={ref} className={styles.layer}>
         <div className={styles.spotlight}>
           <Slot className={styles.slot} id="spotlight" model={tileModel} />
         </div>
@@ -65,18 +62,15 @@ export const makeSpotlightLandscapeLayout: CallLayout<
     { model, Slot },
     ref,
   ) {
-    const { width, height } = useObservableEagerState(minBounds);
+    useLayout();
+    useObservableEagerState(minBounds);
     const tileModels: GridTileModel[] = useMemo(
       () => model.grid.map((vm) => ({ type: "grid", vm })),
       [model.grid],
     );
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [model.spotlight.length, model.grid, width, height],
-    );
 
     return (
-      <div ref={ref} data-generation={generation} className={styles.layer}>
+      <div ref={ref} className={styles.layer}>
         <div
           className={classNames(styles.spotlight, {
             [styles.withIndicators]: model.spotlight.length > 1,

src/grid/SpotlightPortraitLayout.tsx

@@ -26,7 +26,7 @@ import {
 } from "./CallLayout";
 import { SpotlightPortraitLayout as SpotlightPortraitLayoutModel } from "../state/CallViewModel";
 import styles from "./SpotlightPortraitLayout.module.css";
-import { useReactiveState } from "../useReactiveState";
+import { useLayout } from "./Grid";
 
 interface GridCSSProperties extends CSSProperties {
   "--grid-gap": string;
@@ -48,7 +48,7 @@ export const makeSpotlightPortraitLayout: CallLayout<
     { model, Slot },
     ref,
   ) {
-    const { width, height } = useObservableEagerState(minBounds);
+    useLayout();
     const tileModel: TileModel = useMemo(
       () => ({
         type: "spotlight",
@@ -57,13 +57,9 @@ export const makeSpotlightPortraitLayout: CallLayout<
       }),
       [model.spotlight],
     );
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [model.grid.length, width, height, model.spotlight],
-    );
 
     return (
-      <div ref={ref} data-generation={generation} className={styles.layer}>
+      <div ref={ref} className={styles.layer}>
         <div className={styles.spotlight}>
           <Slot className={styles.slot} id="spotlight" model={tileModel} />
         </div>
@@ -75,7 +71,8 @@ export const makeSpotlightPortraitLayout: CallLayout<
     { model, Slot },
     ref,
   ) {
-    const { width, height } = useObservableEagerState(minBounds);
+    useLayout();
+    const { width } = useObservableEagerState(minBounds);
     const { gap, tileWidth, tileHeight } = arrangeTiles(
       width,
       0,
@@ -85,15 +82,10 @@ export const makeSpotlightPortraitLayout: CallLayout<
       () => model.grid.map((vm) => ({ type: "grid", vm })),
       [model.grid],
     );
-    const [generation] = useReactiveState<number>(
-      (prev) => (prev === undefined ? 0 : prev + 1),
-      [model.spotlight.length, model.grid, width, height],
-    );
 
     return (
       <div
         ref={ref}
-        data-generation={generation}
         className={styles.layer}
         style={
           {