Refactor useRovingCellRef, fix test warnings (#2883)

* Refactor useRovingCellRef, fix test warnings

* optimize

* lazy focus listening

* fix type issue

src/HeaderCell.tsx

@@ -155,7 +155,7 @@ export default function HeaderCell<R, SR>({
   }
 
   function handleFocus(event: React.FocusEvent<HTMLDivElement>) {
-    onFocus(event);
+    onFocus?.(event);
     if (shouldFocusGrid) {
       // Select the first header cell if there is no selected cell
       selectCell(0);

src/hooks/useRovingCellRef.ts

@@ -1,39 +1,31 @@
-import { useRef, useState } from 'react';
-import { useLayoutEffect } from './useLayoutEffect';
+import { useCallback, useState } from 'react';
 
 // https://www.w3.org/TR/wai-aria-practices-1.1/#kbd_roving_tabindex
 export function useRovingCellRef(isSelected: boolean) {
-  const ref = useRef<HTMLDivElement>(null);
   // https://www.w3.org/TR/wai-aria-practices-1.1/#gridNav_focus
-  const isChildFocused = useRef(false);
-  const [, forceRender] = useState<unknown>({});
+  const [isChildFocused, setIsChildFocused] = useState(false);
 
-  useLayoutEffect(() => {
-    if (!isSelected) {
-      isChildFocused.current = false;
-      return;
-    }
+  if (isChildFocused && !isSelected) {
+    setIsChildFocused(false);
+  }
 
-    if (isChildFocused.current) {
-      // When the child is focused, we need to rerender
-      // the cell again so tabIndex is updated to -1
-      forceRender({});
-      return;
-    }
-    ref.current?.focus({ preventScroll: true });
-  }, [isSelected]);
+  const ref = useCallback((cell: HTMLDivElement | null) => {
+    if (cell === null || cell.contains(document.activeElement)) return;
+
+    cell.focus({ preventScroll: true });
+  }, []);
 
   function onFocus(event: React.FocusEvent<HTMLDivElement>) {
-    if (event.target !== ref.current) {
-      isChildFocused.current = true;
+    if (event.target !== event.currentTarget) {
+      setIsChildFocused(true);
     }
   }
 
-  const isFocused = isSelected && !isChildFocused.current;
+  const isFocused = isSelected && !isChildFocused;
 
   return {
-    ref,
+    ref: isSelected ? ref : undefined,
     tabIndex: isFocused ? 0 : -1,
-    onFocus
+    onFocus: isSelected ? onFocus : undefined
   };
 }

test/components.test.tsx

@@ -1,8 +1,8 @@
-import { StrictMode } from 'react';
+import { forwardRef, StrictMode } from 'react';
 import { render, screen } from '@testing-library/react';
 
 import DataGrid, { DataGridDefaultComponentsProvider, SelectColumn } from '../src';
-import type { Column, DataGridProps } from '../src';
+import type { Column, DataGridProps, CheckboxFormatterProps } from '../src';
 import { getRows, setup } from './utils';
 
 interface Row {
@@ -25,13 +25,16 @@ function GlobalNoRowsFallback() {
   return <div>Global no rows fallback</div>;
 }
 
-function Checkbox() {
-  return <div>Local checkbox</div>;
-}
+const Checkbox = forwardRef<HTMLDivElement, CheckboxFormatterProps>(function Checkbox(props, ref) {
+  return <div ref={ref}>Local checkbox</div>;
+});
 
-function GlobalCheckbox() {
-  return <div>Global checkbox</div>;
-}
+const GlobalCheckbox = forwardRef<HTMLDivElement, CheckboxFormatterProps>(function GlobalCheckbox(
+  props,
+  ref
+) {
+  return <div ref={ref}>Global checkbox</div>;
+});
 
 function setupProvider<R, SR, K extends React.Key>(props: DataGridProps<R, SR, K>) {
   return render(