1 files changed, 35 insertions, 5 deletions
diff --git a/lib/gitlab/diff/file.rb b/lib/gitlab/diff/file.rb
index 765fb0289a8..a9e209d5b71 100644
--- a/lib/gitlab/diff/file.rb
+++ b/lib/gitlab/diff/file.rb
@@ -78,9 +78,12 @@ module Gitlab
 
       # Returns the raw diff content up to the given line index
       def diff_hunk(diff_line)
-        # Adding 2 because of the @@ diff header and Enum#take should consider
-        # an extra line, because we're passing an index.
-        raw_diff.each_line.take(diff_line.index + 2).join
+        diff_line_index = diff_line.index
+        # @@ (match) header is not kept if it's found in the top of the file,
+        # therefore we should keep an extra line on this scenario.
+        diff_line_index += 1 unless diff_lines.first.match?
+
+        diff_lines.select { |line| line.index <= diff_line_index }.map(&:text).join("\n")
       end
 
       def old_sha
@@ -127,11 +130,13 @@ module Gitlab
 
       # Array of Gitlab::Diff::Line objects
       def diff_lines
-        @diff_lines ||= Gitlab::Diff::Parser.new.parse(raw_diff.each_line).to_a
+        @diff_lines ||=
+          Gitlab::Diff::Parser.new.parse(raw_diff.each_line, diff_file: self).to_a
       end
 
       def highlighted_diff_lines
-        @highlighted_diff_lines ||= Gitlab::Diff::Highlight.new(self, repository: self.repository).highlight
+        @highlighted_diff_lines ||=
+          Gitlab::Diff::Highlight.new(self, repository: self.repository).highlight
       end
 
       # Array[<Hash>] with right/left keys that contains Gitlab::Diff::Line objects which text is hightlighted
@@ -236,8 +241,33 @@ module Gitlab
         simple_viewer.is_a?(DiffViewer::Text) && (ignore_errors || simple_viewer.render_error.nil?)
       end
 
+      # This adds the bottom match line to the array if needed. It contains
+      # the data to load more context lines.
+      def diff_lines_for_serializer
+        lines = highlighted_diff_lines
+
+        return if lines.empty?
+
+        last_line = lines.last
+
+        if last_line.new_pos < total_blob_lines(blob) && !deleted_file?
+          match_line = Gitlab::Diff::Line.new("", 'match', nil, last_line.old_pos, last_line.new_pos)
+          lines.push(match_line)
+        end
+
+        lines
+      end
+
       private
 
+      def total_blob_lines(blob)
+        @total_lines ||= begin
+          line_count = blob.lines.size
+          line_count -= 1 if line_count > 0 && blob.lines.last.blank?
+          line_count
+        end
+      end
+
       # We can't use Object#try because Blob doesn't inherit from Object, but
       # from BasicObject (via SimpleDelegator).
       def try_blobs(meth)