3 files changed, 73 insertions, 16 deletions

diff --git a/app/assets/javascripts/diffs/components/diff_expansion_cell.vue b/app/assets/javascripts/diffs/components/diff_expansion_cell.vue
index 23fbfc2b74b..4eae2e09c08 100644
--- a/app/assets/javascripts/diffs/components/diff_expansion_cell.vue
+++ b/app/assets/javascripts/diffs/components/diff_expansion_cell.vue
@@ -3,7 +3,7 @@ import { mapState, mapActions } from 'vuex';
 import createFlash from '~/flash';
 import { s__ } from '~/locale';
 import Icon from '~/vue_shared/components/icon.vue';
-import { UNFOLD_COUNT } from '../constants';
+import { UNFOLD_COUNT, INLINE_DIFF_VIEW_TYPE, PARALLEL_DIFF_VIEW_TYPE } from '../constants';
 import * as utils from '../store/utils';
 import tooltip from '../../vue_shared/directives/tooltip';
 
@@ -11,6 +11,16 @@ const EXPAND_ALL = 0;
 const EXPAND_UP = 1;
 const EXPAND_DOWN = 2;
 
+const lineNumberByViewType = (viewType, diffLine) => {
+  const numberGetters = {
+    [INLINE_DIFF_VIEW_TYPE]: line => line?.new_line,
+    [PARALLEL_DIFF_VIEW_TYPE]: line => (line?.right || line?.left)?.new_line,
+  };
+  const numberGetter = numberGetters[viewType];
+
+  return numberGetter && numberGetter(diffLine);
+};
+
 export default {
   directives: {
     tooltip,
@@ -67,12 +77,16 @@ export default {
     ...mapActions('diffs', ['loadMoreLines']),
     getPrevLineNumber(oldLineNumber, newLineNumber) {
       const diffFile = utils.findDiffFile(this.diffFiles, this.fileHash);
-      const indexForInline = utils.findIndexInInlineLines(diffFile.highlighted_diff_lines, {
+      const lines = {
+        [INLINE_DIFF_VIEW_TYPE]: diffFile.highlighted_diff_lines,
+        [PARALLEL_DIFF_VIEW_TYPE]: diffFile.parallel_diff_lines,
+      };
+      const index = utils.getPreviousLineIndex(this.diffViewType, diffFile, {
         oldLineNumber,
         newLineNumber,
       });
-      const prevLine = diffFile.highlighted_diff_lines[indexForInline - 2];
-      return (prevLine && prevLine.new_line) || 0;
+
+      return lineNumberByViewType(this.diffViewType, lines[this.diffViewType][index - 2]) || 0;
     },
     callLoadMoreLines(
       endpoint,
@@ -114,7 +128,7 @@ export default {
         this.handleExpandAllLines(expandOptions);
       }
     },
-    handleExpandUpLines(expandOptions = EXPAND_ALL) {
+    handleExpandUpLines(expandOptions) {
       const { endpoint, fileHash, view, oldLineNumber, newLineNumber, offset } = expandOptions;
 
       const bottom = this.isBottom;
diff --git a/app/assets/javascripts/diffs/store/mutations.js b/app/assets/javascripts/diffs/store/mutations.js
index c26411af5d7..086a7872a5d 100644
--- a/app/assets/javascripts/diffs/store/mutations.js
+++ b/app/assets/javascripts/diffs/store/mutations.js
@@ -140,6 +140,7 @@ export default {
     addContextLines({
       inlineLines: diffFile.highlighted_diff_lines,
       parallelLines: diffFile.parallel_diff_lines,
+      diffViewType: state.diffViewType,
       contextLines: lines,
       bottom,
       lineNumbers,
diff --git a/app/assets/javascripts/diffs/store/utils.js b/app/assets/javascripts/diffs/store/utils.js
index 80972d2aeb8..29133c814ea 100644
--- a/app/assets/javascripts/diffs/store/utils.js
+++ b/app/assets/javascripts/diffs/store/utils.js
@@ -13,6 +13,8 @@ import {
   LINES_TO_BE_RENDERED_DIRECTLY,
   MAX_LINES_TO_BE_RENDERED,
   TREE_TYPE,
+  INLINE_DIFF_VIEW_TYPE,
+  PARALLEL_DIFF_VIEW_TYPE,
 } from '../constants';
 
 export function findDiffFile(files, match, matchKey = 'file_hash') {
@@ -93,8 +95,7 @@ export function getNoteFormData(params) {
 export const findIndexInInlineLines = (lines, lineNumbers) => {
   const { oldLineNumber, newLineNumber } = lineNumbers;
 
-  return _.findIndex(
-    lines,
+  return lines.findIndex(
     line => line.old_line === oldLineNumber && line.new_line === newLineNumber,
   );
 };
@@ -102,8 +103,7 @@ export const findIndexInInlineLines = (lines, lineNumbers) => {
 export const findIndexInParallelLines = (lines, lineNumbers) => {
   const { oldLineNumber, newLineNumber } = lineNumbers;
 
-  return _.findIndex(
-    lines,
+  return lines.findIndex(
     line =>
       line.left &&
       line.right &&
@@ -112,13 +112,32 @@ export const findIndexInParallelLines = (lines, lineNumbers) => {
   );
 };
 
+const indexGettersByViewType = {
+  [INLINE_DIFF_VIEW_TYPE]: findIndexInInlineLines,
+  [PARALLEL_DIFF_VIEW_TYPE]: findIndexInParallelLines,
+};
+
+export const getPreviousLineIndex = (diffViewType, file, lineNumbers) => {
+  const findIndex = indexGettersByViewType[diffViewType];
+  const lines = {
+    [INLINE_DIFF_VIEW_TYPE]: file.highlighted_diff_lines,
+    [PARALLEL_DIFF_VIEW_TYPE]: file.parallel_diff_lines,
+  };
+
+  return findIndex && findIndex(lines[diffViewType], lineNumbers);
+};
+
 export function removeMatchLine(diffFile, lineNumbers, bottom) {
   const indexForInline = findIndexInInlineLines(diffFile.highlighted_diff_lines, lineNumbers);
   const indexForParallel = findIndexInParallelLines(diffFile.parallel_diff_lines, lineNumbers);
   const factor = bottom ? 1 : -1;
 
-  diffFile.highlighted_diff_lines.splice(indexForInline + factor, 1);
-  diffFile.parallel_diff_lines.splice(indexForParallel + factor, 1);
+  if (indexForInline > -1) {
+    diffFile.highlighted_diff_lines.splice(indexForInline + factor, 1);
+  }
+  if (indexForParallel > -1) {
+    diffFile.parallel_diff_lines.splice(indexForParallel + factor, 1);
+  }
 }
 
 export function addLineReferences(lines, lineNumbers, bottom, isExpandDown, nextLineNumbers) {
@@ -160,8 +179,8 @@ export function addLineReferences(lines, lineNumbers, bottom, isExpandDown, next
   return linesWithNumbers;
 }
 
-export function addContextLines(options) {
-  const { inlineLines, parallelLines, contextLines, lineNumbers, isExpandDown } = options;
+function addParallelContextLines(options) {
+  const { parallelLines, contextLines, lineNumbers, isExpandDown } = options;
   const normalizedParallelLines = contextLines.map(line => ({
     left: line,
     right: line,
@@ -170,17 +189,40 @@ export function addContextLines(options) {
   const factor = isExpandDown ? 1 : 0;
 
   if (!isExpandDown && options.bottom) {
-    inlineLines.push(...contextLines);
     parallelLines.push(...normalizedParallelLines);
   } else {
-    const inlineIndex = findIndexInInlineLines(inlineLines, lineNumbers);
     const parallelIndex = findIndexInParallelLines(parallelLines, lineNumbers);
 
-    inlineLines.splice(inlineIndex + factor, 0, ...contextLines);
     parallelLines.splice(parallelIndex + factor, 0, ...normalizedParallelLines);
   }
 }
 
+function addInlineContextLines(options) {
+  const { inlineLines, contextLines, lineNumbers, isExpandDown } = options;
+  const factor = isExpandDown ? 1 : 0;
+
+  if (!isExpandDown && options.bottom) {
+    inlineLines.push(...contextLines);
+  } else {
+    const inlineIndex = findIndexInInlineLines(inlineLines, lineNumbers);
+
+    inlineLines.splice(inlineIndex + factor, 0, ...contextLines);
+  }
+}
+
+export function addContextLines(options) {
+  const { diffViewType } = options;
+  const contextLineHandlers = {
+    [INLINE_DIFF_VIEW_TYPE]: addInlineContextLines,
+    [PARALLEL_DIFF_VIEW_TYPE]: addParallelContextLines,
+  };
+  const contextLineHandler = contextLineHandlers[diffViewType];
+
+  if (contextLineHandler) {
+    contextLineHandler(options);
+  }
+}
+
 /**
  * Trims the first char of the `richText` property when it's either a space or a diff symbol.
  * @param {Object} line