fix: keep detecting mode changes the class filter cannot see

Watch childList on body and class on the html element, so element moves and
modes applied above body are still detected. Keep isMotionDisabled reading the
live DOM, since it is public and may be called while the flush cache is
populated.

Author: TrevorBurnham

src/internal/visual-mode/__tests__/use-visual-mode-performance.test.tsx

@@ -1,8 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import React, { useRef } from 'react';
-import { useCurrentMode, useDensityMode } from '../index';
+import React, { useLayoutEffect, useRef } from 'react';
+import { isMotionDisabled, useCurrentMode, useDensityMode, useReducedMotion } from '../index';
 import { render, screen } from '@testing-library/react';
 import { mutate } from './utils';
 
@@ -44,6 +44,11 @@ function spyOnClassListReads() {
 
 afterEach(() => {
   jest.restoreAllMocks();
+  // Reset shared state explicitly rather than in each test, so that one failing assertion
+  // cannot leave a mode class behind and cascade into unrelated failures.
+  document.documentElement.className = '';
+  document.body.className = '';
+  document.body.replaceChildren();
 });
 
 describe('mode detection cost', () => {
@@ -59,8 +64,9 @@ describe('mode detection cost', () => {
   });
 
   test('shares ancestor lookups between subscribers within a single flush', async () => {
+    const depth = 20;
     async function countReadsForOneFlush(detectorCount: number) {
-      const { container, unmount } = render(<ManyDetectors count={detectorCount} depth={20} />);
+      const { container, unmount } = render(<ManyDetectors count={detectorCount} depth={depth} />);
       const reads = spyOnClassListReads();
       await mutate(() => container.classList.add('unrelated-class'));
       const total = reads.mock.calls.length;
@@ -69,14 +75,16 @@ describe('mode detection cost', () => {
       return total;
     }
 
-    const withTen = await countReadsForOneFlush(10);
     const withForty = await countReadsForOneFlush(40);
-    expect(withTen).toBeGreaterThan(0);
-
-    // Each subscriber walks to the document root, so without memoization 4x the subscribers
-    // costs ~4x the class reads. Sharing the resolved path within a flush means only each
-    // subscriber's own element is uncached, making the growth far sublinear in chain depth.
-    expect(withForty).toBeLessThan(withTen * 2);
+    const withEighty = await countReadsForOneFlush(80);
+
+    // Assert the marginal cost of one more subscriber, which is what memoization bounds.
+    // Each additional subscriber should only read its own element's classList, once per
+    // detected mode, because its ancestors were already resolved by an earlier subscriber.
+    // Without memoization each one re-walks the shared chain instead, so the marginal cost
+    // would scale with `depth`.
+    const marginalReadsPerSubscriber = (withEighty - withForty) / 40;
+    expect(marginalReadsPerSubscriber).toBeLessThan(depth / 2);
   });
 
   test('does not reuse cached lookups across separate flushes', async () => {
@@ -111,6 +119,72 @@ describe('mode detection cost', () => {
     expect(screen.getByTestId('in-neither')).toHaveTextContent('light-comfortable');
   });
 
+  test('detects a mode class applied above body', async () => {
+    render(<ModeRender testId="detector" />);
+    expect(screen.getByTestId('detector')).toHaveTextContent('light-comfortable');
+
+    await mutate(() => document.documentElement.classList.add('awsui-dark-mode'));
+    expect(screen.getByTestId('detector')).toHaveTextContent('dark-comfortable');
+
+    await mutate(() => document.documentElement.classList.remove('awsui-dark-mode'));
+    expect(screen.getByTestId('detector')).toHaveTextContent('light-comfortable');
+  });
+
+  test('detects a move into a subtree with a different mode', async () => {
+    const host = document.createElement('div');
+    document.body.appendChild(host);
+    const darkSubtree = document.createElement('div');
+    darkSubtree.className = 'awsui-dark-mode';
+    document.body.appendChild(darkSubtree);
+
+    render(<ModeRender testId="detector" />, { container: host });
+    expect(screen.getByTestId('detector')).toHaveTextContent('light-comfortable');
+
+    // No class changes here: only the ancestor chain does.
+    await mutate(() => darkSubtree.appendChild(host));
+    expect(screen.getByTestId('detector')).toHaveTextContent('dark-comfortable');
+
+    await mutate(() => document.body.appendChild(host));
+    expect(screen.getByTestId('detector')).toHaveTextContent('light-comfortable');
+  });
+
+  test('isMotionDisabled reads the live DOM when called during a fan-out', async () => {
+    const wrapper = document.createElement('div');
+    document.body.appendChild(wrapper);
+    let observedDuringFlush: boolean | undefined = undefined;
+
+    function MotionSubscriber() {
+      const ref = useRef(null);
+      useReducedMotion(ref);
+      return <div ref={ref} />;
+    }
+
+    // This effect runs inside the fan-out, after MotionSubscriber has already resolved (and
+    // cached) the motion lookup for the shared ancestor chain.
+    function MutatingDuringFlush() {
+      const ref = useRef<HTMLDivElement>(null);
+      const colorMode = useCurrentMode(ref);
+      useLayoutEffect(() => {
+        if (colorMode === 'dark' && ref.current) {
+          wrapper.classList.add('awsui-motion-disabled');
+          observedDuringFlush = isMotionDisabled(ref.current);
+        }
+      }, [colorMode]);
+      return <div ref={ref} />;
+    }
+
+    render(
+      <>
+        <MotionSubscriber />
+        <MutatingDuringFlush />
+      </>,
+      { container: wrapper }
+    );
+    await mutate(() => wrapper.classList.add('awsui-dark-mode'));
+
+    expect(observedDuringFlush).toBe(true);
+  });
+
   test('detects a mode applied to an intermediate ancestor rather than body', async () => {
     const { container } = render(
       <div className="outer">

src/internal/visual-mode/index.ts

@@ -13,8 +13,14 @@ import { safeMatchMedia } from '../utils/safe-match-media.js';
 
 /**
  * Ancestor-chain lookups resolved during the current mutation flush, keyed by mode and then
- * by element. Only populated while the singleton observer is fanning out (see
- * `useMutationSingleton`); `null` at all other times, so a stale entry can never be read.
+ * by element. Only populated while the singleton observer is fanning out to its subscribers
+ * (see `useMutationSingleton`), and `null` at all other times.
+ *
+ * Only the detectors below read it. Callers outside the fan-out must not, because React can
+ * run effects while it is populated: an effect that mutates a mode class and then queries
+ * synchronously has to see the DOM as it is, not as the flush found it. The detectors
+ * themselves are unaffected, since they all run before any effect in the batch and any
+ * mutation an effect makes schedules a fresh flush.
  */
 let flushCache: null | Map<string, Map<HTMLElement, boolean>> = null;
 
@@ -73,7 +79,17 @@ function hasMatchingAncestor(mode: string, element: HTMLElement, test: (node: HT
   return resolved;
 }
 
+function hasMotionDisabledAncestor(element: HTMLElement): boolean {
+  return !!findUpUntil(element, node => node.classList.contains('awsui-motion-disabled'));
+}
+
+// Public API, callable at any time, so it always reads the live DOM rather than the flush cache.
 export function isMotionDisabled(element: HTMLElement): boolean {
+  return hasMotionDisabledAncestor(element) || safeMatchMedia(element, '(prefers-reduced-motion: reduce)');
+}
+
+// Equivalent to `isMotionDisabled`, but shares ancestor lookups across subscribers in a flush.
+function detectReducedMotion(element: HTMLElement): boolean {
   return (
     hasMatchingAncestor('motion', element, node => node.classList.contains('awsui-motion-disabled')) ||
     safeMatchMedia(element, '(prefers-reduced-motion: reduce)')
@@ -135,27 +151,39 @@ export function useDensityMode(elementRef: React.RefObject<HTMLElement>) {
 }
 
 export function useReducedMotion(elementRef: React.RefObject<HTMLElement>) {
-  return useModeDetector(elementRef, isMotionDisabled, false);
+  return useModeDetector(elementRef, detectReducedMotion, false);
 }
 
 const useMutationSingleton = createSingletonHandler<void>(handler => {
-  const observer = new MutationObserver(() => {
+  const fanOut = () => {
     // Memoize ancestor lookups for the duration of this fan-out only. Every subscriber runs
-    // synchronously inside handler(), so the cache is scoped to a single consistent view of
-    // the DOM and is discarded before any further mutations can be observed.
+    // synchronously inside handler(), before React processes any effect in the batch, so all
+    // of them see one consistent view of the DOM.
     flushCache = new Map();
     try {
       handler();
     } finally {
       flushCache = null;
     }
-  });
-  // Every mode is expressed as a class name, so class is the only attribute whose change can
-  // affect a detected mode. Filtering here avoids waking all subscribers for unrelated
-  // attribute changes anywhere in the document, such as the `data-awsui-focus-visible`
-  // toggle that focus-visible writes to `<body>` on every keydown and mousedown.
-  observer.observe(document.body, { attributes: true, subtree: true, attributeFilter: ['class'] });
-  return () => observer.disconnect();
+  };
+  const observer = new MutationObserver(fanOut);
+  const htmlObserver = new MutationObserver(fanOut);
+  // A mode is only ever expressed as a class name, so watching `class` is what detects a mode
+  // change. Filtering to it avoids waking every subscriber for unrelated attribute changes
+  // anywhere in the document, such as the `data-awsui-focus-visible` toggle that
+  // focus-visible writes to `<body>` on every keydown and mousedown.
+  //
+  // `childList` is watched as well because moving an element between subtrees changes its
+  // ancestor chain, and therefore its mode, without any class changing. Previously such
+  // moves were only picked up incidentally, by whatever unrelated attribute mutation
+  // happened to follow.
+  observer.observe(document.body, { attributes: true, subtree: true, childList: true, attributeFilter: ['class'] });
+  // Modes are also honoured above `<body>`, which the observer above does not cover.
+  htmlObserver.observe(document.documentElement, { attributes: true, attributeFilter: ['class'] });
+  return () => {
+    observer.disconnect();
+    htmlObserver.disconnect();
+  };
 });
 
 function useMutationObserver(elementRef: React.RefObject<HTMLElement>, onChange: (element: HTMLElement) => void) {