1 files changed, 64 insertions, 39 deletions

diff --git a/lib/gitlab/git/diff_collection.rb b/lib/gitlab/git/diff_collection.rb
index 4e45ec7c174..6a601561c2a 100644
--- a/lib/gitlab/git/diff_collection.rb
+++ b/lib/gitlab/git/diff_collection.rb
@@ -1,3 +1,5 @@
+# Gitaly note: JV: no RPC's here.
+
 module Gitlab
   module Git
     class DiffCollection
@@ -5,39 +7,58 @@ module Gitlab
 
       DEFAULT_LIMITS = { max_files: 100, max_lines: 5000 }.freeze
 
+      attr_reader :limits
+
+      delegate :max_files, :max_lines, :max_bytes, :safe_max_files, :safe_max_lines, :safe_max_bytes, to: :limits
+
+      def self.collection_limits(options = {})
+        limits = {}
+        limits[:max_files] = options.fetch(:max_files, DEFAULT_LIMITS[:max_files])
+        limits[:max_lines] = options.fetch(:max_lines, DEFAULT_LIMITS[:max_lines])
+        limits[:max_bytes] = limits[:max_files] * 5.kilobytes # Average 5 KB per file
+        limits[:safe_max_files] = [limits[:max_files], DEFAULT_LIMITS[:max_files]].min
+        limits[:safe_max_lines] = [limits[:max_lines], DEFAULT_LIMITS[:max_lines]].min
+        limits[:safe_max_bytes] = limits[:safe_max_files] * 5.kilobytes # Average 5 KB per file
+
+        OpenStruct.new(limits)
+      end
+
       def initialize(iterator, options = {})
         @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
-        @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)
-        @deltas_only = !!options.fetch(:deltas_only, false)
+        @limits = self.class.collection_limits(options)
+        @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)
-        elsif @deltas_only
-          each_delta(&block)
-        else
-          Gitlab::GitalyClient.migrate(:commit_raw_diffs) do
-            each_patch(&block)
+        @array.each(&block)
+
+        return if @overflow
+        return if @iterator.nil?
+
+        Gitlab::GitalyClient.migrate(:commit_raw_diffs) do |is_enabled|
+          if is_enabled && @iterator.is_a?(Gitlab::GitalyClient::DiffStitcher)
+            each_gitaly_patch(&block)
+          else
+            each_rugged_patch(&block)
           end
         end
+
+        @populated = true
+
+        # Allow iterator to be garbage-collected. It cannot be reused anyway.
+        @iterator = nil
       end
 
       def empty?
-        !@iterator.any?
+        any? # Make sure the iterator has been exercised
+        @empty
       end
 
       def overflow?
@@ -63,60 +84,63 @@ 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
 
       def over_safe_limits?(files)
-        files >= @safe_max_files || @line_count > @safe_max_lines || @byte_count >= @safe_max_bytes
+        files >= safe_max_files || @line_count > safe_max_lines || @byte_count >= safe_max_bytes
       end
 
-      def each_delta
-        @iterator.each_delta.with_index do |delta, i|
-          diff = Gitlab::Git::Diff.new(delta)
+      def each_gitaly_patch
+        i = @array.length
+
+        @iterator.each do |raw|
+          diff = Gitlab::Git::Diff.new(raw, expanded: !@enforce_limits || @expanded)
+
+          if raw.overflow_marker
+            @overflow = true
+            break
+          end
 
           yield @array[i] = diff
+          i += 1
         end
       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
+      def each_rugged_patch
+        i = @array.length
 
-          # 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.line_count > 0
+            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
@@ -124,6 +148,7 @@ module Gitlab
           end
 
           yield @array[i] = diff
+          i += 1
         end
       end
     end