Handle diffs unfolding for batch loading

The diffs batch endpoint wasn't handling
the diff unfolding when leaving comments in
hidden parts of the diff.

This commit fixes that for both diff notes
and draft notes.

app/controllers/projects/merge_requests/diffs_controller.rb

@@ -25,6 +25,9 @@ class Projects::MergeRequests::DiffsController < Projects::MergeRequests::Applic
     return render_404 unless diffable
 
     diffs = diffable.diffs_in_batch(params[:page], params[:per_page], diff_options: diff_options)
+    positions = @merge_request.note_positions_for_paths(diffs.diff_file_paths, current_user)
+
+    diffs.unfold_diff_files(positions.unfoldable)
 
     options = {
       merge_request: @merge_request,

app/models/merge_request.rb

@@ -454,6 +454,15 @@ class MergeRequest < ApplicationRecord
     merge_request_diffs.where.not(id: merge_request_diff.id)
   end
 
+  # Overwritten in EE
+  def note_positions_for_paths(paths, _user = nil)
+    positions = notes.new_diff_notes.joins(:note_diff_file)
+      .where('note_diff_files.old_path IN (?) OR note_diff_files.new_path IN (?)', paths, paths)
+      .positions
+
+    Gitlab::Diff::PositionCollection.new(positions, diff_head_sha)
+  end
+
   def preloads_discussion_diff_highlighting?
     true
   end

app/models/note.rb

@@ -193,6 +193,12 @@ class Note < ApplicationRecord
       groups
     end
 
+    def positions
+      where.not(position: nil)
+        .select(:id, :type, :position) # ActiveRecord needs id and type for typecasting.
+        .map(&:position)
+    end
+
     def count_for_collection(ids, type)
       user.select('noteable_id', 'COUNT(*) as count')
         .group(:noteable_id)

ee/app/models/draft_note.rb

@@ -28,6 +28,12 @@ class DraftNote < ApplicationRecord
 
   scope :authored_by, ->(u) { where(author_id: u.id) }
 
+  def self.positions
+    where.not(position: nil)
+      .select(:position)
+      .map(&:position)
+  end
+
   def project
     merge_request.target_project
   end

ee/app/models/ee/merge_request.rb

@@ -106,6 +106,18 @@ module EE
       hidden.count
     end
 
+    override :note_positions_for_paths
+    def note_positions_for_paths(file_paths, user = nil)
+      return super unless user
+
+      positions = draft_notes
+        .authored_by(user)
+        .positions
+        .select { |pos| file_paths.include?(pos.file_path) }
+
+      super.concat(positions)
+    end
+
     def participant_approvers
       strong_memoize(:participant_approvers) do
         next [] unless approval_needed?

ee/spec/models/merge_request_spec.rb

@@ -50,6 +50,48 @@ describe MergeRequest do
     end
   end
 
+  describe '#note_positions_for_paths' do
+    let(:user) { create(:user) }
+    let(:merge_request) { create(:merge_request, :with_diffs) }
+    let(:project) { merge_request.project }
+    let!(:diff_note) do
+      create(:diff_note_on_merge_request, project: project, noteable: merge_request)
+    end
+    let!(:draft_note) do
+      create(:draft_note_on_text_diff, author: user, merge_request: merge_request)
+    end
+
+    let(:file_paths) { merge_request.diffs.diff_files.map(&:file_path) }
+
+    subject do
+      merge_request.note_positions_for_paths(file_paths)
+    end
+
+    it 'returns a Gitlab::Diff::PositionCollection' do
+      expect(subject).to be_a(Gitlab::Diff::PositionCollection)
+    end
+
+    context 'when user is given' do
+      subject do
+        merge_request.note_positions_for_paths(file_paths, user)
+      end
+
+      it 'returns notes and draft notes positions' do
+        expect(subject).to match_array([draft_note.position, diff_note.position])
+      end
+    end
+
+    context 'when user is not given' do
+      subject do
+        merge_request.note_positions_for_paths(file_paths)
+      end
+
+      it 'returns notes positions' do
+        expect(subject).to match_array([diff_note.position])
+      end
+    end
+  end
+
   describe '#participant_approvers' do
     let(:approvers) { create_list(:user, 2) }
     let(:code_owners) { create_list(:user, 2) }

lib/gitlab/diff/file_collection/merge_request_diff_batch.rb

@@ -29,6 +29,10 @@ module Gitlab
           }
         end
 
+        def diff_file_paths
+          diff_files.map(&:file_path)
+        end
+
         override :diffs
         def diffs
           strong_memoize(:diffs) do

lib/gitlab/diff/lines_unfolder.rb

@@ -54,7 +54,7 @@ module Gitlab
       def unfold_required?
         strong_memoize(:unfold_required) do
           next false unless @diff_file.text?
-          next false unless @position.on_text? && @position.unchanged?
+          next false unless @position.unfoldable?
           next false if @diff_file.new_file? || @diff_file.deleted_file?
           next false unless @position.old_line
           # Invalid position (MR import scenario)

lib/gitlab/diff/position.rb

@@ -79,6 +79,10 @@ module Gitlab
         formatter.line_age
       end
 
+      def unfoldable?
+        on_text? && unchanged?
+      end
+
       def unchanged?
         type.nil?
       end

