From df04533aa56157985792e1ceb8aaf965a77a193d Mon Sep 17 00:00:00 2001
From: Andrei Zhaleznichenka <zhalea@amazon.de>
Date: Mon, 4 Sep 2023 15:47:53 +0200
Subject: [PATCH] refactor: Table selection reduce coupling with table (#1496)

---
 .../table-fragments/selection-custom.page.tsx | 135 ++++++++++++++++++
 pages/table-fragments/styles.scss             |   5 +
 src/cards/index.tsx                           |  11 +-
 src/table/__integ__/sticky-header.test.ts     |   2 +-
 src/table/internal.tsx                        |   5 +-
 .../selection/__integ__/selection.test.ts     |  66 +++++++++
 .../__tests__/use-selection.test.tsx          |   4 +-
 src/table/selection/index.ts                  |   8 ++
 src/table/selection/interfaces.ts             |  12 ++
 .../selection-control.tsx}                    |   4 +-
 .../styles.scss                               |   0
 .../selection/use-selection-focus-move.ts     |  57 ++++++++
 src/table/{ => selection}/use-selection.ts    |  96 ++-----------
 src/table/selection/utils.ts                  |  27 ++++
 src/table/thead.tsx                           |   3 +-
 src/test-utils/dom/table/index.ts             |   2 +-
 16 files changed, 335 insertions(+), 102 deletions(-)
 create mode 100644 pages/table-fragments/selection-custom.page.tsx
 create mode 100644 src/table/selection/__integ__/selection.test.ts
 rename src/table/{ => selection}/__tests__/use-selection.test.tsx (81%)
 create mode 100644 src/table/selection/index.ts
 create mode 100644 src/table/selection/interfaces.ts
 rename src/table/{selection-control/index.tsx => selection/selection-control.tsx} (97%)
 rename src/table/{selection-control => selection}/styles.scss (100%)
 create mode 100644 src/table/selection/use-selection-focus-move.ts
 rename src/table/{ => selection}/use-selection.ts (63%)
 create mode 100644 src/table/selection/utils.ts

