From 0e804f728c26f4efe1a2cc32884f2643ef9bc8c4 Mon Sep 17 00:00:00 2001
From: Nicolas Stepien <stepien.nicolas@gmail.com>
Date: Thu, 7 Apr 2022 10:28:54 +0100
Subject: [PATCH 1/4] Refactor useRovingCellRef, fix test warnings

---
 src/hooks/useRovingCellRef.ts | 37 +++++++++++++++--------------------
 test/components.test.tsx      | 19 ++++++++++--------
 2 files changed, 27 insertions(+), 29 deletions(-)

diff --git a/src/hooks/useRovingCellRef.ts b/src/hooks/useRovingCellRef.ts
index 00458f4667..e997f628c2 100644
--- a/src/hooks/useRovingCellRef.ts
+++ b/src/hooks/useRovingCellRef.ts
@@ -1,35 +1,30 @@
-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 || !isSelected || cell.contains(document.activeElement)) return;
+
+      cell.focus({ preventScroll: true });
+    },
+    [isSelected]
+  );
 
   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,
diff --git a/test/components.test.tsx b/test/components.test.tsx
index 90716824ff..a08661c1dd 100644
--- a/test/components.test.tsx
+++ b/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(

From a07b3ed56b559d099884e16944d3b4ca2cf2b06f Mon Sep 17 00:00:00 2001
From: Nicolas Stepien <stepien.nicolas@gmail.com>
Date: Thu, 7 Apr 2022 15:26:03 +0100
Subject: [PATCH 2/4] optimize

---
 src/hooks/useRovingCellRef.ts | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/src/hooks/useRovingCellRef.ts b/src/hooks/useRovingCellRef.ts
index e997f628c2..a33f282952 100644
--- a/src/hooks/useRovingCellRef.ts
+++ b/src/hooks/useRovingCellRef.ts
@@ -9,14 +9,11 @@ export function useRovingCellRef(isSelected: boolean) {
     setIsChildFocused(false);
   }
 
-  const ref = useCallback(
-    (cell: HTMLDivElement | null) => {
-      if (cell === null || !isSelected || cell.contains(document.activeElement)) return;
+  const ref = useCallback((cell: HTMLDivElement | null) => {
+    if (cell === null || cell.contains(document.activeElement)) return;
 
-      cell.focus({ preventScroll: true });
-    },
-    [isSelected]
-  );
+    cell.focus({ preventScroll: true });
+  }, []);
 
   function onFocus(event: React.FocusEvent<HTMLDivElement>) {
     if (event.target !== event.currentTarget) {
@@ -27,7 +24,7 @@ export function useRovingCellRef(isSelected: boolean) {
   const isFocused = isSelected && !isChildFocused;
 
   return {
-    ref,
+    ref: isSelected ? ref : undefined,
     tabIndex: isFocused ? 0 : -1,
     onFocus
   };

From 3763416a97a1d20aa1f5a65ee56519748f55bfc0 Mon Sep 17 00:00:00 2001
From: Nicolas Stepien <stepien.nicolas@gmail.com>
Date: Thu, 7 Apr 2022 15:39:32 +0100
Subject: [PATCH 3/4] lazy focus listening

---
 src/hooks/useRovingCellRef.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/hooks/useRovingCellRef.ts b/src/hooks/useRovingCellRef.ts
index a33f282952..6933d4733a 100644
--- a/src/hooks/useRovingCellRef.ts
+++ b/src/hooks/useRovingCellRef.ts
@@ -26,6 +26,6 @@ export function useRovingCellRef(isSelected: boolean) {
   return {
     ref: isSelected ? ref : undefined,
     tabIndex: isFocused ? 0 : -1,
-    onFocus
+    onFocus: isSelected ? onFocus : undefined
   };
 }

From 00a999c3ec205539697ac3e944aca8899bcb977a Mon Sep 17 00:00:00 2001
From: Nicolas Stepien <stepien.nicolas@gmail.com>
Date: Thu, 7 Apr 2022 15:47:06 +0100
Subject: [PATCH 4/4] fix type issue

---
 src/HeaderCell.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/HeaderCell.tsx b/src/HeaderCell.tsx
index 44b633c28b..6752c34c9d 100644
--- a/src/HeaderCell.tsx
+++ b/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);