1 files changed, 33 insertions, 28 deletions

diff --git a/lib/gitlab/git/diff_collection.rb b/lib/gitlab/git/diff_collection.rb
index bcbad8ec829..334e06a6eca 100644
--- a/lib/gitlab/git/diff_collection.rb
+++ b/lib/gitlab/git/diff_collection.rb
@@ -9,32 +9,29 @@ module Gitlab
         @iterator = iterator
         @max_files = options.fetch(:max_files, DEFAULT_LIMITS[:max_files])
         @max_lines = options.fetch(:max_lines, DEFAULT_LIMITS[:max_lines])
-        @max_bytes = @max_files * 5120 # Average 5 KB per file
+        @max_bytes = @max_files * 5.kilobytes # Average 5 KB per file
         @safe_max_files = [@max_files, DEFAULT_LIMITS[:max_files]].min
         @safe_max_lines = [@max_lines, DEFAULT_LIMITS[:max_lines]].min
-        @safe_max_bytes = @safe_max_files * 5120 # Average 5 KB per file
-        @all_diffs = !!options.fetch(:all_diffs, false)
-        @no_collapse = !!options.fetch(:no_collapse, true)
+        @safe_max_bytes = @safe_max_files * 5.kilobytes # Average 5 KB per file
+        @enforce_limits = !!options.fetch(:limits, true)
+        @expanded = !!options.fetch(:expanded, true)
 
         @line_count = 0
         @byte_count = 0
         @overflow = false
+        @empty = true
         @array = Array.new
       end
 
       def each(&block)
-        if @populated
-          # @iterator.each is slower than just iterating the array in place
-          @array.each(&block)
-        else
-          Gitlab::GitalyClient.migrate(:commit_raw_diffs) do
-            each_patch(&block)
-          end
+        Gitlab::GitalyClient.migrate(:commit_raw_diffs) do
+          each_patch(&block)
         end
       end
 
       def empty?
-        !@iterator.any?
+        any? # Make sure the iterator has been exercised
+        @empty
       end
 
       def overflow?
@@ -60,17 +57,17 @@ module Gitlab
         collection = each_with_index do |element, i|
           @array[i] = yield(element)
         end
-        @populated = true
         collection
       end
 
+      alias_method :to_ary, :to_a
+
       private
 
       def populate!
         return if @populated
 
         each { nil } # force a loop through all diffs
-        @populated = true
         nil
       end
 
@@ -79,33 +76,35 @@ module Gitlab
       end
 
       def each_patch
-        @iterator.each_with_index do |raw, i|
-          # First yield cached Diff instances from @array
-          if @array[i]
-            yield @array[i]
-            next
-          end
+        i = 0
+        @array.each do |diff|
+          yield diff
+          i += 1
+        end
+
+        return if @overflow
+        return if @iterator.nil?
 
-          # We have exhausted @array, time to create new Diff instances or stop.
-          break if @overflow
+        @iterator.each do |raw|
+          @empty = false
 
-          if !@all_diffs && i >= @max_files
+          if @enforce_limits && i >= @max_files
             @overflow = true
             break
           end
 
-          collapse = !@all_diffs && !@no_collapse
+          expanded = !@enforce_limits || @expanded
 
-          diff = Gitlab::Git::Diff.new(raw, collapse: collapse)
+          diff = Gitlab::Git::Diff.new(raw, expanded: expanded)
 
-          if collapse && over_safe_limits?(i)
-            diff.prune_collapsed_diff!
+          if !expanded && over_safe_limits?(i)
+            diff.collapse!
           end
 
           @line_count += diff.line_count
           @byte_count += diff.diff.bytesize
 
-          if !@all_diffs && (@line_count >= @max_lines || @byte_count >= @max_bytes)
+          if @enforce_limits && (@line_count >= @max_lines || @byte_count >= @max_bytes)
             # This last Diff instance pushes us over the lines limit. We stop and
             # discard it.
             @overflow = true
@@ -113,7 +112,13 @@ module Gitlab
           end
 
           yield @array[i] = diff
+          i += 1
         end
+
+        @populated = true
+
+        # Allow iterator to be garbage-collected. It cannot be reused anyway.
+        @iterator = nil
       end
     end
   end