diff --git a/pages/table-fragments/selection-custom.page.tsx b/pages/table-fragments/selection-custom.page.tsx
new file mode 100644
index 0000000000..176ba85c04
--- /dev/null
+++ b/pages/table-fragments/selection-custom.page.tsx
@@ -0,0 +1,135 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import React, { useContext, useState } from 'react';
+import { ColumnLayout, Container, ContentLayout, FormField, Header, Link, Select } from '~components';
+import { SelectionControl, focusMarkers, useSelectionFocusMove, useSelection } from '~components/table/selection';
+import styles from './styles.scss';
+import { generateItems, Instance } from '../table/generate-data';
+import AppContext, { AppContextType } from '../app/app-context';
+import ScreenreaderOnly from '~components/internal/components/screenreader-only';
+import clsx from 'clsx';
+
+type PageContext = React.Context<
+  AppContextType<{
+    selectionType: 'single' | 'multi';
+  }>
+>;
+
+const items = generateItems(25);
+
+const selectionTypeOptions = [{ value: 'single' }, { value: 'multi' }];
+
+export default function Page() {
+  const { urlParams, setUrlParams } = useContext(AppContext as PageContext);
+  const selectionType = urlParams.selectionType ?? 'single';
+
+  const [selectedItems, setSelectedItems] = useState<Instance[]>([]);
+  const { getSelectAllProps, getItemSelectionProps, updateShiftToggle } = useSelection({
+    items,
+    selectedItems,
+    selectionType,
+    onSelectionChange: event => setSelectedItems(event.detail.selectedItems),
+    ariaLabels: {
+      selectionGroupLabel: 'group label',
+      allItemsSelectionLabel: ({ selectedItems }) => `${selectedItems.length} item selected`,
+      itemSelectionLabel: ({ selectedItems }, item) =>
+        `${item.id} is ${selectedItems.indexOf(item) < 0 ? 'not ' : ''}selected`,
+    },
+  });
+
+  const { moveFocusDown, moveFocusUp, moveFocus } = useSelectionFocusMove(selectionType, items.length);
+
+  const columnDefinitions = [
+    {
+      key: 'selection',
+      header:
+        selectionType === 'multi' ? (
+          <SelectionControl
+            onFocusDown={event => moveFocus?.(event.target as HTMLElement, -1, +1)}
+            {...getSelectAllProps()}
+          />
+        ) : (
+          <ScreenreaderOnly>selection cell</ScreenreaderOnly>
+        ),
+      cell: (item: Instance) => (
+        <SelectionControl
+          onFocusDown={moveFocusDown}
+          onFocusUp={moveFocusUp}
+          onShiftToggle={updateShiftToggle}
+          {...getItemSelectionProps(item)}
+        />
+      ),
+    },
+    {
+      key: 'id',
+      header: 'ID',
+      cell: (item: Instance) => item.id,
+    },
+    {
+      key: 'imageId',
+      header: 'Image ID',
+      cell: (item: Instance) => <Link>{item.imageId}</Link>,
+    },
+    {
+      key: 'state',
+      header: 'State',
+      cell: (item: Instance) => item.state,
+    },
+    { key: 'dnsName', header: 'DNS name', cell: (item: Instance) => item.dnsName ?? '?' },
+    { key: 'type', header: 'Type', cell: (item: Instance) => item.type },
+  ];
+
+  return (
+    <ContentLayout header={<Header variant="h1">Rows selection with a custom table</Header>}>
+      <Container
+        disableContentPaddings={true}
+        header={
+          <ColumnLayout columns={3}>
+            <FormField label="Table role">
+              <Select
+                data-testid="selection-type"
+                options={selectionTypeOptions}
+                selectedOption={selectionTypeOptions.find(option => option.value === selectionType) ?? null}
+                onChange={event =>
+                  setUrlParams({ selectionType: event.detail.selectedOption.value as 'single' | 'multi' })
+                }
+              />
+            </FormField>
+          </ColumnLayout>
+        }
+        {...focusMarkers.root}
+      >
+        <div className={styles['custom-table']}>
+          <table className={styles['custom-table-table']} role="grid">
+            <thead>
+              <tr {...focusMarkers.all}>
+                {columnDefinitions.map(column => (
+                  <th
+                    key={column.key}
+                    className={clsx(
+                      styles['custom-table-cell'],
+                      column.key === 'selection' && styles['custom-table-selection-cell']
+                    )}
+                  >
+                    {column.header}
+                  </th>
+                ))}
+              </tr>
+            </thead>
+            <tbody>
+              {items.map((item, index) => (
+                <tr key={item.id} {...focusMarkers.item} data-rowindex={index + 1}>
+                  {columnDefinitions.map(column => (
+                    <td key={column.key} className={styles['custom-table-cell']}>
+                      {column.cell(item)}
+                    </td>
+                  ))}
+                </tr>
+              ))}
+            </tbody>
+          </table>
+        </div>
+      </Container>
+    </ContentLayout>
+  );
+}
diff --git a/pages/table-fragments/styles.scss b/pages/table-fragments/styles.scss
index 1a4ec51d01..936a318114 100644
--- a/pages/table-fragments/styles.scss
+++ b/pages/table-fragments/styles.scss
@@ -26,6 +26,7 @@
   }
 
   &-cell {
+    position: relative;
     border-left: 1px solid tokens.$color-border-divider-default;
     border-bottom: 1px solid tokens.$color-border-divider-default;
     padding: 8px 16px;
@@ -59,6 +60,10 @@
     }
   }
 
