Use internal reference extractor in banzai unfold pipeline

2 files changed, 27 insertions, 21 deletions

diff --git a/lib/banzai/filter/reference_unfold_filter.rb b/lib/banzai/filter/reference_unfold_filter.rb
index a6145261651..d8ed44a0e4e 100644
--- a/lib/banzai/filter/reference_unfold_filter.rb
+++ b/lib/banzai/filter/reference_unfold_filter.rb
@@ -1,45 +1,39 @@
+require 'html/pipeline/filter'
+
 module Banzai
   module Filter
     ##
     # Filter than unfolds local references.
     #
-    # Replaces all local references with project cross reference version
-    # in all objects passed to this filter in context.
-    #
-    # Requires objects array with each element implementing `Referable`.
     #
-    class ReferenceUnfoldFilter < ReferenceFilter
+    class ReferenceUnfoldFilter < HTML::Pipeline::Filter
       def initialize(*)
         super
 
-        @objects = context[:objects]
-        @project = context[:project]
-
-        unless @objects.all? { |object| object.respond_to?(:to_reference) }
-          raise StandardError, "No `to_reference` method implemented in one of the objects !"
+        unless result[:references].is_a?(Hash)
+          raise StandardError, 'References not processed!'
         end
 
-        unless @project.kind_of?(Project)
-          raise StandardError, 'No valid project passed in context!'
-        end
+        @text = context[:text].dup
+        @new_project = context[:new_project]
+        @referables = result[:references].values.flatten
       end
 
       def call
-        @objects.each do |object|
-          pattern = /#{Regexp.escape(object.to_reference)}/
-          replace_text_nodes_matching(pattern) do |content|
-            content.gsub(pattern, object.to_reference(@project))
-          end
+        @referables.each do |referable|
+          pattern = /#{Regexp.escape(referable.to_reference)}/
+          @text.gsub!(pattern, referable.to_reference(@new_project))
         end
 
-        doc
+        @text
       end
 
       private
 
       def validate
         needs :project
-        needs :objects
+        needs :new_project
+        needs :text
       end
     end
   end
diff --git a/lib/banzai/pipeline/reference_unfold_pipeline.rb b/lib/banzai/pipeline/reference_unfold_pipeline.rb
index 5b555d7c2d7..597bf7befd1 100644
--- a/lib/banzai/pipeline/reference_unfold_pipeline.rb
+++ b/lib/banzai/pipeline/reference_unfold_pipeline.rb
@@ -2,7 +2,19 @@ module Banzai
   module Pipeline
     class ReferenceUnfoldPipeline < BasePipeline
       def self.filters
-        [Filter::ReferenceUnfoldFilter]
+        FullPipeline.filters +
+        [Filter::ReferenceGathererFilter,
+         Filter::ReferenceUnfoldFilter]
+      end
+
+      def self.call(text, context = {})
+        context = context.merge(text: text)
+        super
+      end
+
+      class << self
+        alias_method :to_document, :call
+        alias_method :to_html, :call
       end
     end
   end