Refactor hunk expansion info computation

app/src/lib/diff-parser.ts

@@ -6,6 +6,7 @@ import {
   DiffLineType,
 } from '../models/diff'
 import { assertNever } from '../lib/fatal-error'
+import { getHunkHeaderExpansionType } from '../ui/diff/text-diff-expansion'
 
 // https://en.wikipedia.org/wiki/Diff_utility
 //
@@ -280,7 +281,11 @@ export class DiffParser {
    *                      have no real meaning in the context of a diff and
    *                      are only used to aid the app in line-selections.
    */
-  private parseHunk(linesConsumed: number): DiffHunk {
+  private parseHunk(
+    linesConsumed: number,
+    hunkIndex: number,
+    previousHunk: DiffHunk | null
+  ): DiffHunk {
     const headerLine = this.readLine()
     if (!headerLine) {
       throw new Error('Expected hunk header but reached end of diff')
@@ -366,7 +371,8 @@ export class DiffParser {
       header,
       lines,
       linesConsumed,
-      linesConsumed + lines.length - 1
+      linesConsumed + lines.length - 1,
+      getHunkHeaderExpansionType(hunkIndex, header, previousHunk)
     )
   }
 
@@ -397,10 +403,12 @@ export class DiffParser {
 
       const hunks = new Array<DiffHunk>()
       let linesConsumed = 0
+      let previousHunk: DiffHunk | null = null
 
       do {
-        const hunk = this.parseHunk(linesConsumed)
+        const hunk = this.parseHunk(linesConsumed, hunks.length, previousHunk)
         hunks.push(hunk)
+        previousHunk = hunk
         linesConsumed += hunk.lines.length
       } while (this.peek())
 

app/src/models/diff/raw-diff.ts

@@ -1,5 +1,31 @@
 import { DiffLine } from './diff-line'
 
+export enum DiffHunkExpansionType {
+  /** The hunk header cannot be expanded at all. */
+  None = 'None',
+
+  /**
+   * The hunk header can be expanded up exclusively. Only the first hunk can be
+   * expanded up exclusively.
+   */
+  Up = 'Up',
+
+  /**
+   * The hunk header can be expanded down exclusively. Only the last hunk (if
+   * it's the dummy hunk with only one line) can be expanded down exclusively.
+   */
+  Down = 'Down',
+
+  /** The hunk header can be expanded both up and down. */
+  Both = 'Both',
+
+  /**
+   * The hunk header represents a short gap that, when expanded, will
+   * result in merging this hunk and the hunk above.
+   */
+  Short = 'Short',
+}
+
 /** each diff is made up of a number of hunks */
 export class DiffHunk {
   /**
@@ -12,7 +38,8 @@ export class DiffHunk {
     public readonly header: DiffHunkHeader,
     public readonly lines: ReadonlyArray<DiffLine>,
     public readonly unifiedDiffStart: number,
-    public readonly unifiedDiffEnd: number
+    public readonly unifiedDiffEnd: number,
+    public readonly expansionType: DiffHunkExpansionType
   ) {}
 }
 

app/src/ui/diff/diff-helpers.tsx

@@ -8,6 +8,7 @@ import {
   WorkingDirectoryFileChange,
   CommittedFileChange,
 } from '../../models/status'
+import { DiffHunkExpansionType } from '../../models/diff/raw-diff'
 
 /**
  * DiffRowType defines the different types of
@@ -181,6 +182,9 @@ interface IDiffRowHunk {
    * The actual contents of the line.
    */
   readonly content: string
+
+  /** How the hunk can be expanded. */
+  readonly expansionType: DiffHunkExpansionType
 }
 
 export type DiffRow =

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

@@ -864,7 +864,11 @@ function getDiffRowsFromHunk(
     }
 
     if (line.type === DiffLineType.Hunk) {
-      rows.push({ type: DiffRowType.Hunk, content: line.content })
+      rows.push({
+        type: DiffRowType.Hunk,
+        content: line.content,
+        expansionType: hunk.expansionType,
+      })
       continue
     }
 

app/src/ui/diff/text-diff-expansion.ts

@@ -1,5 +1,7 @@
+import { enableTextDiffExpansion } from '../../lib/feature-flag'
 import {
   DiffHunk,
+  DiffHunkExpansionType,
   DiffHunkHeader,
   DiffLine,
   DiffLineType,
@@ -12,24 +14,6 @@ export const DiffExpansionStep = 20
 /** Type of expansion: could be up or down. */
 export type ExpansionKind = 'up' | 'down'
 
-/** Interface to represent the expansion capabilities of a hunk header. */
-interface IHunkHeaderExpansionInfo {
-  /** True if the hunk header can be expanded down. */
-  readonly isExpandableDown: boolean
-
-  /** True if the hunk header can be expanded up. */
-  readonly isExpandableUp: boolean
-
-  /** True if the hunk header can be expanded both up and down. */
-  readonly isExpandableBoth: boolean
-
-  /**
-   * True if the hunk header represents a short gap that, when expanded, will
-   * result in merging this hunk and the hunk above.
-   */
-  readonly isExpandableShort: boolean
-}
-
 /** Builds the diff text string given a list of hunks. */
 function getDiffTextFromHunks(hunks: ReadonlyArray<DiffHunk>) {
   // Grab all hunk lines and rebuild the diff text from it
@@ -70,7 +54,18 @@ function mergeDiffHunks(hunk1: DiffHunk, hunk2: DiffHunk): DiffHunk {
     [newFirstHunkLine, ...allHunk1LinesButFirst, ...allHunk2LinesButFirst],
     hunk1.unifiedDiffStart,
     // This -1 represents the header line of the second hunk that we removed
-    hunk2.unifiedDiffEnd - 1
+    hunk2.unifiedDiffEnd - 1,
+    // The expansion type of the resulting hunk will match the expansion type
+    // of the first hunk:
+    // - If the first hunk can be expanded up, it means it's the very first
+    //   hunk, so the resulting hunk will be the first too.
+    // - If the first hunk can be expanded but short, that doesn't change after
+    //   merging it with the second one.
+    // - If it can be expanded up and down (meaning it's a long gap), that
+    //   doesn't change after merging it with the second one.
+    // - It can never be expanded down exclusively, because only the last dummy
+    //   hunk can do that, and that will never be the first hunk in a merge.
+    hunk1.expansionType
   )
 }
 
@@ -79,24 +74,24 @@ function mergeDiffHunks(hunk1: DiffHunk, hunk2: DiffHunk): DiffHunk {
  * the space represented by the hunk header is short and expansion there would
  * mean merging with the hunk above.
  *
- * @param hunks All hunks in the diff.
- * @param hunk  The specific hunk for which we want to know the expansion info.
+ * @param hunkIndex     Index of the hunk to evaluate within the whole diff.
+ * @param hunkHeader    Header of the hunk to evaluate.
+ * @param previousHunk  Hunk previous to the one to evaluate. Null if the
+ *                      evaluated hunk is the first one.
  */
-export function getHunkHeaderExpansionInfo(
-  hunks: ReadonlyArray<DiffHunk>,
-  hunk: DiffHunk
-): IHunkHeaderExpansionInfo {
-  let isExpandableDown = false
-  let isExpandableUp = false
-  let isExpandableBoth = false
-  let isExpandableShort = false
+export function getHunkHeaderExpansionType(
+  hunkIndex: number,
+  hunkHeader: DiffHunkHeader,
+  previousHunk: DiffHunk | null
+): DiffHunkExpansionType {
+  if (!enableTextDiffExpansion()) {
+    return DiffHunkExpansionType.None
+  }
 
-  const hunkIndex = hunks.indexOf(hunk)
-  const previousHunk = hunks[hunkIndex - 1]
   const distanceToPrevious =
-    previousHunk === undefined
+    previousHunk === null
       ? Infinity
-      : hunk.header.oldStartLine -
+      : hunkHeader.oldStartLine -
         previousHunk.header.oldStartLine -
         previousHunk.header.oldLineCount
 
@@ -108,21 +103,15 @@ export function getHunkHeaderExpansionInfo(
   // short and therefore the direction of expansion doesn't matter.
   if (hunkIndex === 0) {
     // The top hunk can only be expanded if there is content above it
-    isExpandableUp =
-      hunk.header.oldStartLine > 1 && hunk.header.newStartLine > 1
-  } else if (hunkIndex === hunks.length - 1 && hunk.lines.length === 1) {
-    isExpandableDown = true
+    if (hunkHeader.oldStartLine > 1 && hunkHeader.newStartLine > 1) {
+      return DiffHunkExpansionType.Up
+    } else {
+      return DiffHunkExpansionType.None
+    }
   } else if (distanceToPrevious <= DiffExpansionStep) {
-    isExpandableShort = true
+    return DiffHunkExpansionType.Short
   } else {
-    isExpandableBoth = true
-  }
-
-  return {
-    isExpandableDown,
-    isExpandableUp,
-    isExpandableBoth,
-    isExpandableShort,
+    return DiffHunkExpansionType.Both
   }
 }
 
@@ -270,12 +259,20 @@ export function expandTextDiffHunk(
 
   let numberOfNewDiffLines = updatedHunkLines.length - hunk.lines.length
 
+  const previousHunk = hunkIndex === 0 ? null : diff.hunks[hunkIndex - 1]
+  const expansionType = getHunkHeaderExpansionType(
+    hunkIndex,
+    newHunkHeader,
+    previousHunk
+  )
+
   // Update the hunk with all the new info (header, lines, start/end...)
   let updatedHunk = new DiffHunk(
     newHunkHeader,
     updatedHunkLines,
     hunk.unifiedDiffStart,
-    hunk.unifiedDiffEnd + numberOfNewDiffLines
+    hunk.unifiedDiffEnd + numberOfNewDiffLines,
+    expansionType
   )
 
   let previousHunksEndIndex = 0 // Exclusive
@@ -312,17 +309,33 @@ export function expandTextDiffHunk(
   const followingHunks =
     newHunkLastLine >= newContentLines.length
       ? []
-      : diff.hunks
-          .slice(followingHunksStartIndex)
-          .map(
-            hunk =>
-              new DiffHunk(
-                hunk.header,
-                hunk.lines,
-                hunk.unifiedDiffStart + numberOfNewDiffLines,
-                hunk.unifiedDiffEnd + numberOfNewDiffLines
-              )
+      : diff.hunks.slice(followingHunksStartIndex).map((hunk, hunkIndex) => {
+          const isLastDummyHunk =
+            hunkIndex + followingHunksStartIndex === diff.hunks.length - 1 &&
+            hunk.lines.length === 1 &&
+            hunk.lines[0].type === DiffLineType.Hunk
+
+          // Only compute the new expansion type if the hunk is the first one
+          // (of the remaining hunks) and it's not the last dummy hunk.
+          const shouldComputeNewExpansionType =
+            hunkIndex === 0 && !isLastDummyHunk
+
+          return new DiffHunk(
+            hunk.header,
+            hunk.lines,
+            hunk.unifiedDiffStart + numberOfNewDiffLines,
+            hunk.unifiedDiffEnd + numberOfNewDiffLines,
+            // If it's the first hunk after the one we expanded, recalculate
+            // its expansion type.
+            shouldComputeNewExpansionType
+              ? getHunkHeaderExpansionType(
+                  followingHunksStartIndex,
+                  hunk.header,
+                  updatedHunk
+                )
+              : hunk.expansionType
           )
+        })
 
   // Create the new list of hunks of the diff, and the new diff text
   const newHunks = [...previousHunks, updatedHunk, ...followingHunks]
@@ -386,7 +399,8 @@ export function getTextDiffWithBottomDummyHunk(
     dummyHeader,
     [dummyLine],
     lastHunk.unifiedDiffEnd + 1,
-    lastHunk.unifiedDiffEnd + 1
+    lastHunk.unifiedDiffEnd + 1,
+    DiffHunkExpansionType.Down
   )
 
   const newHunks = [...hunks, dummyHunk]

app/src/ui/diff/text-diff.tsx

@@ -8,6 +8,7 @@ import {
   DiffSelection,
   DiffLine,
   ITextDiff,
+  DiffHunkExpansionType,
 } from '../../models/diff'
 import {
   WorkingDirectoryFileChange,
@@ -50,7 +51,6 @@ import { canSelect } from './diff-helpers'
 import {
   expandTextDiffHunk,
   ExpansionKind,
-  getHunkHeaderExpansionInfo,
   getTextDiffWithBottomDummyHunk,
 } from './text-diff-expansion'
 import { createOcticonElement } from '../octicons/octicon'
@@ -338,6 +338,7 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
   /** The current, active, diff gutter selection if any */
   private selection: ISelection | null = null
 
+  /** The content lines of the "new" file */
   private newContentLines: ReadonlyArray<string> | null = null
 
   /** Whether a particular range should be highlighted due to hover */
@@ -909,7 +910,7 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
           ) {
             const marker = lineInfo.gutterMarkers[diffGutterName]
             if (marker instanceof HTMLElement) {
-              this.updateGutterMarker(marker, hunks, hunk, diffLine)
+              this.updateGutterMarker(marker, hunk, diffLine)
             }
           } else {
             batchedOps.push(() => {
@@ -949,7 +950,6 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
   }
 
   private getGutterLineClassNameInfo(
-    hunks: ReadonlyArray<DiffHunk>,
     hunk: DiffHunk,
     diffLine: DiffLine
   ): { [className: string]: boolean } {
@@ -962,10 +962,6 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
       isIncludeable &&
       diffLine.originalLineNumber !== null &&
       inSelection(this.hunkHighlightRange, diffLine.originalLineNumber)
-    const hunkExpansionInfo =
-      enableTextDiffExpansion() && diffLine.type === DiffLineType.Hunk
-        ? getHunkHeaderExpansionInfo(hunks, hunk)
-        : undefined
 
     return {
       'diff-line-gutter': true,
@@ -976,10 +972,10 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
       'read-only': this.props.readOnly,
       'diff-line-selected': isIncluded,
       'diff-line-hover': hover,
-      'expandable-down': hunkExpansionInfo?.isExpandableDown === true,
-      'expandable-up': hunkExpansionInfo?.isExpandableUp === true,
-      'expandable-both': hunkExpansionInfo?.isExpandableBoth === true,
-      'expandable-short': hunkExpansionInfo?.isExpandableShort === true,
+      'expandable-down': hunk.expansionType === DiffHunkExpansionType.Down,
+      'expandable-up': hunk.expansionType === DiffHunkExpansionType.Up,
+      'expandable-both': hunk.expansionType === DiffHunkExpansionType.Both,
+      'expandable-short': hunk.expansionType === DiffHunkExpansionType.Short,
       includeable: isIncludeable && !this.props.readOnly,
     }
   }
@@ -1074,7 +1070,7 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
       )
     }
 
-    this.updateGutterMarker(marker, hunks, hunk, diffLine)
+    this.updateGutterMarker(marker, hunk, diffLine)
 
     return marker
   }
@@ -1157,11 +1153,10 @@ export class TextDiff extends React.Component<ITextDiffProps, ITextDiffState> {
 
   private updateGutterMarker(
     marker: HTMLElement,
-    hunks: ReadonlyArray<DiffHunk>,
     hunk: DiffHunk,
     diffLine: DiffLine
   ) {
-    const classNameInfo = this.getGutterLineClassNameInfo(hunks, hunk, diffLine)
+    const classNameInfo = this.getGutterLineClassNameInfo(hunk, diffLine)
     for (const [className, include] of Object.entries(classNameInfo)) {
       if (include) {
         marker.classList.add(className)