+  &-selection-cell {
+    width: 40px;
+  }
+
   &-sorting-header {
     padding: 0;
     margin: 0;
diff --git a/src/cards/index.tsx b/src/cards/index.tsx
index 5c75bb1128..1fca41033c 100644
--- a/src/cards/index.tsx
+++ b/src/cards/index.tsx
@@ -8,8 +8,13 @@ import { getCardsPerRow } from './cards-layout-helper';
 import { getBaseProps } from '../internal/base-component';
 import ToolsHeader from '../table/tools-header';
 import { getItemKey } from '../table/utils';
-import { focusMarkers, useFocusMove, useSelection } from '../table/use-selection';
-import SelectionControl, { SelectionControlProps } from '../table/selection-control';
+import {
+  SelectionControl,
+  SelectionControlProps,
+  focusMarkers,
+  useSelectionFocusMove,
+  useSelection,
+} from '../table/selection';
 import InternalContainer from '../container/internal';
 import InternalStatusIndicator from '../status-indicator/internal';
 import { applyDisplayName } from '../internal/utils/apply-display-name';
@@ -219,7 +224,7 @@ const CardsList = <T,>({
 }) => {
   const selectable = !!selectionType;
 
-  const { moveFocusDown, moveFocusUp } = useFocusMove(selectionType, items.length);
+  const { moveFocusDown, moveFocusUp } = useSelectionFocusMove(selectionType, items.length);
 
   let visibleSectionsDefinition = cardDefinition.sections || [];
   visibleSectionsDefinition = visibleSections
diff --git a/src/table/__integ__/sticky-header.test.ts b/src/table/__integ__/sticky-header.test.ts
index c989a4ff89..8b3c21737a 100644
--- a/src/table/__integ__/sticky-header.test.ts
+++ b/src/table/__integ__/sticky-header.test.ts
@@ -4,7 +4,7 @@ import useBrowser from '@cloudscape-design/browser-test-tools/use-browser';
 import createWrapper from '../../../lib/components/test-utils/selectors';
 import { BasePageObject } from '@cloudscape-design/browser-test-tools/page-objects';
 import styles from '../../../lib/components/table/styles.selectors.js';
-import selectionStyles from '../../../lib/components/table/selection-control/styles.selectors.js';
+import selectionStyles from '../../../lib/components/table/selection/styles.selectors.js';
 
 const tableWrapper = createWrapper().findTable();
 const tableScrollWrapper = tableWrapper.findByClassName(styles.wrapper);
diff --git a/src/table/internal.tsx b/src/table/internal.tsx
index 0b56f5707f..f23af672d8 100644
--- a/src/table/internal.tsx
+++ b/src/table/internal.tsx
@@ -11,10 +11,9 @@ import Thead, { TheadProps } from './thead';
 import { TableBodyCell } from './body-cell';
 import InternalStatusIndicator from '../status-indicator/internal';
 import { supportsStickyPosition } from '../internal/utils/dom';
-import SelectionControl from './selection-control';
 import { checkSortingState, getColumnKey, getItemKey, getVisibleColumnDefinitions, toContainerVariant } from './utils';
 import { useRowEvents } from './use-row-events';
-import { focusMarkers, useFocusMove, useSelection } from './use-selection';
+import { SelectionControl, focusMarkers, useSelectionFocusMove, useSelection } from './selection';
 import { fireNonCancelableEvent } from '../internal/events';
 import { isDevelopment } from '../internal/is-development';
 import { ColumnWidthDefinition, ColumnWidthsProvider, DEFAULT_COLUMN_WIDTH } from './use-column-widths';
@@ -120,7 +119,7 @@ const InternalTable = React.forwardRef(
 
     const handleScroll = useScrollSync([wrapperRefObject, scrollbarRef, secondaryWrapperRef]);
 
-    const { moveFocusDown, moveFocusUp, moveFocus } = useFocusMove(selectionType, items.length);
+    const { moveFocusDown, moveFocusUp, moveFocus } = useSelectionFocusMove(selectionType, items.length);
     const { onRowClickHandler, onRowContextMenuHandler } = useRowEvents({ onRowClick, onRowContextMenu });
 
     const visibleColumnDefinitions = getVisibleColumnDefinitions({
diff --git a/src/table/selection/__integ__/selection.test.ts b/src/table/selection/__integ__/selection.test.ts
new file mode 100644
index 0000000000..ad10da46f3
--- /dev/null
+++ b/src/table/selection/__integ__/selection.test.ts
@@ -0,0 +1,66 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { BasePageObject } from '@cloudscape-design/browser-test-tools/page-objects';
+import useBrowser from '@cloudscape-design/browser-test-tools/use-browser';
+
+class SelectionTestPage extends BasePageObject {
+  isChecked(selector: string) {
+    return this.browser.execute(selector => (document.querySelector(selector) as HTMLInputElement).checked, selector);
+  }
+}
+
+const setupTest = (selectionType: 'single' | 'multi', testFn: (page: SelectionTestPage) => Promise<void>) => {
+  return useBrowser(async browser => {
+    const page = new SelectionTestPage(browser);
+    await browser.url(`#/light/table-fragments/selection-custom?selectionType=${selectionType}`);
+    await page.waitForVisible('table');
+    await testFn(page);
+  });
+};
+
+describe('selection', () => {
+  test(
+    'selects first item with single selection',
+    setupTest('single', async page => {
+      await page.click('[data-testid="selection-type"]');
+      await page.keys(['Tab', 'Space']);
+      await expect(page.isChecked('tr[data-rowindex="1"] input')).resolves.toBe(true);
+      await expect(page.isChecked('tr[data-rowindex="2"] input')).resolves.toBe(false);
+    })
+  );
+
+  test(
+    'selects second item with single selection',
+    setupTest('single', async page => {
+      await page.click('[data-testid="selection-type"]');
+      await page.keys(['Tab', 'ArrowDown']);
+      await expect(page.isChecked('tr[data-rowindex="1"] input')).resolves.toBe(false);
+      await expect(page.isChecked('tr[data-rowindex="2"] input')).resolves.toBe(true);
+    })
+  );
+
+  test(
+    'selects first two items with multi selection',
+    setupTest('multi', async page => {
+      await page.click('[data-testid="selection-type"]');
+      await page.keys(['Tab', 'ArrowDown', 'Space']);
+      await page.keys(['ArrowDown', 'Space']);
+      await expect(page.isChecked('tr[data-rowindex="1"] input')).resolves.toBe(true);
+      await expect(page.isChecked('tr[data-rowindex="2"] input')).resolves.toBe(true);
+      await expect(page.isChecked('tr[data-rowindex="3"] input')).resolves.toBe(false);
+    })
+  );
+
+  test(
+    'selects all items but first with multi selection',
+    setupTest('multi', async page => {
+      await page.click('[data-testid="selection-type"]');
+      await page.keys(['Tab', 'Space']);
+      await page.keys(['ArrowDown', 'Space']);
+      await expect(page.isChecked('tr[data-rowindex="1"] input')).resolves.toBe(false);
+      await expect(page.isChecked('tr[data-rowindex="2"] input')).resolves.toBe(true);
+      await expect(page.isChecked('tr[data-rowindex="3"] input')).resolves.toBe(true);
+    })
+  );
+});
diff --git a/src/table/__tests__/use-selection.test.tsx b/src/table/selection/__tests__/use-selection.test.tsx
similarity index 81%
rename from src/table/__tests__/use-selection.test.tsx
rename to src/table/selection/__tests__/use-selection.test.tsx
index 3e1c8d7edb..510fb96372 100644
--- a/src/table/__tests__/use-selection.test.tsx
+++ b/src/table/selection/__tests__/use-selection.test.tsx
@@ -1,8 +1,8 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { useSelection } from '../../../lib/components/table/use-selection';
-import { renderHook } from '../../__tests__/render-hook';
+import { useSelection } from '../../../../lib/components/table/selection';
+import { renderHook } from '../../../__tests__/render-hook';
 
 describe('useSelection', () => {
   it('satisfies istanbul coverage', () => {
diff --git a/src/table/selection/index.ts b/src/table/selection/index.ts
new file mode 100644
index 0000000000..91c49e48cd
--- /dev/null
+++ b/src/table/selection/index.ts
@@ -0,0 +1,8 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+export { SelectionProps } from './interfaces';
+export { SelectionControl, SelectionControlProps } from './selection-control';
+export { useSelectionFocusMove } from './use-selection-focus-move';
+export { useSelection } from './use-selection';
+export { focusMarkers } from './utils';
diff --git a/src/table/selection/interfaces.ts b/src/table/selection/interfaces.ts
new file mode 100644
index 0000000000..19208200ed
--- /dev/null
+++ b/src/table/selection/interfaces.ts
@@ -0,0 +1,12 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+export interface SelectionProps {
+  name: string;
+  disabled: boolean;
+  selectionType: 'single' | 'multi';
+  indeterminate?: boolean;
+  checked: boolean;
+  onChange: () => void;
+  ariaLabel?: string;
+}
diff --git a/src/table/selection-control/index.tsx b/src/table/selection/selection-control.tsx
similarity index 97%
rename from src/table/selection-control/index.tsx
rename to src/table/selection/selection-control.tsx
index 6041beb78d..7204fe49fb 100644
--- a/src/table/selection-control/index.tsx
+++ b/src/table/selection/selection-control.tsx
@@ -8,7 +8,7 @@ import InternalCheckbox from '../../checkbox/internal';
 import RadioButton from '../../radio-group/radio-button';
 
 import styles from './styles.css.js';
-import { SelectionProps } from '../use-selection';
+import { SelectionProps } from './interfaces';
 
 export interface SelectionControlProps extends SelectionProps {
   onShiftToggle?(shiftPressed: boolean): void;
@@ -19,7 +19,7 @@ export interface SelectionControlProps extends SelectionProps {
   focusedComponent?: null | string;
 }
 
-export default function SelectionControl({
+export function SelectionControl({
   selectionType,
   indeterminate = false,
   onShiftToggle,
diff --git a/src/table/selection-control/styles.scss b/src/table/selection/styles.scss
similarity index 100%
rename from src/table/selection-control/styles.scss
rename to src/table/selection/styles.scss
diff --git a/src/table/selection/use-selection-focus-move.ts b/src/table/selection/use-selection-focus-move.ts
new file mode 100644
index 0000000000..41350333cd
--- /dev/null
+++ b/src/table/selection/use-selection-focus-move.ts
@@ -0,0 +1,57 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import { KeyboardEvent } from 'react';
+import { findUpUntil } from '../../internal/utils/dom';
+import { TableProps } from '../interfaces';
+import selectionStyles from './styles.css.js';
+import { SELECTION_ITEM } from './utils';
+
+// The hooks moves focus between multi-selection checkboxes.
+// Not eligible for tables with grid navigation.
+export function useSelectionFocusMove(selectionType: TableProps['selectionType'], totalItems: number) {
+  if (selectionType !== 'multi') {
+    return {};
+  }
+  function moveFocus(sourceElement: HTMLElement, fromIndex: number, direction: -1 | 1) {
+    let index = fromIndex;
+    const rootContainer = findRootContainer(sourceElement);
+
+    while (index >= -1 && index < totalItems) {
+      index += direction;
+      const control = findSelectionControlByIndex(rootContainer, index);
+      if (control && !control.disabled) {
+        control.focus();
+        break;
+      }
+    }
+  }
+  const [moveFocusDown, moveFocusUp] = ([1, -1] as const).map(direction => {
+    return (event: KeyboardEvent) => {
+      const target = event.currentTarget as HTMLElement;
+      const itemNode = findUpUntil(target, node => node.dataset.selectionItem === 'item')!;
+      const fromIndex = Array.prototype.indexOf.call(itemNode.parentElement!.children, itemNode);
+      moveFocus(target, fromIndex, direction);
+    };
+  });
+  return {
+    moveFocusDown,
+    moveFocusUp,
+    moveFocus,
+  };
+}
+
+function findSelectionControlByIndex(rootContainer: HTMLElement, index: number) {
+  if (index === -1) {
+    // find "select all" checkbox
+    return rootContainer.querySelector<HTMLInputElement>(
+      `[data-${SELECTION_ITEM}="all"] .${selectionStyles.root} input`
+    );
+  }
+  return rootContainer.querySelectorAll<HTMLInputElement>(
+    `[data-${SELECTION_ITEM}="item"] .${selectionStyles.root} input`
+  )[index];
+}
+
+function findRootContainer(element: HTMLElement) {
+  return findUpUntil(element, node => node.dataset.selectionRoot === 'true')!;
+}
diff --git a/src/table/use-selection.ts b/src/table/selection/use-selection.ts
similarity index 63%
rename from src/table/use-selection.ts
rename to src/table/selection/use-selection.ts
index e3791d33e9..e84a73b188 100644
--- a/src/table/use-selection.ts
+++ b/src/table/selection/use-selection.ts
@@ -1,93 +1,13 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import { KeyboardEvent, useState } from 'react';
-import { fireNonCancelableEvent } from '../internal/events';
-import { useUniqueId } from '../internal/hooks/use-unique-id';
-import { findUpUntil } from '../internal/utils/dom';
-import { TableProps } from './interfaces';
-import { getTrackableValue } from './utils';
-import selectionStyles from './selection-control/styles.css.js';
-import { joinStrings } from '../internal/utils/strings';
-
-export interface SelectionProps {
-  name: string;
-  disabled: boolean;
-  selectionType: 'single' | 'multi';
-  indeterminate?: boolean;
-  checked: boolean;
-  onChange: () => void;
-  ariaLabel?: string;
-}
-
-const SELECTION_ITEM = 'selection-item';
-const SELECTION_ROOT = 'selection-root';
-
-function findSelectionControlByIndex(rootContainer: HTMLElement, index: number) {
-  if (index === -1) {
-    // find "select all" checkbox
-    return rootContainer.querySelector<HTMLInputElement>(
-      `[data-${SELECTION_ITEM}="all"] .${selectionStyles.root} input`
-    );
-  }
-  return rootContainer.querySelectorAll<HTMLInputElement>(
-    `[data-${SELECTION_ITEM}="item"] .${selectionStyles.root} input`
-  )[index];
-}
-
-function findRootContainer(element: HTMLElement) {
-  return findUpUntil(element, node => node.dataset.selectionRoot === 'true')!;
-}
-
-export function useFocusMove(selectionType: TableProps['selectionType'], totalItems: number) {
-  if (selectionType !== 'multi') {
-    return {};
-  }
-  function moveFocus(sourceElement: HTMLElement, fromIndex: number, direction: -1 | 1) {
-    let index = fromIndex;
-    const rootContainer = findRootContainer(sourceElement);
-
-    while (index >= -1 && index < totalItems) {
-      index += direction;
-      const control = findSelectionControlByIndex(rootContainer, index);
-      if (control && !control.disabled) {
-        control.focus();
-        break;
-      }
-    }
-  }
-  const [moveFocusDown, moveFocusUp] = ([1, -1] as const).map(direction => {
-    return (event: KeyboardEvent) => {
-      const target = event.currentTarget as HTMLElement;
-      const itemNode = findUpUntil(target, node => node.dataset.selectionItem === 'item')!;
-      const fromIndex = Array.prototype.indexOf.call(itemNode.parentElement!.children, itemNode);
-      moveFocus(target, fromIndex, direction);
-    };
-  });
-  return {
-    moveFocusDown,
-    moveFocusUp,
-    moveFocus,
-  };
-}
-
-// A set, that compares items by their "trackables" (the results of applying `trackBy` to them)
-class ItemSet<T> {
-  constructor(trackBy: TableProps.TrackBy<T> | undefined, items: ReadonlyArray<T>) {
-    this.trackBy = trackBy;
-    items.forEach(this.put);
-  }
-  private trackBy: TableProps.TrackBy<T> | undefined;
-  private map: Map<unknown, T> = new Map();
-  put = (item: T) => this.map.set.call(this.map, getTrackableValue(this.trackBy, item), item);
-  has = (item: T) => this.map.has.call(this.map, getTrackableValue(this.trackBy, item));
-  forEach = this.map.forEach.bind(this.map);
-}
-
-export const focusMarkers = {
-  item: { ['data-' + SELECTION_ITEM]: 'item' },
-  all: { ['data-' + SELECTION_ITEM]: 'all' },
-  root: { ['data-' + SELECTION_ROOT]: 'true' },
-};
+import { useState } from 'react';
+import { fireNonCancelableEvent } from '../../internal/events';
+import { useUniqueId } from '../../internal/hooks/use-unique-id';
+import { TableProps } from '../interfaces';
+import { getTrackableValue } from '../utils';
+import { joinStrings } from '../../internal/utils/strings';
+import { SelectionProps } from './interfaces';
+import { ItemSet } from './utils';
 
 export function useSelection<T>({
   items,
diff --git a/src/table/selection/utils.ts b/src/table/selection/utils.ts
new file mode 100644
index 0000000000..e7d4664830
--- /dev/null
+++ b/src/table/selection/utils.ts
@@ -0,0 +1,27 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { TableProps } from '../interfaces';
+import { getTrackableValue } from '../utils';
+
+export const SELECTION_ITEM = 'selection-item';
+export const SELECTION_ROOT = 'selection-root';
+
+// A set, that compares items by their "trackables" (the results of applying `trackBy` to them)
+export class ItemSet<T> {
+  constructor(trackBy: TableProps.TrackBy<T> | undefined, items: ReadonlyArray<T>) {
+    this.trackBy = trackBy;
+    items.forEach(this.put);
+  }
+  private trackBy: TableProps.TrackBy<T> | undefined;
+  private map: Map<unknown, T> = new Map();
+  put = (item: T) => this.map.set.call(this.map, getTrackableValue(this.trackBy, item), item);
+  has = (item: T) => this.map.has.call(this.map, getTrackableValue(this.trackBy, item));
+  forEach = this.map.forEach.bind(this.map);
+}
+
+export const focusMarkers = {
+  item: { ['data-' + SELECTION_ITEM]: 'item' },
+  all: { ['data-' + SELECTION_ITEM]: 'all' },
+  root: { ['data-' + SELECTION_ROOT]: 'true' },
+};
diff --git a/src/table/thead.tsx b/src/table/thead.tsx
index 4be7f2c1ae..04a0f73083 100644
--- a/src/table/thead.tsx
+++ b/src/table/thead.tsx
@@ -3,8 +3,7 @@
 import clsx from 'clsx';
 import React from 'react';
 import { TableProps } from './interfaces';
-import SelectionControl from './selection-control';
-import { focusMarkers, SelectionProps } from './use-selection';
+import { SelectionControl, focusMarkers, SelectionProps } from './selection';
 import { fireNonCancelableEvent, NonCancelableEventHandler } from '../internal/events';
 import { getColumnKey } from './utils';
 import { TableHeaderCell } from './header-cell';
diff --git a/src/test-utils/dom/table/index.ts b/src/test-utils/dom/table/index.ts
index 35d360c837..99a3413c93 100644
--- a/src/test-utils/dom/table/index.ts
+++ b/src/test-utils/dom/table/index.ts
@@ -4,7 +4,7 @@ import { ComponentWrapper, ElementWrapper } from '@cloudscape-design/test-utils-
 import styles from '../../../table/styles.selectors.js';
 import headerCellStyles from '../../../table/header-cell/styles.selectors.js';
 import bodyCellStyles from '../../../table/body-cell/styles.selectors.js';
-import selectionStyles from '../../../table/selection-control/styles.selectors.js';
+import selectionStyles from '../../../table/selection/styles.selectors.js';
 import resizerStyles from '../../../table/resizer/styles.selectors.js';
 import CollectionPreferencesWrapper from '../collection-preferences';
 import ContainerWrapper from '../container';