Code cleanup

3 files changed, 17 insertions, 23 deletions

diff --git a/lib/gitlab/markdown/reference_filter.rb b/lib/gitlab/markdown/reference_filter.rb
index ede9f8865ff..adaca78ba27 100644
--- a/lib/gitlab/markdown/reference_filter.rb
+++ b/lib/gitlab/markdown/reference_filter.rb
@@ -14,10 +14,18 @@ module Gitlab
     class ReferenceFilter < HTML::Pipeline::Filter
       LazyReference = Struct.new(:klass, :ids) do
         def self.load(refs)
-          refs.group_by(&:klass).flat_map do |klass, refs|
+          lazy_references, values = refs.partition { |ref| ref.is_a?(self) }
+          
+          lazy_values = lazy_references.group_by(&:klass).flat_map do |klass, refs|
             ids = refs.flat_map(&:ids)
             klass.where(id: ids)
           end
+
+          values + lazy_values
+        end
+
+        def load
+          self.klass.where(id: self.ids)
         end
       end
 
diff --git a/lib/gitlab/markdown/reference_gatherer_filter.rb b/lib/gitlab/markdown/reference_gatherer_filter.rb
index 18df5db94d6..31fb71a98a3 100644
--- a/lib/gitlab/markdown/reference_gatherer_filter.rb
+++ b/lib/gitlab/markdown/reference_gatherer_filter.rb
@@ -12,8 +12,7 @@ module Gitlab
       def initialize(*)
         super
 
-        result[:lazy_references]  ||= Hash.new { |hash, type| hash[type] = [] }
-        result[:references]       ||= Hash.new { |hash, type| hash[type] = [] }
+        result[:references] ||= Hash.new { |hash, type| hash[type] = [] }
       end
 
       def call
@@ -41,23 +40,16 @@ module Gitlab
 
         references.each do |type, values|
           Array.wrap(values).each do |value|
-            refs = 
-              if value.is_a?(ReferenceFilter::LazyReference)
-                result[:lazy_references]
-              else
-                result[:references]
-              end
-
-            refs[type] << value
+            result[:references][type] << value
           end
         end
       end
 
       # Will load all references of one type using one query.
       def load_lazy_references
-        result[:lazy_references].each do |type, refs|
-          values = ReferenceFilter::LazyReference.load(refs)
-          result[:references][type].concat(values)
+        refs = result[:references]
+        refs.each do |type, values|
+          refs[type] = ReferenceFilter::LazyReference.load(values)
         end
       end
 
diff --git a/lib/gitlab/reference_extractor.rb b/lib/gitlab/reference_extractor.rb
index 8100f2675a7..d6b739d7b9a 100644
--- a/lib/gitlab/reference_extractor.rb
+++ b/lib/gitlab/reference_extractor.rb
@@ -54,19 +54,13 @@ module Gitlab
 
       pipeline = HTML::Pipeline.new([filter, Gitlab::Markdown::ReferenceGathererFilter], context)
 
-      values = []
-      lazy_references = []
-
-      @texts.each do |text|
+      values = @texts.flat_map do |text|
         result = pipeline.call(text)
 
-        values.concat(result[:references][filter_type])
-        lazy_references.concat(result[:lazy_references][filter_type])
+        result[:references][filter_type]
       end
 
-      lazy_values = Gitlab::Markdown::ReferenceFilter::LazyReference.load(lazy_references)
-      values.concat(lazy_values)
-      values
+      Gitlab::Markdown::ReferenceFilter::LazyReference.load(values)
     end
   end
 end