1 files changed, 62 insertions, 23 deletions

diff --git a/app/presenters/blobs/unfold_presenter.rb b/app/presenters/blobs/unfold_presenter.rb
index 7b13db3bb74..a256dd05a4d 100644
--- a/app/presenters/blobs/unfold_presenter.rb
+++ b/app/presenters/blobs/unfold_presenter.rb
@@ -1,34 +1,42 @@
 # frozen_string_literal: true
 
-require 'gt_one_coercion'
-
 module Blobs
   class UnfoldPresenter < BlobPresenter
-    include Virtus.model
+    include ActiveModel::Attributes
+    include ActiveModel::AttributeAssignment
     include Gitlab::Utils::StrongMemoize
 
-    attribute :full, Boolean, default: false
-    attribute :since, GtOneCoercion
-    attribute :to, GtOneCoercion
-    attribute :bottom, Boolean
-    attribute :unfold, Boolean, default: true
-    attribute :offset, Integer
-    attribute :indent, Integer, default: 0
+    attribute :full, :boolean, default: false
+    attribute :since, :integer, default: 1
+    attribute :to, :integer, default: 1
+    attribute :bottom, :boolean, default: false
+    attribute :unfold, :boolean, default: true
+    attribute :offset, :integer, default: 0
+    attribute :indent, :integer, default: 0
+
+    alias_method :full?, :full
+    alias_method :bottom?, :bottom
+    alias_method :unfold?, :unfold
 
     def initialize(blob, params)
-      @subject = blob
-      @all_lines = highlight.lines
-      super(params)
+      super(blob)
+      self.attributes = params
 
-      if full?
-        self.attributes = { since: 1, to: @all_lines.size, bottom: false, unfold: false, offset: 0, indent: 0 }
-      end
+      # Load all blob data first as we need to ensure they're all loaded first
+      # so we can accurately show the rest of the diff when unfolding.
+      load_all_blob_data
+
+      @all_lines = blob.data.lines
+
+      handle_full_or_end!
     end
 
-    # Converts a String array to Gitlab::Diff::Line array, with match line added
+    # Returns an array of Gitlab::Diff::Line with match line added
     def diff_lines
-      diff_lines = lines.map do |line|
-        Gitlab::Diff::Line.new(line, nil, nil, nil, nil, rich_text: line)
+      diff_lines = lines.map.with_index do |line, index|
+        full_line = limited_blob_lines[index].delete("\n")
+
+        Gitlab::Diff::Line.new(full_line, nil, nil, nil, nil, rich_text: line)
       end
 
       add_match_line(diff_lines)
@@ -38,9 +46,7 @@ module Blobs
 
     def lines
       strong_memoize(:lines) do
-        lines = @all_lines
-        lines = lines[since - 1..to - 1] unless full?
-        lines.map(&:html_safe)
+        limit(highlight.lines).map(&:html_safe)
       end
     end
 
@@ -54,10 +60,31 @@ module Blobs
 
     private
 
+    def handle_full_or_end!
+      return unless full? || to == -1
+
+      self.since = 1 if full?
+
+      self.attributes = {
+        to: all_lines_size,
+        bottom: false,
+        unfold: false,
+        offset: 0,
+        indent: 0
+      }
+    end
+
+    def all_lines_size
+      strong_memoize(:all_lines_size) do
+        @all_lines.size
+      end
+    end
+
     def add_match_line(diff_lines)
       return unless unfold?
+      return if bottom? && to >= all_lines_size
 
-      if bottom? && to < @all_lines.size
+      if bottom? && to < all_lines_size
         old_pos = to - offset
         new_pos = to
       elsif since != 1
@@ -71,5 +98,17 @@ module Blobs
 
       bottom? ? diff_lines.push(match_line) : diff_lines.unshift(match_line)
     end
+
+    def limited_blob_lines
+      strong_memoize(:limited_blob_lines) do
+        limit(@all_lines)
+      end
+    end
+
+    def limit(lines)
+      return lines if full?
+
+      lines[since - 1..to - 1]
+    end
   end
 end