lightning-base-components 1.21.2-alpha → 1.21.4-alpha

package/src/lightning/datatable/rowSelection.js

@@ -1,33 +1,11 @@
-import {
-    resolveRowClassNames,
-    getRows,
-    getRowByKey,
-    getRowsTotal,
-    getRowIndexByKey,
-    rowKeyExists,
-} from './rows';
-import {
-    getColumns,
-    getStateColumnIndex,
-    SELECTABLE_ROW_CHECKBOX,
-} from './columns';
+import { getRowByKey, getRowIndexByKey, HEADER_ROW_KEY } from './indexes';
+import { getStateColumnIndex } from './columns';
 import { isNonNegativeInteger } from './utils';
 import {
     getCurrentSelectionLength,
-    isSelectedRow,
-    isDisabledRow,
-    getRowSelectionInputType,
-    getMaxRowSelection,
     getSelectedRowsKeys,
-} from './rowSelectionShared';
-
-export {
-    getCurrentSelectionLength,
-    isSelectedRow,
     isDisabledRow,
-    getRowSelectionInputType,
-    getMaxRowSelection,
-    getSelectedRowsKeys,
+    SELECTABLE_HEADER_TYPE,
 } from './rowSelectionShared';
 
 const MAX_ROW_SELECTION_DEFAULT = undefined;
