[ui] Fix issue where visibleChanged does not cause rerender (#13170)

@ch-brian noticed a weird issue where a TextInput in a Modal was losing
focus on the second Modal open. The issue was that visibleChanged was
not synced to the render cycle, causing an incorrect value read on the
next render, triggering the auto focus function of the Modal to the
default focus target at the wrong time.

GitOrigin-RevId: 2a8b49e741b97ee9f8183ed4553db31076acee5f

packages/ui/src/components/Modal.tsx

@@ -2,7 +2,13 @@
 import styled from "@emotion/styled";
 
 import type { ComponentProps, MutableRefObject, ReactNode } from "react";
-import React, { Fragment, useEffect, useLayoutEffect, useRef } from "react";
+import React, {
+  Fragment,
+  useEffect,
+  useLayoutEffect,
+  useRef,
+  useState,
+} from "react";
 import ReactDOM from "react-dom";
 import { useLiveRef } from "../hooks/useLiveRef.js";
 import { Breakpoints, bp, useBreakpoints } from "../styles/breakpoints.js";
@@ -133,6 +139,7 @@ export function Modal({
   topLeftIcon,
 }: ModalProps) {
   const visibleChangedRef = useRef(false);
+  const [visibleChanged, setVisibleChanged] = useState(false);
   const nodeRef = useRef<null | HTMLDivElement>(null);
   const [defaultFirstFocusRef, defaultFirstFocusRefCb] = useLiveRef();
   const ref = firstFocusRef || defaultFirstFocusRef;
@@ -146,9 +153,15 @@ export function Modal({
   useEffect(() => {
     if (visible !== visibleChangedRef.current) {
       visibleChangedRef.current = visible;
+      setVisibleChanged(true);
     }
   }, [visible]);
-  const visibleChanged = visible !== visibleChangedRef.current;
+
+  useEffect(() => {
+    if (visibleChanged) {
+      setVisibleChanged(false);
+    }
+  }, [visibleChanged]);
 
   useEffect(() => {
     prevFocusedElement.current = document.activeElement;

packages/ui/src/hooks/useLiveRef.tsx

@@ -1,12 +1,10 @@
-import type { MutableRefObject, RefCallback } from "react";
-import { useCallback, useRef, useState } from "react";
+import { useCallback, useMemo, useRef, useState } from "react";
+
+const defaultRef = { current: null };
 
 /* For use when you need a rerender when the ref is actually assigned.
    Usually ref assignment does not trigger rerenders. */
-export function useLiveRef<T = HTMLElement>(): [
-  MutableRefObject<T | null>,
-  RefCallback<T>,
-] {
+export function useLiveRef<T = HTMLElement>() {
   const [ready, setReady] = useState(false);
   const ref = useRef<T | null>(null);
   const refCb = useCallback((node: T | null) => {
@@ -16,5 +14,10 @@ export function useLiveRef<T = HTMLElement>(): [
     }
   }, []);
 
-  return [ready ? ref : { current: null }, refCb];
+  const value = useMemo(
+    () => [ready ? ref : defaultRef, refCb] as const,
+    [ready, ref, refCb],
+  );
+
+  return value;
 }