From 92fd6c06701b3c6a751524bb419ddb911fe3317e Mon Sep 17 00:00:00 2001
From: Yuki Nishijima <yuki24@hey.com>
Date: Tue, 16 Feb 2021 23:39:03 -0500
Subject: [PATCH] Speed up collection rendering and support collection caching

---
 lib/jbuilder/collection_renderer.rb | 109 ++++++++++++++++++++++++++++
 lib/jbuilder/jbuilder_template.rb   |  32 +++++---
 2 files changed, 129 insertions(+), 12 deletions(-)
 create mode 100644 lib/jbuilder/collection_renderer.rb

diff --git a/lib/jbuilder/collection_renderer.rb b/lib/jbuilder/collection_renderer.rb
new file mode 100644
index 00000000..ee980d6a
--- /dev/null
+++ b/lib/jbuilder/collection_renderer.rb
@@ -0,0 +1,109 @@
+require 'delegate'
+
+begin
+  require 'action_view/renderer/collection_renderer'
+rescue LoadError
+  require 'action_view/renderer/partial_renderer'
+end
+
+class Jbuilder
+  module CollectionRenderable
+    private
+
+    def build_rendered_template(content, template, layout = nil)
+      super(content || json.attributes!, template)
+    end
+
+    def build_rendered_collection(templates, _spacer)
+      json.merge!(templates.map(&:body))
+    end
+
+    def json
+      @options[:locals].fetch(:json)
+    end
+
+    class ScopedIterator < ::SimpleDelegator # :nodoc:
+      include Enumerable
+
+      def initialize(obj, scope)
+        super(obj)
+        @scope = scope
+      end
+
+      # Rails 6.0 support:
+      def each
+        return enum_for(:each) unless block_given?
+
+        __getobj__.each do |object|
+          @scope.call { yield(object) }
+        end
+      end
+
+      # Rails 6.1 support:
+      def each_with_info
+        return enum_for(:each_with_info) unless block_given?
+
+        __getobj__.each_with_info do |object, (path, as, counter, iteration)|
+          @scope.call { yield(object, [path, as, counter, iteration]) }
+        end
+      end
+    end
+
+    private_constant :ScopedIterator
+  end
+
+  collection_renderer_available = begin
+                                     ::ActionView::CollectionRenderer
+
+                                     true
+                                   rescue ::NameError
+                                     false
+                                   end
+
+  if collection_renderer_available
+    # Rails 6.1 support:
+    class CollectionRenderer < ::ActionView::CollectionRenderer
+      include CollectionRenderable
+
+      def initialize(lookup_context, options, &scope)
+        super(lookup_context, options)
+        @scope = scope
+      end
+
+      private
+
+        def collection_with_template(view, template, layout, collection)
+          super(view, template, layout, ScopedIterator.new(collection, @scope))
+        end
+    end
+  else
+    # Rails 6.0 support:
+    class CollectionRenderer < ::ActionView::PartialRenderer
+      include CollectionRenderable
+
+      def initialize(lookup_context, options, &scope)
+        super(lookup_context)
+        @options = options
+        @scope = scope
+      end
+
+      def render_collection_with_partial(collection, partial, context, block)
+        render(context, @options.merge(collection: collection, partial: partial), block)
+      end
+
+      private
+
+        def collection_without_template(view)
+          @collection = ScopedIterator.new(@collection, @scope)
+
+          super(view)
+        end
+
+        def collection_with_template(view, template)
+          @collection = ScopedIterator.new(@collection, @scope)
+
+          super(view, template)
+        end
+    end
+  end
+end
diff --git a/lib/jbuilder/jbuilder_template.rb b/lib/jbuilder/jbuilder_template.rb
index 26f94ae9..27c9d3a5 100644
--- a/lib/jbuilder/jbuilder_template.rb
+++ b/lib/jbuilder/jbuilder_template.rb
@@ -1,4 +1,5 @@
 require 'jbuilder/jbuilder'
+require 'jbuilder/collection_renderer'
 require 'action_dispatch/http/mime_type'
 require 'active_support/cache'
 
@@ -104,20 +105,27 @@ def set!(name, object = BLANK, *args)
   private
 
   def _render_partial_with_options(options)
-    options.reverse_merge! locals: options.except(:partial, :as, :collection)
+    options.reverse_merge! locals: options.except(:partial, :as, :collection, :cached)
     options.reverse_merge! ::JbuilderTemplate.template_lookup_options
-    as = options[:as]
-
-    if as && options.key?(:collection)
-      as = as.to_sym
-      collection = options.delete(:collection)
-      locals = options.delete(:locals)
-      array! collection do |member|
-        member_locals = locals.clone
-        member_locals.merge! collection: collection
-        member_locals.merge! as => member
-        _render_partial options.merge(locals: member_locals)
+
+    if options.key?(:collection) && options[:collection].nil?
+      array!
+    elsif options[:as] && options.key?(:collection)
+      collection = options.delete(:collection) || []
+      partial = options.delete(:partial)
+      options[:locals].merge!(json: self)
+
+      if options.has_key?(:layout)
+        raise NotImplementedError, "The `:layout option' is not supported in collection rendering."
       end
+
+      if options.has_key?(:spacer_template)
+        raise NotImplementedError, "The `:spacer_template' option is not supported in collection rendering."
+      end
+
+      CollectionRenderer
+        .new(@context.lookup_context, options) { |&block| _scope(&block) }
+        .render_collection_with_partial(collection, partial, @context, nil)
     else
       _render_partial options
     end