stuff

app/src/ui/diff/side-by-side-diff-row.tsx

@@ -136,11 +136,11 @@ interface ISideBySideDiffRowProps {
   /** Called when the user changes the hide whitespace in diffs setting. */
   readonly onHideWhitespaceInDiffChanged: (checked: boolean) => void
 
-  /* This tracks the last expanded hunk index so that we can refocus the expander after rerender */
-  readonly lastExpandedHunk: {
-    index: number
-    expansionType: DiffHunkExpansionType
-  } | null
+  readonly onHunkExpansionRef: (
+    hunkIndex: number,
+    expansionType: DiffHunkExpansionType,
+    element: HTMLButtonElement | null
+  ) => void
 }
 
 interface ISideBySideDiffRowState {
@@ -422,25 +422,6 @@ export class SideBySideDiffRow extends React.Component<
       expansionType
     )
 
-    /**
-     * For accessibility, when a button is focused, it should maintain focus.
-     * This sets the autofocus of the button if the last expanded button at the
-     * position was the same type. The +1 is to handle the last hunk index which
-     * is one off, and if there are two hunks with the same expansion types on
-     * after each other we just want the first one and autofocus will go to the
-     * first one automatically.
-     *
-     * Other notes: the expand up buttons already worked. This is
-     * for expand all and expand down buttons.
-     */
-    const { lastExpandedHunk } = this.props
-    const focusButton =
-      lastExpandedHunk !== null &&
-      expansionType === lastExpandedHunk.expansionType
-        ? hunkIndex === lastExpandedHunk.index ||
-          hunkIndex === lastExpandedHunk.index + 1
-        : false
-
     return (
       <div
         className="hunk-expansion-handle selectable hoverable"
@@ -451,8 +432,8 @@ export class SideBySideDiffRow extends React.Component<
           onContextMenu={this.props.onContextMenuExpandHunk}
           tooltip={elementInfo.title}
           toolTipDirection={TooltipDirection.SOUTH}
-          autoFocus={focusButton}
           ariaLabel={elementInfo.title}
+          onButtonRef={this.onHunkExpansionRef(hunkIndex, expansionType)}
         >
           <Octicon symbol={elementInfo.icon} />
         </Button>
@@ -460,6 +441,12 @@ export class SideBySideDiffRow extends React.Component<
     )
   }
 
+  private onHunkExpansionRef =
+    (hunkIndex: number, expansionType: DiffHunkExpansionType) =>
+    (button: HTMLButtonElement | null) => {
+      this.props.onHunkExpansionRef(hunkIndex, expansionType, button)
+    }
+
   private renderHunkHeaderGutter(
     hunkIndex: number,
     expansionType: DiffHunkExpansionType

app/src/ui/diff/side-by-side-diff.tsx

@@ -63,6 +63,7 @@ import {
 import { IMenuItem } from '../../lib/menu-item'
 import { HiddenBidiCharsWarning } from './hidden-bidi-chars-warning'
 import { escapeRegExp } from 'lodash'
+import { findDOMNode } from 'react-dom'
 
 const DefaultRowHeight = 20
 
@@ -227,6 +228,11 @@ export class SideBySideDiff extends React.Component<
   private textSelectionStartRow: number | undefined = undefined
   private textSelectionEndRow: number | undefined = undefined
 
+  private readonly hunkExpansionRefs = new Map<
+    number,
+    { expansionType: DiffHunkExpansionType; button: HTMLButtonElement }
+  >()
+
   public constructor(props: ISideBySideDiffProps) {
     super(props)
 
@@ -412,6 +418,41 @@ export class SideBySideDiff extends React.Component<
         }
       }
     }
+
+    if (this.state.lastExpandedHunk !== prevState.lastExpandedHunk) {
+      this.focusAfterLastExpandedHunkChange()
+    }
+  }
+
+  private focusAfterLastExpandedHunkChange() {
+    if (this.state.lastExpandedHunk === null) {
+      return
+    }
+
+    const diffNode = findDOMNode(this.virtualListRef.current)
+    const diff = diffNode instanceof HTMLElement ? diffNode : null
+
+    if (this.hunkExpansionRefs.size === 0) {
+      diff?.focus()
+      return
+    }
+
+    const { expansionType, index } = this.state.lastExpandedHunk
+
+    const hunk = this.hunkExpansionRefs.get(index)
+    const lastHunk = this.hunkExpansionRefs.get(this.hunkExpansionRefs.size - 1)
+
+    if (hunk?.expansionType === expansionType) {
+      hunk?.button.focus()
+      return
+    }
+
+    if (lastHunk?.expansionType === expansionType) {
+      lastHunk?.button.focus()
+      return
+    }
+
+    diff?.focus()
   }
 
   private canExpandDiff() {
@@ -587,13 +628,25 @@ export class SideBySideDiff extends React.Component<
             }
             beforeClassNames={beforeClassNames}
             afterClassNames={afterClassNames}
-            lastExpandedHunk={this.state.lastExpandedHunk}
+            onHunkExpansionRef={this.onHunkExpansionRef}
           />
         </div>
       </CellMeasurer>
     )
   }
 
+  private onHunkExpansionRef = (
+    hunkIndex: number,
+    expansionType: DiffHunkExpansionType,
+    button: HTMLButtonElement | null
+  ) => {
+    if (button === null) {
+      this.hunkExpansionRefs.delete(hunkIndex)
+    } else {
+      this.hunkExpansionRefs.set(hunkIndex, { expansionType, button })
+    }
+  }
+
   private getRowHeight = (row: { index: number }) => {
     return listRowsHeightCache.rowHeight(row) ?? DefaultRowHeight
   }