lib/gitlab/diff/position_collection.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Diff
+    class PositionCollection
+      include Enumerable
+
+      # collection - An array of Gitlab::Diff::Position
+      def initialize(collection, diff_head_sha)
+        @collection = collection
+        @diff_head_sha = diff_head_sha
+      end
+
+      def each(&block)
+        @collection.each(&block)
+      end
+
+      def concat(positions)
+        tap { @collection.concat(positions) }
+      end
+
+      # Doing a lightweight filter in-memory given we're not prepared for querying
+      # positions (https://gitlab.com/gitlab-org/gitlab/issues/33271).
+      def unfoldable
+        select do |position|
+          position.unfoldable? && position.head_sha == @diff_head_sha
+        end
+      end
+    end
+  end
+end

spec/controllers/projects/merge_requests/diffs_controller_spec.rb

@@ -258,5 +258,26 @@ describe Projects::MergeRequests::DiffsController do
 
     it_behaves_like 'forked project with submodules'
     it_behaves_like 'persisted preferred diff view cookie'
+
+    context 'diff unfolding' do
+      let!(:unfoldable_diff_note) do
+        create(:diff_note_on_merge_request, :folded_position, project: project, noteable: merge_request)
+      end
+
+      let!(:diff_note) do
+        create(:diff_note_on_merge_request, project: project, noteable: merge_request)
+      end
+
+      it 'unfolds correct diff file positions' do
+        expect_next_instance_of(Gitlab::Diff::FileCollection::MergeRequestDiffBatch) do |instance|
+          expect(instance)
+            .to receive(:unfold_diff_files)
+            .with([unfoldable_diff_note.position])
+            .and_call_original
+        end
+
+        go
+      end
+    end
   end
 end

spec/factories/notes.rb

@@ -62,6 +62,18 @@ FactoryBot.define do
         )
       end
 
+      trait :folded_position do
+        position do
+          Gitlab::Diff::Position.new(
+            old_path: "files/ruby/popen.rb",
+            new_path: "files/ruby/popen.rb",
+            old_line: 1,
+            new_line: 1,
+            diff_refs: diff_refs
+          )
+        end
+      end
+
       trait :resolved do
         resolved_at { Time.now }
         resolved_by { create(:user) }

spec/lib/gitlab/diff/position_collection_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::Diff::PositionCollection do
+  let(:merge_request) { build(:merge_request) }
+
+  def build_text_position(attrs = {})
+    attributes = {
+      old_path: "files/ruby/popen.rb",
+      new_path: "files/ruby/popen.rb",
+      old_line: nil,
+      new_line: 14,
+      diff_refs: merge_request.diff_refs
+    }.merge(attrs)
+
+    Gitlab::Diff::Position.new(attributes)
+  end
+
+  def build_image_position(attrs = {})
+    attributes = {
+      old_path: "files/images/any_image.png",
+      new_path: "files/images/any_image.png",
+      width: 10,
+      height: 10,
+      x: 1,
+      y: 1,
+      diff_refs: merge_request.diff_refs,
+      position_type: "image"
+    }.merge(attrs)
+
+    Gitlab::Diff::Position.new(attributes)
+  end
+
+  let(:text_position) { build_text_position }
+  let(:folded_text_position) { build_text_position(old_line: 1, new_line: 1) }
+  let(:image_position) { build_image_position }
+  let(:head_sha) { merge_request.diff_head_sha }
+
+  let(:collection) do
+    described_class.new([text_position, folded_text_position, image_position], head_sha)
+  end
+
+  describe '#to_a' do
+    it 'returns all positions' do
+      expect(collection.to_a).to eq([text_position, folded_text_position, image_position])
+    end
+  end
+
+  describe '#unfoldable' do
+    it 'returns unfoldable diff positions' do
+      expect(collection.unfoldable).to eq([folded_text_position])
+    end
+
+    context 'when given head_sha does not match with positions head_sha' do
+      let(:head_sha) { 'unknown' }
+
+      it 'returns no position' do
+        expect(collection.unfoldable).to be_empty
+      end
+    end
+  end
+
+  describe '#concat' do
+    let(:new_text_position) { build_text_position(old_line: 1, new_line: 1) }
+
+    it 'returns a Gitlab::Diff::Position' do
+      expect(collection.concat([new_text_position])).to be_a(described_class)
+    end
+
+    it 'concatenates the new position to the collection' do
+      collection.concat([new_text_position])
+
+      expect(collection.to_a).to eq([text_position, folded_text_position, image_position, new_text_position])
+    end
+  end
+end

spec/models/merge_request_spec.rb

@@ -650,6 +650,46 @@ describe MergeRequest do
     end
   end
 
+  describe '#note_positions_for_paths' do
+    let(:merge_request) { create(:merge_request, :with_diffs) }
+    let(:project) { merge_request.project }
+    let!(:diff_note) do
+      create(:diff_note_on_merge_request, project: project, noteable: merge_request)
+    end
+
+    let(:file_paths) { merge_request.diffs.diff_files.map(&:file_path) }
+
+    subject do
+      merge_request.note_positions_for_paths(file_paths)
+    end
+
+    it 'returns a Gitlab::Diff::PositionCollection' do
+      expect(subject).to be_a(Gitlab::Diff::PositionCollection)
+    end
+
+    context 'within all diff files' do
+      it 'returns correct positions' do
+        expect(subject).to match_array([diff_note.position])
+      end
+    end
+
+    context 'within specific diff file' do
+      let(:file_paths) { [diff_note.position.file_path] }
+
+      it 'returns correct positions' do
+        expect(subject).to match_array([diff_note.position])
+      end
+    end
+
+    context 'within no diff files' do
+      let(:file_paths) { [] }
+
+      it 'returns no positions' do
+        expect(subject.to_a).to be_empty
+      end
+    end
+  end
+
   describe '#discussions_diffs' do
     let(:merge_request) { create(:merge_request) }