[refactor] simplify logic (less useEffect, useRef, and move code) (#368)

* [refactor] extract onScollKeyDown to Scroller

* [refactor] remove unnecessary nesting in cell navigation context

* don't scroll when focusing the first row, and update comments

* [refactor] remove shouldScroll and complex useEffect

* [refactor] move code to focus the first cell to CellNavigationProvider, without useEffect, and add tests

* [refactor] remove reference to viewportRef

* [refactor] use ref callback and remove useEffect

* [refactor] replace useRef + useEffect with ref callbacks

* apply relevant copilot comments

Author: severo

src/components/Cell/Cell.tsx

@@ -1,5 +1,5 @@
 import type { KeyboardEvent, MouseEvent, ReactNode } from 'react'
-import { useCallback, useContext, useMemo, useRef } from 'react'
+import { useCallback, useContext, useMemo } from 'react'
 
 import { ColumnWidthsContext } from '../../contexts/ColumnWidthsContext.js'
 import type { ResolvedValue } from '../../helpers/dataframe/index.js'
@@ -44,8 +44,7 @@ interface Props {
  * @param {number} [props.rowNumber] the row index in the original data, undefined if the value has not been fetched yet
  */
 export default function Cell({ cell, onDoubleClickCell, onMouseDownCell, onKeyDownCell, stringify, columnIndex, visibleColumnIndex, className, ariaColIndex, ariaRowIndex, rowNumber, renderCellContent }: Props) {
-  const ref = useRef<HTMLTableCellElement | null>(null)
-  const { tabIndex, navigateToCell } = useCellFocus({ ref, ariaColIndex, ariaRowIndex })
+  const { tabIndex, navigateToCell, focusCellIfNeeded } = useCellFocus({ ariaColIndex, ariaRowIndex })
 
   // Get the column width from the context (use visibleColumnIndex for styling)
   const columnStyle = useContext(ColumnWidthsContext).getStyle?.(visibleColumnIndex)
@@ -90,7 +89,7 @@ export default function Cell({ cell, onDoubleClickCell, onMouseDownCell, onKeyDo
 
   return (
     <td
-      ref={ref}
+      ref={focusCellIfNeeded}
       role="cell"
       aria-busy={cell === undefined}
       aria-rowindex={ariaRowIndex}

src/components/ColumnHeader/ColumnHeader.tsx

@@ -31,11 +31,17 @@ interface Props {
 
 export default function ColumnHeader({ columnIndex, columnName, columnConfig, canMeasureWidth, direction, toggleOrderBy, orderByIndex, orderBySize, ariaColIndex, ariaRowIndex, className, children }: Props) {
   const ref = useRef<HTMLTableCellElement | null>(null)
-  const { tabIndex, navigateToCell } = useCellFocus({ ref, ariaColIndex, ariaRowIndex })
+  const { tabIndex, navigateToCell, focusCellIfNeeded } = useCellFocus({ ariaColIndex, ariaRowIndex })
   const { sortable } = columnConfig
   const { isOpen, position, menuId, close, handleMenuClick } = useColumnMenu(ref, navigateToCell)
   const { getHideColumn, showAllColumns } = useContext(ColumnVisibilityStatesContext)
 
+  const refCallback = useCallback((node: HTMLTableCellElement | null) => {
+    focusCellIfNeeded(node)
+    // set the current ref, it will be used to position the menu in handleMenuClick
+    ref.current = node
+  }, [focusCellIfNeeded])
+
   const handleClick = useCallback(() => {
     navigateToCell()
     if (sortable) toggleOrderBy?.()
@@ -150,7 +156,7 @@ export default function ColumnHeader({ columnIndex, columnName, columnConfig, ca
 
   return (
     <th
-      ref={ref}
+      ref={refCallback}
       scope="col"
       role="columnheader"
       aria-sort={direction ?? (sortable ? 'none' : undefined)}

src/components/ColumnMenu/ColumnMenu.tsx

@@ -126,7 +126,7 @@ export default function ColumnMenu({
         break
       case 'Enter':
       case ' ':
-      // Handled by the menu item
+        // Handled by the menu item
         break
       case 'ArrowUp':
       case 'ArrowLeft':
@@ -170,12 +170,14 @@ export default function ColumnMenu({
     }
     return () => {
       hideColumn()
-      // We focus the top left cell, which will always exist, because this column will disappear
-      focusFirstCell?.()
+
+      // The header cell was focused when opening the menu.
+      // Move the focus to the left cell to keep ability to navigate with keyboard.
+      focusFirstCell()
+
       close()
     }
   }, [hideColumn, close, focusFirstCell])
-
   const showAllColumnsAndClose = useMemo(() => {
     if (!showAllColumns) {
       return undefined

src/components/HighTable/Scroller.tsx

@@ -1,9 +1,10 @@
 import type { KeyboardEvent } from 'react'
-import { useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react'
+import { useCallback, useContext, useMemo, useState } from 'react'
 
 import { CellNavigationContext } from '../../contexts/CellNavigationContext.js'
 import { DataContext } from '../../contexts/DataContext.js'
 import { RowsAndColumnsContext } from '../../contexts/RowsAndColumnsContext.js'
+import { ScrollerContext } from '../../contexts/ScrollerContext.js'
 import styles from '../../HighTable.module.css'
 import { ariaOffset, rowHeight } from './constants.js'
 
@@ -18,15 +19,10 @@ export default function Scroller({
   setViewportWidth,
   children,
 }: Props) {
-  // TODO(SL): replace with a callback function (https://react.dev/reference/react-dom/components/common#ref-callback)
-  const viewportRef = useRef<HTMLDivElement>(null)
-
-  const [scrollTop, setScrollTop] = useState<number | undefined>(undefined)
   const [scrollToTop, setScrollToTop] = useState<((top: number) => void) | undefined>(undefined)
 
   const { numRows } = useContext(DataContext)
-  const { onScrollKeyDown } = useContext(CellNavigationContext)
-  const { shouldScroll, setShouldScroll, cellPosition } = useContext(CellNavigationContext)
+  const { setShouldFocus, rowIndex } = useContext(CellNavigationContext)
   const { fetchedRowsRange, renderedRowsRange, setVisibleRowsRange } = useContext(RowsAndColumnsContext)
 
   /**
@@ -61,52 +57,55 @@ export default function Scroller({
   }, [numRows, scrollHeight, setVisibleRowsRange])
 
   /**
-   * Handle keyboard events for scrolling
+   * Vertically scroll to bring a specific row into view
    */
-  const onKeyDown = useCallback((event: KeyboardEvent) => {
-    if (event.target !== viewportRef.current) {
-      // don't handle the event if the target is not the scroller
+  const scrollRowIntoView = useCallback(({ rowIndex }: { rowIndex: number }) => {
+    if (scrollToTop === undefined || fetchedRowsRange === undefined) {
       return
     }
-    onScrollKeyDown?.(event)
-  }, [onScrollKeyDown, viewportRef])
-
-  /**
-   * React to cell navigation changes to scroll to the focused cell
-   *
-   * scroll if the navigation cell changed, or if entering navigation mode
-   * this excludes the case where the whole table is focused (not in cell navigation mode), the user
-   * is scrolling with the mouse or the arrow keys, and the cell exits the viewport: don't want to scroll
-   * back to it
-   */
-  useEffect(() => {
-    if (!shouldScroll || scrollTop === undefined || scrollToTop === undefined || fetchedRowsRange === undefined) {
+    if (rowIndex < 1) {
+      throw new Error(`invalid rowIndex ${rowIndex}`)
+    }
+    if (rowIndex === 1) {
+      // always visible
       return
     }
-    setShouldScroll?.(false)
-    const row = cellPosition.rowIndex - ariaOffset
-    let nextScrollTop = scrollTop
+    // should be zero-based
+    const row = rowIndex - ariaOffset
     // if the row is outside of the fetched rows range, scroll to the estimated position of the cell,
     // to wait for the cell to be fetched and rendered
     // TODO(SL): should fetchedRowsRange be replaced with visibleRowsRange?
     if (row < fetchedRowsRange.start || row >= fetchedRowsRange.end) {
-      nextScrollTop = row * rowHeight
-    }
-    if (nextScrollTop !== scrollTop) {
-      // scroll to the cell
-      scrollToTop(nextScrollTop)
+      scrollToTop(row * rowHeight)
     }
-  }, [cellPosition, shouldScroll, fetchedRowsRange, setShouldScroll, scrollToTop, scrollTop])
+  }, [fetchedRowsRange, scrollToTop])
 
   /**
-   * Track viewport size and scroll position
+   * Handle keyboard events for scrolling
    */
-  useEffect(() => {
-    const viewport = viewportRef.current
-    if (!viewport) {
-      throw new Error('Viewport element is not available. Viewport size will not be tracked accurately.')
+  const onKeyDown = useCallback((event: KeyboardEvent) => {
+    if (event.target !== event.currentTarget) {
+      // don't handle the event if the target is not the scroller
+      return
     }
+    const { key } = event
+    // the user can scroll with the keyboard using the arrow keys. Here we only handle the Tab,
+    // Enter and Space keys, to enter the cell navigation mode instead of scrolling the table
+    // TODO(SL): exclude other meta keys (return without handling if shiftKey, ctrlKey, altKey, metaKey)
+    if ((key === 'Tab' && !event.shiftKey) || key === 'Enter' || key === ' ') {
+      event.stopPropagation()
+      event.preventDefault()
+      // scroll to the active cell
+      scrollRowIntoView({ rowIndex })
+      // focus the cell (once it exists)
+      setShouldFocus(true)
+    }
+  }, [rowIndex, setShouldFocus, scrollRowIntoView])
 
+  /**
+   * Track viewport size and scroll position
+   */
+  const viewportRef = useCallback((viewport: HTMLDivElement) => {
     // Use arrow functions to get correct viewport type (not null)
     // eslint-disable-next-line func-style
     const updateViewportSize = () => {
@@ -117,7 +116,7 @@ export default function Scroller({
     // eslint-disable-next-line func-style
     const handleScroll = () => {
       // TODO(SL): throttle? see https://github.com/hyparam/hightable/pull/347
-      setScrollTop(viewport.scrollTop)
+      // recompute the rows range if the scroll position changed
       computeAndSetRowsRange(viewport)
     }
 
@@ -168,7 +167,9 @@ export default function Scroller({
     <div className={styles.tableScroll} ref={viewportRef} role="group" aria-labelledby="caption" onKeyDown={onKeyDown} tabIndex={0}>
       <div style={{ height: `${scrollHeight}px` }}>
         <div style={{ top: `${top}px` }}>
-          {children}
+          <ScrollerContext.Provider value={{ scrollRowIntoView }}>
+            {children}
+          </ScrollerContext.Provider>
         </div>
       </div>
     </div>

src/components/HighTable/Slice.tsx

@@ -1,10 +1,11 @@
 import type { KeyboardEvent, MouseEvent, ReactNode } from 'react'
-import { useCallback, useContext, useEffect, useMemo } from 'react'
+import { useCallback, useContext, useMemo } from 'react'
 
 import { CellNavigationContext } from '../../contexts/CellNavigationContext.js'
 import { DataContext } from '../../contexts/DataContext.js'
 import { OrderByContext } from '../../contexts/OrderByContext.js'
 import { RowsAndColumnsContext } from '../../contexts/RowsAndColumnsContext.js'
+import { ScrollerContext } from '../../contexts/ScrollerContext.js'
 import { SelectionContext } from '../../contexts/SelectionContext.js'
 import { stringify as stringifyDefault } from '../../utils/stringify.js'
 import Cell, { type CellContentProps } from '../Cell/Cell.js'
@@ -15,7 +16,6 @@ import TableHeader from '../TableHeader/TableHeader.js'
 import { ariaOffset, defaultNumRowsPerPage } from './constants.js'
 
 export interface SliceProps {
-  focus?: boolean // focus table on mount? (default true)
   numRowsPerPage?: number // number of rows per page for keyboard navigation (default 20)
   // TODO(SL): replace col: number with col: string?
   onDoubleClickCell?: (event: MouseEvent, col: number, row: number) => void
@@ -30,7 +30,6 @@ type Props = {
 } & SliceProps
 
 export default function Slice({
-  focus = true,
   numRowsPerPage = defaultNumRowsPerPage,
   onDoubleClickCell,
   onKeyDownCell,
@@ -40,13 +39,82 @@ export default function Slice({
   stringify = stringifyDefault,
 }: Props) {
   const { data, version, numRows } = useContext(DataContext)
-  const { onTableKeyDown: onNavigationTableKeyDown, focusFirstCell } = useContext(CellNavigationContext)
+  const { rowIndex, colCount, rowCount, setColIndex, setRowIndex, setShouldFocus } = useContext(CellNavigationContext)
   const { orderBy, onOrderByChange } = useContext(OrderByContext)
   const { selectable, toggleAllRows, pendingSelectionGesture, onTableKeyDown: onSelectionTableKeyDown, allRowsSelected, isRowSelected, toggleRowNumber, toggleRangeToRowNumber } = useContext(SelectionContext)
   const { columnsParameters, renderedRowsRange, fetchedRowsRange } = useContext(RowsAndColumnsContext)
+  const { scrollRowIntoView } = useContext(ScrollerContext)
+
+  // TODO(SL): we depend on rowIndex to trigger the scroll effect, which means we recreate the
+  // callback every time the rowIndex changes. Can we avoid that?
+  // For now, we don't need to depend on colIndex, we can set the state using the update function form.
+  const onNavigationTableKeyDown = useCallback((event: KeyboardEvent, { numRowsPerPage }: {
+    numRowsPerPage: number // number of rows to skip when navigating with the keyboard (PageUp/PageDown)
+  }) => {
+    const { key, altKey, ctrlKey, metaKey, shiftKey } = event
+    // if the user is pressing Alt, Meta or Shift, do not handle the event
+    if (altKey || metaKey || shiftKey) {
+      return
+    }
+    let newRowIndex: number | undefined = undefined
+    if (key === 'ArrowRight') {
+      if (ctrlKey) {
+        setColIndex(colCount)
+      } else {
+        setColIndex(prev => prev < colCount ? prev + 1 : prev)
+      }
+    } else if (key === 'ArrowLeft') {
+      if (ctrlKey) {
+        setColIndex(1)
+      } else {
+        setColIndex(prev => prev > 1 ? prev - 1 : prev)
+      }
+    } else if (key === 'ArrowDown') {
+      if (ctrlKey) {
+        newRowIndex = rowCount
+      } else {
+        newRowIndex = rowIndex < rowCount ? rowIndex + 1 : rowIndex
+      }
+    } else if (key === 'ArrowUp') {
+      if (ctrlKey) {
+        newRowIndex = 1
+      } else {
+        newRowIndex = rowIndex > 1 ? rowIndex - 1 : rowIndex
+      }
+    } else if (key === 'Home') {
+      if (ctrlKey) {
+        newRowIndex = 1
+      }
+      setColIndex(1)
+    } else if (key === 'End') {
+      if (ctrlKey) {
+        newRowIndex = rowCount
+      }
+      setColIndex(colCount)
+    } else if (key === 'PageDown') {
+      newRowIndex = rowIndex + numRowsPerPage <= rowCount ? rowIndex + numRowsPerPage : rowCount
+      // TODO(SL): same for horizontal scrolling with Alt+PageDown?
+    } else if (key === 'PageUp') {
+      newRowIndex = rowIndex - numRowsPerPage >= 1 ? rowIndex - numRowsPerPage : 1
+      // TODO(SL): same for horizontal scrolling with Alt+PageUp?
+    } else if (key !== ' ') {
+      // if the key is not one of the above, do not handle it
+      // special case: no action is associated with the Space key, but it's captured
+      // anyway to prevent the default action (scrolling the page) and stay in navigation mode
+      return
+    }
+    // avoid scrolling the table when the user is navigating with the keyboard
+    event.stopPropagation()
+    event.preventDefault()
+    if (newRowIndex !== undefined) {
+      setRowIndex(newRowIndex)
+      scrollRowIntoView?.({ rowIndex: newRowIndex }) // ensure the cell is visible
+    }
+    setShouldFocus(true)
+  }, [rowIndex, colCount, rowCount, setColIndex, setRowIndex, setShouldFocus, scrollRowIntoView])
 
   const onTableKeyDown = useCallback((event: KeyboardEvent) => {
-    onNavigationTableKeyDown?.(event, { numRowsPerPage })
+    onNavigationTableKeyDown(event, { numRowsPerPage })
     onSelectionTableKeyDown?.(event)
   }, [onNavigationTableKeyDown, onSelectionTableKeyDown, numRowsPerPage])
 
@@ -63,16 +131,6 @@ export default function Slice({
     }
   }, [toggleRowNumber, toggleRangeToRowNumber])
 
-  // focus table on mount, or on later changes, so arrow keys work
-  // Note that the dependency upon data and numRows was removed, because focusFirstCell should depend on them
-  // TODO(SL): move to CellNavigationProvider?
-  useEffect(() => {
-    if (focus) {
-      // Try focusing the first cell
-      focusFirstCell?.()
-    }
-  }, [focus, focusFirstCell])
-
   // Prepare the slice of data to render
   // TODO(SL): also compute progress percentage here, to show a loading indicator
   const slice = useMemo(() => {

src/components/HighTable/Wrapper.tsx

@@ -7,6 +7,7 @@ import type { ColumnConfiguration } from '../../helpers/columnConfiguration.js'
 import type { Selection } from '../../helpers/selection.js'
 import type { OrderBy } from '../../helpers/sort.js'
 import styles from '../../HighTable.module.css'
+import type { CellNavigationProviderProps } from '../../providers/CellNavigationProvider.js'
 import { CellNavigationProvider } from '../../providers/CellNavigationProvider.js'
 import { ColumnParametersProvider } from '../../providers/ColumnParametersProvider.js'
 import { ColumnVisibilityStatesProvider } from '../../providers/ColumnVisibilityStatesProvider.js'
@@ -32,12 +33,13 @@ export type WrapperProps = {
   onColumnsVisibilityChange?: (columns: Record<string, MaybeHiddenColumn>) => void // callback which is called whenever the set of hidden columns changes.
   onOrderByChange?: (orderBy: OrderBy) => void // callback to call when a user interaction changes the order. The interactions are disabled if undefined.
   onSelectionChange?: (selection: Selection) => void // callback to call when a user interaction changes the selection. The selection is expressed as data indexes (not as indexes in the table). The interactions are disabled if undefined.
-} & RowsAndColumnsProviderProps & SliceProps
+} & RowsAndColumnsProviderProps & CellNavigationProviderProps & SliceProps
 
 export default function Wrapper({
   columnConfiguration,
   cacheKey,
   className = '',
+  focus,
   orderBy,
   padding,
   overscan,
@@ -98,7 +100,7 @@ export default function Wrapper({
                 {/* Create a new selection context if the dataframe has changed */}
                 <SelectionProvider key={key} selection={selection} onSelectionChange={onSelectionChange} data={data} numRows={numRows}>
                   {/* Create a new navigation context if the dataframe has changed, because the focused cell might not exist anymore */}
-                  <CellNavigationProvider key={key}>
+                  <CellNavigationProvider key={key} focus={focus}>
                     <RowsAndColumnsProvider key={key} padding={padding} overscan={overscan}>
 
                       <Scroller setViewportWidth={setViewportWidth} headerHeight={headerHeight}>