1 files changed, 12 insertions, 5 deletions
diff --git a/app/models/merge_request.rb b/app/models/merge_request.rb
index ac26d29ad19..90061fe181e 100644
--- a/app/models/merge_request.rb
+++ b/app/models/merge_request.rb
@@ -454,17 +454,24 @@ class MergeRequest < ApplicationRecord
     true
   end
 
+  def preload_discussions_diff_highlight
+    preloadable_files = note_diff_files.for_commit_or_unresolved
+
+    discussions_diffs.load_highlight(preloadable_files.pluck(:id))
+  end
+
   def discussions_diffs
     strong_memoize(:discussions_diffs) do
-      note_diff_files = NoteDiffFile
-        .joins(:diff_note)
-        .merge(notes.or(commit_notes))
-        .includes(diff_note: :project)
-
       Gitlab::DiscussionsDiff::FileCollection.new(note_diff_files.to_a)
     end
   end
 
+  def note_diff_files
+    NoteDiffFile
+      .where(diff_note: discussion_notes)
+      .includes(diff_note: :project)
+  end
+
   def diff_size
     # Calling `merge_request_diff.diffs.real_size` will also perform
     # highlighting, which we don't need here.