Move to Gitlab::Diff::FileCollection

Instead calling diff_collection.count use diff_collection.size which is cache on the diff_collection

6 files changed, 156 insertions, 1 deletions

diff --git a/lib/gitlab/diff/file.rb b/lib/gitlab/diff/file.rb
index 77b3798d78f..e47df508ca2 100644
--- a/lib/gitlab/diff/file.rb
+++ b/lib/gitlab/diff/file.rb
@@ -63,7 +63,7 @@ module Gitlab
         diff_refs.try(:head_sha)
       end
 
-      attr_writer :diff_lines, :highlighted_diff_lines
+      attr_writer :highlighted_diff_lines
 
       # Array of Gitlab::Diff::Line objects
       def diff_lines
diff --git a/lib/gitlab/diff/file_collection.rb b/lib/gitlab/diff/file_collection.rb
new file mode 100644
index 00000000000..ce6717c7205
--- /dev/null
+++ b/lib/gitlab/diff/file_collection.rb
@@ -0,0 +1,9 @@
+module Gitlab
+  module Diff
+    module FileCollection
+      def self.default_options
+        ::Commit.max_diff_options.merge(ignore_whitespace_change: false, no_collapse: false)
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/diff/file_collection/base.rb b/lib/gitlab/diff/file_collection/base.rb
new file mode 100644
index 00000000000..20562773c14
--- /dev/null
+++ b/lib/gitlab/diff/file_collection/base.rb
@@ -0,0 +1,30 @@
+module Gitlab
+  module Diff
+    module FileCollection
+
+      class Base
+        attr_reader :project, :diff_options, :diff_view, :diff_refs
+
+        delegate :count, :size, :real_size, to: :diff_files
+
+        def initialize(diffs, project:, diff_options:, diff_refs: nil)
+          @diffs        = diffs
+          @project      = project
+          @diff_options = diff_options
+          @diff_refs    = diff_refs
+        end
+
+        def diff_files
+          @diffs.decorate! { |diff| decorate_diff!(diff) }
+        end
+
+        private
+
+        def decorate_diff!(diff)
+          return diff if diff.is_a?(Gitlab::Diff::File)
+          Gitlab::Diff::File.new(diff, diff_refs: @diff_refs, repository: @project.repository)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/diff/file_collection/commit.rb b/lib/gitlab/diff/file_collection/commit.rb
new file mode 100644
index 00000000000..2a46109ad99
--- /dev/null
+++ b/lib/gitlab/diff/file_collection/commit.rb
@@ -0,0 +1,14 @@
+module Gitlab
+  module Diff
+    module FileCollection
+      class Commit < Base
+        def initialize(commit, diff_options:)
+          super(commit.diffs(diff_options),
+            project: commit.project,
+            diff_options: diff_options,
+            diff_refs: commit.diff_refs)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/diff/file_collection/compare.rb b/lib/gitlab/diff/file_collection/compare.rb
new file mode 100644
index 00000000000..1bcda145f15
--- /dev/null
+++ b/lib/gitlab/diff/file_collection/compare.rb
@@ -0,0 +1,14 @@
+module Gitlab
+  module Diff
+    module FileCollection
+      class Compare < Base
+        def initialize(compare, project:, diff_options:, diff_refs: nil)
+          super(compare.diffs(diff_options),
+            project:      project,
+            diff_options: diff_options,
+            diff_refs:    diff_refs)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/diff/file_collection/merge_request.rb b/lib/gitlab/diff/file_collection/merge_request.rb
new file mode 100644
index 00000000000..7c40622d594
--- /dev/null
+++ b/lib/gitlab/diff/file_collection/merge_request.rb
@@ -0,0 +1,88 @@
+module Gitlab
+  module Diff
+    module FileCollection
+      class MergeRequest < Base
+        def initialize(merge_request, diff_options:)
+          @merge_request = merge_request
+
+          super(merge_request.diffs(diff_options),
+            project: merge_request.project,
+            diff_options: diff_options,
+            diff_refs: merge_request.diff_refs)
+        end
+
+        def diff_files
+          super.tap { |_| store_highlight_cache }
+        end
+
+        private
+
+        # Extracted method to highlight in the same iteration to the diff_collection. Iteration in the DiffCollections
+        # seems particularly slow on big diffs (event when already populated).
+        def decorate_diff!(diff)
+          highlight! super
+        end
+
+        def highlight!(diff_file)
+          if cacheable?
+            cache_highlight!(diff_file)
+          else
+            highlight_diff_file!(diff_file)
+          end
+        end
+
+        def highlight_diff_file!(diff_file)
+          diff_file.highlighted_diff_lines = Gitlab::Diff::Highlight.new(diff_file, repository: diff_file.repository).highlight
+          diff_file
+        end
+
+        def highlight_diff_file_from_cache!(diff_file, cache_diff_lines)
+          diff_file.highlighted_diff_lines = cache_diff_lines.map do |line|
+            Gitlab::Diff::Line.init_from_hash(line)
+          end
+        end
+
+        #
+        # If we find the highlighted diff files lines on the cache we replace existing diff_files lines (no highlighted)
+        # for the highlighted ones, so we just skip their execution.
+        # If the highlighted diff files lines are not cached we calculate and cache them.
+        #
+        # The content of the cache is a Hash where the key correspond to the file_path and the values are Arrays of
+        # hashes that represent serialized diff lines.
+        #
+        def cache_highlight!(diff_file)
+          file_path = diff_file.file_path
+
+          if highlight_cache[file_path]
+            highlight_diff_file_from_cache!(diff_file, highlight_cache[file_path])
+          else
+            highlight_diff_file!(diff_file)
+            highlight_cache[file_path] = diff_file.highlighted_diff_lines.map(&:to_hash)
+          end
+
+          diff_file
+        end
+
+        def highlight_cache
+          return @highlight_cache if defined?(@highlight_cache)
+
+          @highlight_cache = Rails.cache.read(cache_key) || {}
+          @highlight_cache_was_empty = highlight_cache.empty?
+          @highlight_cache
+        end
+
+        def store_highlight_cache
+          Rails.cache.write(cache_key, highlight_cache) if @highlight_cache_was_empty
+        end
+
+        def cacheable?
+          @merge_request.merge_request_diff.present?
+        end
+
+        def cache_key
+          [@merge_request.merge_request_diff, 'highlighted-diff-files', diff_options]
+        end
+      end
+    end
+  end
+end