@@ -79,15 +57,13 @@ export function handleDeselectAllRows(event) {
  */
 export function handleSelectRow(event) {
     event.stopPropagation();
+    const { state } = this;
     const { rowKeyValue, isMultiple } = event.detail;
-    let fromRowKey = rowKeyValue;
-
-    if (isMultiple) {
-        fromRowKey = getLastRowSelection(this.state) || rowKeyValue;
-    }
-
-    markSelectedRowsInterval(this.state, fromRowKey, rowKeyValue);
-    setLastRowSelection(this.state, rowKeyValue);
+    const fromRowKey = isMultiple
+        ? getLastRowSelection(state) || rowKeyValue
+        : rowKeyValue;
+    markSelectedRowsInterval(state, fromRowKey, rowKeyValue);
+    state.lastSelectedRowKey = rowKeyValue;
     this.fireSelectedRowsChange(this.getSelectedRows(), {
         action: ROWS_ACTION.ROW_SELECT,
         value: rowKeyValue,
@@ -105,15 +81,13 @@ export function handleSelectRow(event) {
  */
 export function handleDeselectRow(event) {
     event.stopPropagation();
+    const { state } = this;
     const { rowKeyValue, isMultiple } = event.detail;
-    let fromRowKey = rowKeyValue;
-
-    if (isMultiple) {
-        fromRowKey = getLastRowSelection(this.state) || rowKeyValue;
-    }
-
-    markDeselectedRowsInterval(this.state, fromRowKey, rowKeyValue);
-    setLastRowSelection(this.state, rowKeyValue);
+    const fromRowKey = isMultiple
+        ? getLastRowSelection(state) || rowKeyValue
+        : rowKeyValue;
+    markDeselectedRowsInterval(state, fromRowKey, rowKeyValue);
+    state.lastSelectedRowKey = rowKeyValue;
     this.fireSelectedRowsChange(this.getSelectedRows(), {
         action: ROWS_ACTION.ROW_DESELECT,
         value: rowKeyValue,
@@ -139,22 +113,22 @@ export function handleRowSelectionChange() {
  *     2. If max-row-selection has been reached, mark the remaining rows
  *        to reflect that they are not selected and disable them.
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  */
 export function markAllRowsSelected(state) {
-    const rows = getRows(state);
-    const maxRowSelection = getMaxRowSelection(state);
-
-    resetSelectedRowsKeys(state);
-    rows.forEach((row, index) => {
-        if (index < maxRowSelection || maxRowSelection === undefined) {
-            setRowSelectedAttributes(true, row);
-            addKeyToSelectedRowKeys(state, row.key);
+    const { maxRowSelection, rows } = state;
+    const selectedRowsKeys = {};
+    for (let i = 0, { length: rowCount } = rows; i < rowCount; i += 1) {
+        const row = rows[i];
+        if (maxRowSelection === undefined || i < maxRowSelection) {
+            selectedRowsKeys[row.key] = true;
+            setRowSelectedAttributes(row, true);
         } else {
             row.isDisabled = true;
-            setRowSelectedAttributes(false, row);
+            setRowSelectedAttributes(row, false);
         }
-    });
+    }
+    state.selectedRowsKeys = selectedRowsKeys;
 }
 
 /**
@@ -163,17 +137,16 @@ export function markAllRowsSelected(state) {
  * set `isSelected` and `ariaSelected` to false and enable the row. Also resolve
  * the `classnames` for the row to reflect that it is not selected.
  *
- * @param {Object} state - datatable's state object
- * @returns
+ * @param {Object} state - The datatable state
  */
 export function markAllRowsDeselected(state) {
-    const rows = getRows(state);
-
-    resetSelectedRowsKeys(state);
-    rows.forEach((row) => {
+    const { rows } = state;
+    for (let i = 0, { length: rowCount } = rows; i < rowCount; i += 1) {
+        const row = rows[i];
         row.isDisabled = false;
-        setRowSelectedAttributes(false, row);
-    });
+        setRowSelectedAttributes(row, false);
+    }
+    state.selectedRowsKeys = {};
 }
 
 /**
@@ -190,22 +163,20 @@ export function markAllRowsDeselected(state) {
  * This does not handle the case when the header checkbox that selects
  * all rows of the table is clicked. That is handled by - `handleSelectAllRows`
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  * @param {String} startRowKey - row key value of the first row that was selected (start of the interval)
  * @param {String} endRowKey - row key value of the last row that was selected (end of the interval)
  */
 function markSelectedRowsInterval(state, startRowKey, endRowKey) {
-    const rows = getRows(state);
+    const { rows } = state;
     const { start, end } = getRowIntervalIndexes(state, startRowKey, endRowKey);
-    const maxRowSelection = getMaxRowSelection(state) || getRowsTotal(state);
-    let i = start,
-        maxSelectionReached;
-
-    while (i <= end && !maxSelectionReached) {
-        markRowSelected(state, rows[i].key);
-        maxSelectionReached =
-            getCurrentSelectionLength(state) >= maxRowSelection;
-        i++;
+    const total = state.maxRowSelection || rows.length;
+    let rowIndex = start;
+    let maxSelectionReached;
+    while (rowIndex <= end && !maxSelectionReached) {
+        markRowSelected(state, rows[rowIndex].key);
+        maxSelectionReached = getCurrentSelectionLength(state) >= total;
+        rowIndex += 1;
     }
 }
 
@@ -222,16 +193,22 @@ function markSelectedRowsInterval(state, startRowKey, endRowKey) {
  * This does not handle the case when the header checkbox is clicked to de-select all rows
  * That is handledd by - `handleDeselectAllRows`
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  * @param {String} startRowKey - row key value of the first row that was selected (start of the interval)
  * @param {String} endRowKey - row key value of the last row that was selected (end of the interval)
  */
 function markDeselectedRowsInterval(state, startRowKey, endRowKey) {
-    const rows = getRows(state);
+    const { rows, selectedRowsKeys } = state;
     const { start, end } = getRowIntervalIndexes(state, startRowKey, endRowKey);
 
-    for (let i = start; i <= end; i++) {
-        markRowDeselected(state, rows[i].key);
+    for (let rowIndex = start; rowIndex <= end; rowIndex += 1) {
+        const row = rows[rowIndex];
+        setRowSelectedAttributes(row, false);
+        // Setting to `false` instead of using `delete` for better performance.
+        selectedRowsKeys[row.key] = false;
+        if (getCurrentSelectionLength(state) === state.maxRowSelection - 1) {
+            markDeselectedRowEnabled(state);
+        }
     }
 }
 
@@ -248,56 +225,33 @@ function markDeselectedRowsInterval(state, startRowKey, endRowKey) {
  *         a. If another row was previously selected before, de-select that row
  *         b. Add the row key value of that row to the state
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  * @param {String} rowKeyValue - row key value of row to mark selected
  */
 export function markRowSelected(state, rowKeyValue) {
     const row = getRowByKey(state, rowKeyValue);
-    const maxRowSelection = getMaxRowSelection(state) || getRowsTotal(state);
-    const previousSelectionLength = getCurrentSelectionLength(state);
+    const prevSelectionLength = getCurrentSelectionLength(state);
+    const total = state.maxRowSelection || state.rows.length;
 
-    setRowSelectedAttributes(true, row);
+    setRowSelectedAttributes(row, true);
 
-    if (maxRowSelection > 1) {
-        addKeyToSelectedRowKeys(state, row.key);
-        if (previousSelectionLength + 1 === maxRowSelection) {
+    let { selectedRowsKeys } = state;
+    if (total > 1) {
+        selectedRowsKeys[row.key] = true;
+        if (prevSelectionLength + 1 === total) {
             markDeselectedRowDisabled(state);
         }
     } else {
-        if (previousSelectionLength === 1) {
-            const previousSelectedRow = getRowByKey(
+        if (prevSelectionLength === 1) {
+            const prevSelectedRow = getRowByKey(
                 state,
-                Object.keys(state.selectedRowsKeys)[0]
+                Object.keys(selectedRowsKeys)[0]
             );
-            setRowSelectedAttributes(false, previousSelectedRow);
-            resetSelectedRowsKeys(state);
+            setRowSelectedAttributes(prevSelectedRow, false);
+            selectedRowsKeys = {};
+            state.selectedRowsKeys = selectedRowsKeys;
         }
-        addKeyToSelectedRowKeys(state, row.key);
-    }
-}
-
-/**
- * Marks a row with the specified row key value as de-selected. This is done by:
- *     1. Sets `isSelected`, `ariaSelected` to false and resolves `classnames`
- *        to that which reflect that the row is not selected, on the row object.
- *        These are used by the template to render the appropriate values.
- *     2. The row key value of this row is removed from list of selected row keys
- *     3. Before the current de-selection, if the remaining unselected rows were
- *        disabled (max-row-selection was reached), enable all unselected rows
- *        so that they can be selected.
- *
- * @param {Object} state - datatable's state object
- * @param {String} rowKeyValue - row key value of row to mark selected
- */
-export function markRowDeselected(state, rowKeyValue) {
-    const row = getRowByKey(state, rowKeyValue);
-    const maxRowSelection = getMaxRowSelection(state);
-
-    setRowSelectedAttributes(false, row);
-    removeKeyFromSelectedRowKeys(state, row.key);
-
-    if (getCurrentSelectionLength(state) === maxRowSelection - 1) {
-        markDeselectedRowEnabled(state);
+        selectedRowsKeys[row.key] = true;
     }
 }
 
@@ -307,14 +261,14 @@ export function markRowDeselected(state, rowKeyValue) {
  * Sets `isSelected`, `ariaSelected` and `classnames` on the row object
  * which are used by the template to render the appropriate values.
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  * @param {Array} keys - a list of row key values to be marked selected
  */
 function markRowsSelectedByKeys(state, keys) {
-    keys.forEach((rowKeyValue) => {
-        const row = getRowByKey(state, rowKeyValue);
-        setRowSelectedAttributes(true, row);
-    });
+    for (let i = 0, { length } = keys; i < length; i += 1) {
+        const row = getRowByKey(state, keys[i]);
+        setRowSelectedAttributes(row, true);
+    }
 }
 
 /**
@@ -324,14 +278,14 @@ function markRowsSelectedByKeys(state, keys) {
  * to one which reflects that the row is not selected on the row object.
  * These are used by the template to render the appropriate values.
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  * @param {Array} keys - a list of row key values to be marked selected
  */
 function markRowsDeselectedByKeys(state, keys) {
-    keys.forEach((rowKeyValue) => {
-        const row = getRowByKey(state, rowKeyValue);
-        setRowSelectedAttributes(false, row);
-    });
+    for (let i = 0, { length } = keys; i < length; i += 1) {
+        const row = getRowByKey(state, keys[i]);
+        setRowSelectedAttributes(row, false);
+    }
 }
 
 /**
@@ -343,12 +297,13 @@ function markRowsDeselectedByKeys(state, keys) {
  * @param {Object} state - datatable's state object
  */
 export function markDeselectedRowDisabled(state) {
-    const rows = getRows(state);
-    rows.forEach((row) => {
-        if (!isSelectedRow(state, row.key)) {
+    const { rows, selectedRowsKeys } = state;
+    for (let i = 0, { length: rowCount } = rows; i < rowCount; i += 1) {
+        const row = rows[i];
+        if (!selectedRowsKeys[row.key]) {
             row.isDisabled = true;
         }
-    });
+    }
 }
 
 /**
@@ -358,15 +313,16 @@ export function markDeselectedRowDisabled(state) {
  * previously selected but a row was deselected, now allowing
  * any other row to be selected - for this, all rows should be enabled
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  */
 export function markDeselectedRowEnabled(state) {
-    const rows = getRows(state);
-    rows.forEach((row) => {
-        if (!isSelectedRow(state, row.key)) {
+    const { rows, selectedRowsKeys } = state;
+    for (let i = 0, { length: rowCount } = rows; i < rowCount; i += 1) {
+        const row = rows[i];
+        if (!selectedRowsKeys[row.key]) {
             row.isDisabled = false;
         }
-    });
+    }
 }
 
 /************************** SELECTED ROW KEYS **************************/
@@ -391,12 +347,12 @@ export function markDeselectedRowEnabled(state) {
  * If the previous selection had reached the max limit and the new selection
  * is less than the limit, re-enable the de-selected rows to allow for new selection.
  *
- * @param {Object} state - datatable's state object
+ * @param {Object} state - The datatable state
  * @param {Array} value - key-field values of rows to set as selected
  */
 export function setSelectedRowsKeys(state, value) {
     if (Array.isArray(value)) {
-        const maxRowSelection = getMaxRowSelection(state);
+        const { maxRowSelection } = state;
         const previousSelectionLength = getCurrentSelectionLength(state);
         let selectedRows = filterValidKeys(state, value);
         if (selectedRows.length > maxRowSelection) {
@@ -409,8 +365,7 @@ export function setSelectedRowsKeys(state, value) {
 
         // Convert the selectedRows Array to an Object that state.selectedRowKeys expects
         // ['a', 'b'] -> { a : true, b : true}
-        const normalizedSelectedRowsKeys =
-            normalizeSelectedRowsKey(selectedRows);
+        const selectedRowsKeys = normalizeSelectedRowsKey(selectedRows);
 
         // Compute differences between currently selected rows and
         // newly selected row keys. The diff will tell which new rows
@@ -419,18 +374,19 @@ export function setSelectedRowsKeys(state, value) {
         const selectionOperations = getSelectedDiff(state, selectedRows);
         const deselectionOperations = getDeselectedDiff(
             state,
-            normalizedSelectedRowsKeys
+            selectedRowsKeys
         );
         markRowsSelectedByKeys(state, selectionOperations);
         markRowsDeselectedByKeys(state, deselectionOperations);
-        state.selectedRowsKeys = normalizedSelectedRowsKeys;
+        state.selectedRowsKeys = selectedRowsKeys;
 
+        const { length: selectedRowsCount } = selectedRows;
         // If we select the max number of rows allowed and if max-row-selection > 1 (multi-select),
         // disable all the other rows to prevent further selection
-        if (selectedRows.length === maxRowSelection && maxRowSelection > 1) {
+        if (selectedRowsCount === maxRowSelection && maxRowSelection > 1) {
             markDeselectedRowDisabled(state);
         } else if (
-            selectedRows.length < maxRowSelection &&
+            selectedRowsCount < maxRowSelection &&
             previousSelectionLength === maxRowSelection
         ) {
             // If the previous selection had reached the max limit and the new selection
@@ -448,34 +404,44 @@ export function setSelectedRowsKeys(state, value) {
 
 export function syncSelectedRowsKeys(state, selectedRows) {
     let changed = false;
-    const { selectedRowsKeys, keyField } = state;
-    const maxRowSelection = getMaxRowSelection(state) || getRowsTotal(state);
-    if (Object.keys(selectedRowsKeys).length !== selectedRows.length) {
+    const { keyField, selectedRowsKeys } = state;
+    const { length: selectedRowCount } = selectedRows;
+    if (Object.keys(selectedRowsKeys).length !== selectedRowCount) {
         changed = true;
+        // Untracked state change.
         state.selectedRowsKeys = updateSelectedRowsKeysFromSelectedRows(
             selectedRows,
             keyField
         );
     } else {
-        changed = selectedRows.some((row) => !selectedRowsKeys[row[keyField]]);
-        if (changed) {
-            state.selectedRowsKeys = updateSelectedRowsKeysFromSelectedRows(
-                selectedRows,
-                keyField
-            );
+        for (let i = 0; i < selectedRowCount; i += 1) {
+            const row = selectedRows[i];
+            if (!selectedRowsKeys[row[keyField]]) {
+                changed = true;
+                // Untracked state change.
+                state.selectedRowsKeys = updateSelectedRowsKeysFromSelectedRows(
+                    selectedRows,
+                    keyField
+                );
+                break;
+            }
         }
     }
-    if (maxRowSelection > 1 && changed) {
-        if (selectedRows.length < maxRowSelection) {
-            markDeselectedRowEnabled(state);
-        } else {
-            markDeselectedRowDisabled(state);
+    if (changed) {
+        const total = state.maxRowSelection || state.rows.length;
+        if (total > 1) {
+            // Tracked state changes.
+            if (selectedRowCount < total) {
+                markDeselectedRowEnabled(state);
+            } else {
+                markDeselectedRowDisabled(state);
+            }
         }
     }
+    // Tracked state change.
     updateBulkSelectionState(state);
-
     return {
-        ifChanged: (callback) => {
+        ifChanged(callback) {
             if (changed && typeof callback === 'function') {
                 callback(selectedRows);
             }
@@ -484,35 +450,35 @@ export function syncSelectedRowsKeys(state, selectedRows) {
 }
 
 function updateSelectedRowsKeysFromSelectedRows(selectedRows, keyField) {
-    return selectedRows.reduce((selectedRowsKeys, row) => {
+    const selectedRowsKeys = {};
+    for (let i = 0, { length } = selectedRows; i < length; i += 1) {
+        const row = selectedRows[i];
         selectedRowsKeys[row[keyField]] = true;
-        return selectedRowsKeys;
-    }, {});
-}
-
-function addKeyToSelectedRowKeys(state, key) {
-    state.selectedRowsKeys[key] = true;
-}
-
-function removeKeyFromSelectedRowKeys(state, key) {
-    // not using delete this.state.selectedRowsKeys[key]
-    // because that causes perf issues
-    state.selectedRowsKeys[key] = false;
+    }
+    return selectedRowsKeys;
 }
 
-function normalizeSelectedRowsKey(value) {
-    return value.reduce((map, key) => {
-        map[key] = true;
-        return map;
-    }, {});
+function normalizeSelectedRowsKey(keys) {
+    const selectedRowsKeys = {};
+    for (let i = 0, { length } = keys; i < length; i += 1) {
+        selectedRowsKeys[keys[i]] = true;
+    }
+    return selectedRowsKeys;
 }
 
 function filterValidKeys(state, keys) {
-    return keys.filter((key) => rowKeyExists(state, key));
-}
-
-export function resetSelectedRowsKeys(state) {
-    state.selectedRowsKeys = {};
+    const filtered = [];
+    const { rows } = state;
+    if (rows.length) {
+        const { indexes } = state;
+        for (let i = 0, { length } = keys; i < length; i += 1) {
+            const key = keys[i];
+            if (indexes[key]) {
+                filtered.push(key);
+            }
+        }
+    }
+    return filtered;
 }
 
 /************************** LAST SELECTED ROW KEYS **************************/
@@ -521,11 +487,11 @@ export function resetSelectedRowsKeys(state) {
  * Returns the row key value of the row that was last selected
  * Returns undefined if the row key value is invalid
  *
- * @param {Object} state - the datatable state.
+ * @param {Object} state - The datatable state.
  * @returns {String | undefined } the row key or undefined.
  */
 function getLastRowSelection(state) {
-    const lastSelectedRowKey = state.lastSelectedRowKey;
+    const { lastSelectedRowKey } = state;
     const keyIsValid =
         lastSelectedRowKey !== undefined &&
         getRowIndexByKey(state, lastSelectedRowKey) !== undefined;
@@ -533,10 +499,6 @@ function getLastRowSelection(state) {
     return keyIsValid ? lastSelectedRowKey : undefined;
 }
 
-function setLastRowSelection(state, rowKeyValue) {
-    state.lastSelectedRowKey = rowKeyValue;
-}
-
 /************************** INPUT TYPES **************************/
 
 /**
@@ -551,30 +513,35 @@ function setLastRowSelection(state, rowKeyValue) {
  *
  * @param {Number} previousMaxRowSelection
  * @param {Number} newMaxRowSelection
- * @returns
+ * @returns {Boolean}
  */
 export function inputTypeNeedsToChange(
     previousMaxRowSelection,
     newMaxRowSelection
 ) {
     return (
-        (previousMaxRowSelection === 1 &&
-            isMultiSelection(newMaxRowSelection)) ||
-        (isMultiSelection(previousMaxRowSelection) &&
-            newMaxRowSelection === 1) ||
         previousMaxRowSelection === 0 ||
-        newMaxRowSelection === 0
+        newMaxRowSelection === 0 ||
+        (previousMaxRowSelection === 1 &&
+            // is newMaxRowSelection multi-selection
+            (newMaxRowSelection > 1 || newMaxRowSelection === undefined)) ||
+        (newMaxRowSelection === 1 &&
+            // is previousMaxRowSelection multi-selection
+            (previousMaxRowSelection > 1 ||
+                previousMaxRowSelection === undefined))
     );
 }
 
 export function updateRowSelectionInputType(state) {
-    const type = getRowSelectionInputType(state);
-    const rows = getRows(state);
+    const type = state.maxRowSelection === 1 ? 'radio' : 'checkbox';
+    const { rows } = state;
 
-    rows.forEach((row) => {
+    for (let i = 0, { length: rowCount } = rows; i < rowCount; i += 1) {
+        // Tracked state changes.
+        const row = rows[i];
         row.inputType = type;
         row.isDisabled = isDisabledRow(state, row.key);
-    });
+    }
 }
 
 /************************** MAX ROW SELECTION **************************/
@@ -587,34 +554,36 @@ export function updateRowSelectionInputType(state) {
  * and input type is radio if maxRowSelection = 1.
  * Invalid values are set to default and an error is logged
  *
- * @param {Object} state - the datatable state.
+ * @param {Object} state - The datatable state
  * @param {Number | String} - value to set for maxRowSelection
  */
 export function setMaxRowSelection(state, value) {
     const previousSelectedRowsKeys = getSelectedRowsKeys(state);
+    // Tracked state changes.
     markAllRowsDeselected(state);
     if (isNonNegativeInteger(value)) {
-        const previousMaxRowSelection = getMaxRowSelection(state);
-        state.maxRowSelection = Number(value);
-        const newMaxRowSelection = getMaxRowSelection(state);
-        // reselect up to maxRowSelection rows
+        const previousMaxRowSelection = state.maxRowSelection;
+        const newMaxRowSelection = Number(value);
+        // Untracked state change.
+        state.maxRowSelection = newMaxRowSelection;
+        // Reselect up to maxRowSelection rows.
         const numberOfRows = Math.min(
             previousSelectedRowsKeys.length,
             newMaxRowSelection
         );
-        for (let i = 0; i < numberOfRows; i++) {
+        for (let i = 0; i < numberOfRows; i += 1) {
+            // Tracked state changes.
             markRowSelected(state, previousSelectedRowsKeys[i]);
         }
         if (
-            inputTypeNeedsToChange(
-                previousMaxRowSelection,
-                getMaxRowSelection(state)
-            )
+            inputTypeNeedsToChange(previousMaxRowSelection, newMaxRowSelection)
         ) {
+            // Tracked state changes.
             updateRowSelectionInputType(state);
             updateBulkSelectionState(state);
         }
     } else {
+        // Untracked state change.
         state.maxRowSelection = MAX_ROW_SELECTION_DEFAULT;
         // suppress console error if no value is passed in
         if (value !== null && value !== undefined) {
@@ -629,32 +598,31 @@ export function setMaxRowSelection(state, value) {
 /************************** BULK SELECTION STATE **************************/
 
 export function updateBulkSelectionState(state) {
-    const selectBoxesColumnIndex = getSelectBoxesColumnIndex(state);
-    if (selectBoxesColumnIndex >= 0) {
-        state.columns[selectBoxesColumnIndex] = Object.assign(
-            {},
-            state.columns[selectBoxesColumnIndex],
-            {
-                bulkSelection: getBulkSelectionState(state),
-                isBulkSelectionDisabled: isBulkSelectionDisabled(state),
+    const columns = state.columns || {};
+    for (let colIndex = 0; colIndex < columns.length; colIndex += 1) {
+        const col = columns[colIndex];
+        if (col.type === SELECTABLE_HEADER_TYPE) {
+            const {
+                maxRowSelection,
+                rows: { length: rowCount },
+            } = state;
+            const total = maxRowSelection || rowCount;
+            const selected = getCurrentSelectionLength(state);
+            // Force a rerender of this column by replacing the tracked object.
+            // Required to pass treegrid unit tests.
+            const updatedCol = Object.assign({}, col);
+            if (selected) {
+                updatedCol.bulkSelection = selected === total ? 'all' : 'some';
+            } else {
+                updatedCol.bulkSelection = 'none';
             }
-        );
-    }
-}
-
-function getBulkSelectionState(state) {
-    const selected = getCurrentSelectionLength(state);
-    const total = getMaxRowSelection(state) || getRowsTotal(state);
-    if (selected === 0) {
-        return 'none';
-    } else if (selected === total) {
-        return 'all';
+            updatedCol.isBulkSelectionDisabled =
+                maxRowSelection === 0 || rowCount === 0;
+            // Replace old column with shallow cloned and updated column.
+            columns[colIndex] = updatedCol;
+            return;
+        }
     }
-    return 'some';
-}
-
-function isBulkSelectionDisabled(state) {
-    return getRowsTotal(state) === 0 || getMaxRowSelection(state) === 0;
 }
 
 /************************** HELPER FUNCTIONS **************************/
@@ -665,14 +633,16 @@ function isBulkSelectionDisabled(state) {
  * Returns an object that contains the start index which is the lower index value of the two
  * and the end index which is the higher value of the two.
  *
- * @param {Object} state - datatable's state object
- * @param {String} startRowKey - row key value of the first row that was selected (start of the interval)
- * @param {String} endRowKey - row key value of the last row that was selected (end of the interval)
- * @returns {Object} - object with start index and end index
+ * @param {Object} state - The datatable state
+ * @param {String} startRowKey - The row key value of the first row that was selected (start of the interval)
+ * @param {String} endRowKey - The row key value of the last row that was selected (end of the interval)
+ * @returns {Object} - An object with start index and end index
  */
 function getRowIntervalIndexes(state, startRowKey, endRowKey) {
     const start =
-        startRowKey === 'HEADER' ? 0 : getRowIndexByKey(state, startRowKey);
+        startRowKey === HEADER_ROW_KEY
+            ? 0
+            : getRowIndexByKey(state, startRowKey);
     const end = getRowIndexByKey(state, endRowKey);
 
     return {
@@ -687,9 +657,9 @@ function getRowIntervalIndexes(state, startRowKey, endRowKey) {
  *
  * Note: This change is volatile, and will be reset (lost) in the next index regeneration.
  *
- * @param {Object} state - the state of the datatable
- * @param {String} rowKeyValue - the row key of the cell to mark selected
- * @param {String} colKeyValue - the col key of the cell to mark selected
+ * @param {Object} state - The datatable state
+ * @param {String} rowKeyValue - The row key of the cell to mark selected
+ * @param {String} colKeyValue - The column key of the cell to mark selected
  */
 export function setAriaSelectedOnCell(state, rowKeyValue, colKeyValue) {
     const row = getRowByKey(state, rowKeyValue);
@@ -707,9 +677,9 @@ export function setAriaSelectedOnCell(state, rowKeyValue, colKeyValue) {
  *
  * Note: This change is volatile, and will be reset (lost) in the next index regeneration.
  *
- * @param {Object} state - the state of the datatable
- * @param {String} rowKeyValue - the row key of the cell to select
- * @param {String} colKeyValue - the col key of the cell to select
+ * @param {Object} state - The datatable state
+ * @param {String} rowKeyValue - The row key of the cell to select
+ * @param {String} colKeyValue - The column key of the cell to select
  */
 export function unsetAriaSelectedOnCell(state, rowKeyValue, colKeyValue) {
     const row = getRowByKey(state, rowKeyValue);
@@ -725,63 +695,38 @@ export function unsetAriaSelectedOnCell(state, rowKeyValue, colKeyValue) {
  * to one which reflects the selected value of the row on the row object.
  * These are used by the template to render the appropriate values.
  *
- * @param {Boolean} selectedValue - is the row selected or not
- * @param {Object} row - the row on which to set the selected attributes
+ * @param {Object} row - The row
+ * @param {Boolean} selectedValue - Whether the row is selected
  */
-function setRowSelectedAttributes(selectedValue, row) {
-    row.isSelected = selectedValue;
+function setRowSelectedAttributes(row, selectedValue) {
     row.ariaSelected = selectedValue;
-    row.classnames = resolveRowClassNames(row);
+    row.classnames = `slds-hint-parent${
+        selectedValue ? ' slds-is-selected' : ''
+    }`;
+    row.isSelected = selectedValue;
 }
 
-function getSelectedDiff(state, value) {
-    const selectedRowsKeys = state.selectedRowsKeys;
-    return value.filter((key) => !selectedRowsKeys[key]);
+function getSelectedDiff(state, selectedRows) {
+    const filtered = [];
+    const { selectedRowsKeys } = state;
+    for (let i = 0, { length } = selectedRows; i < length; i += 1) {
+        const rowKeyValue = selectedRows[i];
+        if (!selectedRowsKeys[rowKeyValue]) {
+            filtered.push(rowKeyValue);
+        }
+    }
+    return filtered;
 }
 
 function getDeselectedDiff(state, value) {
-    const currentSelectedRowsKeys = state.selectedRowsKeys;
-    return Object.keys(currentSelectedRowsKeys).filter(
-        (key) => currentSelectedRowsKeys[key] && !value[key]
-    );
-}
-
-function getSelectBoxesColumnIndex(state) {
-    const columns = getColumns(state) || [];
-    let selectBoxColumnIndex = -1;
-
-    columns.some((column, index) => {
-        if (column.type === SELECTABLE_ROW_CHECKBOX) {
-            selectBoxColumnIndex = index;
-            return true;
+    const filtered = [];
+    const { selectedRowsKeys } = state;
+    const keys = Object.keys(selectedRowsKeys);
+    for (let i = 0, { length } = keys; i < length; i += 1) {
+        const rowKeyValue = keys[i];
+        if (selectedRowsKeys[rowKeyValue] && !value[rowKeyValue]) {
+            filtered.push(rowKeyValue);
         }
-
-        return false;
-    });
-
-    return selectBoxColumnIndex;
-}
-
-/**
- * Represents whether the select all rows checkbox on the header
- * should be visible or not.
- * If max-row-selection = 1, then the select all rows checkbox
- * should not be visible/rendered.
- *
- * @param {Object} state - datatable's state object
- * @returns
- */
-export function getHideSelectAllCheckbox(state) {
-    return getMaxRowSelection(state) === 1;
-}
-
-/**
- * Represents whether the datatable should enable multiselection or not
- * depending on the max-row-selection value passed in
- *
- * @param {Number} value - max-row-selection value
- * @returns
- */
-function isMultiSelection(value) {
-    return value > 1 || value === undefined;
+    }
+    return filtered;
 }