Add custom polymorphic mapping

Allow the use of `direct` to specify custom mappings for polymorphic_url, e.g:

  resource :basket
  direct(class: "Basket") { [:basket] }

This will then generate the following:

  >> link_to "Basket", @basket
  => <a href="/basket">Basket</a>

More importantly it will generate the correct url when used with `form_for`.

Fixes #1769.

3 files changed, 141 insertions, 56 deletions

diff --git a/actionpack/lib/action_dispatch/routing/mapper.rb b/actionpack/lib/action_dispatch/routing/mapper.rb
index 3756ef15a2..6123a1f5f5 100644
--- a/actionpack/lib/action_dispatch/routing/mapper.rb
+++ b/actionpack/lib/action_dispatch/routing/mapper.rb
@@ -2021,8 +2021,8 @@ module ActionDispatch
       end
 
       module DirectUrls
-        # Define a custom url helper that will be added to the url helpers
-        # module. This allows you override and/or replace the default behavior
+        # Define custom routing behavior that will be added to the application's
+        # routes. This allows you override and/or replace the default behavior
         # of routing helpers, e.g:
         #
         #   direct :homepage do
@@ -2037,8 +2037,35 @@ module ActionDispatch
         #     { controller: 'pages', action: 'index', subdomain: 'www' }
         #   end
         #
-        # The return value must be a valid set of arguments for `url_for` which
-        # will actually build the url string. This can be one of the following:
+        # The above example show how to define a custom url helper but it's also
+        # possible to alter the behavior of `polymorphic_url` and consequently the
+        # behavior of `link_to` and `form_for` when passed a model instance, e.g:
+        #
+        #   direct class: "Basket" do
+        #     [:basket]
+        #   end
+        #
+        # NOTE: This custom behavior only applies to simple polymorphic urls where
+        # a single model instance is passed and not more complicated forms, e.g:
+        #
+        #   # config/routes.rb
+        #   resource :profile
+        #   namespace :admin do
+        #     resources :users
+        #   end
+        #
+        #   direct(class: "User") { [:profile] }
+        #
+        #   # app/views/application/_menu.html.erb
+        #   link_to 'Profile', @current_user
+        #   link_to 'Profile', [:admin, @current_user]
+        #
+        # The first `link_to` will generate '/profile' but the second will generate
+        # the standard polymorphic url of '/admin/users/1'.
+        #
+        # The return value from the block passed to `direct` must be a valid set of
+        # arguments for `url_for` which will actually build the url string. This can
+        # be one of the following:
         #
         #   * A string, which is treated as a generated url
         #   * A hash, e.g. { controller: 'pages', action: 'index' }
@@ -2046,6 +2073,9 @@ module ActionDispatch
         #   * An Active Model instance
         #   * An Active Model class
         #
+        # NOTE: Other url helpers can be called in the block but be careful not to invoke
+        # your custom url helper again otherwise it will result in a stack overflow error
+        #
         # You can also specify default options that will be passed through to
         # your url helper definition, e.g:
         #
@@ -2053,14 +2083,28 @@ module ActionDispatch
         #     [ :products, options.merge(params.permit(:page, :size)) ]
         #   end
         #
-        # NOTE: It is the url helper's responsibility to return the correct
-        # set of options to be passed to the `url_for` call.
-        def direct(name, options = nil, &block)
-          case name
+        # You can pass options to a polymorphic mapping do - the arity for the block
+        # needs to be two as the instance is passed as the first argument, e.g:
+        #
+        #   direct class: 'Basket', anchor: 'items' do |basket, options|
+        #     [:basket, options]
+        #   end
+        #
+        # This generates the url '/basket#items' because when the last item in an
+        # array passed to `polymorphic_url` is a hash then it's treated as options
+        # to the url helper that gets called.
+        #
+        # NOTE: The `direct` method doesn't observe the current scope in routes.rb
+        # and because of this it's recommended to define them outside of any blocks
+        # such as `namespace` or `scope`.
+        def direct(name_or_hash, options = nil, &block)
+          case name_or_hash
+          when Hash
+            @set.add_polymorphic_mapping(name_or_hash, &block)
           when String, Symbol
-            @set.add_url_helper(name, options, &block)
+            @set.add_url_helper(name_or_hash, options, &block)
           else
-            raise ArgumentError, "The direct method only accepts a string or symbol"
+            raise ArgumentError, "The direct method only accepts a hash, string or symbol"
           end
         end
       end
diff --git a/actionpack/lib/action_dispatch/routing/polymorphic_routes.rb b/actionpack/lib/action_dispatch/routing/polymorphic_routes.rb
index 432b9bf4c1..512e23c833 100644
--- a/actionpack/lib/action_dispatch/routing/polymorphic_routes.rb
+++ b/actionpack/lib/action_dispatch/routing/polymorphic_routes.rb
@@ -103,6 +103,10 @@ module ActionDispatch
           return polymorphic_url record, options
         end
 
+        if mapping = polymorphic_mapping(record_or_hash_or_array)
+          return mapping.call(self, [record_or_hash_or_array], options)
+        end
+
         opts   = options.dup
         action = opts.delete :action
         type   = opts.delete(:routing_type) || :url
@@ -123,6 +127,10 @@ module ActionDispatch
           return polymorphic_path record, options
         end
 
+        if mapping = polymorphic_mapping(record_or_hash_or_array)
+          return mapping.call(self, [record_or_hash_or_array], options, only_path: true)
+        end
+
         opts   = options.dup
         action = opts.delete :action
         type   = :path
@@ -156,6 +164,11 @@ module ActionDispatch
           polymorphic_path(record_or_hash, options.merge(action: action))
         end
 
+        def polymorphic_mapping(record)
+          return false unless record.respond_to?(:to_model)
+          _routes.polymorphic_mappings[record.to_model.model_name.name]
+        end
+
         class HelperMethodBuilder # :nodoc:
           CACHE = { "path" => {}, "url" => {} }
 
@@ -255,9 +268,13 @@ module ActionDispatch
             [named_route, args]
           end
 
-          def handle_model_call(target, model)
-            method, args = handle_model model
-            target.send(method, *args)
+          def handle_model_call(target, record)
+            if mapping = polymorphic_mapping(target, record)
+              mapping.call(target, [record], {}, only_path: suffix == "path")
+            else
+              method, args = handle_model(record)
+              target.send(method, *args)
+            end
           end
 
           def handle_list(list)
@@ -303,6 +320,10 @@ module ActionDispatch
 
           private
 
+            def polymorphic_mapping(target, record)
+              target._routes.polymorphic_mappings[record.to_model.model_name.name]
+            end
+
             def get_method_for_class(klass)
               name = @key_strategy.call klass.model_name
               get_method_for_string name
diff --git a/actionpack/lib/action_dispatch/routing/route_set.rb b/actionpack/lib/action_dispatch/routing/route_set.rb
index 923d063655..8bdf0d1a53 100644
--- a/actionpack/lib/action_dispatch/routing/route_set.rb
+++ b/actionpack/lib/action_dispatch/routing/route_set.rb
@@ -160,7 +160,7 @@ module ActionDispatch
 
         def add_url_helper(name, defaults, &block)
           @custom_helpers << name
-          helper = CustomUrlHelper.new(name, defaults, &block)
+          helper = DirectUrlHelper.new(name, defaults, &block)
 
           @path_helpers_module.module_eval do
             define_method(:"#{name}_path") do |*args|
@@ -177,47 +177,6 @@ module ActionDispatch
           end
         end
 
-        class CustomUrlHelper
-          attr_reader :name, :defaults, :block
-
-          def initialize(name, defaults, &block)
-            @name = name
-            @defaults = defaults
-            @block = block
-          end
-
-          def call(t, args, options, outer_options = {})
-            url_options = eval_block(t, args, options)
-
-            case url_options
-            when String
-              t.url_for(url_options)
-            when Hash
-              t.url_for(url_options.merge(outer_options))
-            when ActionController::Parameters
-              if url_options.permitted?
-                t.url_for(url_options.to_h.merge(outer_options))
-              else
-                raise ArgumentError, "Generating an URL from non sanitized request parameters is insecure!"
-              end
-            when Array
-              opts = url_options.extract_options!
-              t.url_for(url_options.push(opts.merge(outer_options)))
-            else
-              t.url_for([url_options, outer_options])
-            end
-          end
-
-          private
-            def eval_block(t, args, options)
-              t.instance_exec(*args, merge_defaults(options), &block)
-            end
-
-            def merge_defaults(options)
-              defaults ? defaults.merge(options) : options
-            end
-        end
-
         class UrlHelper
           def self.create(route, options, route_name, url_strategy)
             if optimize_helper?(route)
@@ -380,7 +339,7 @@ module ActionDispatch
       attr_accessor :formatter, :set, :named_routes, :default_scope, :router
       attr_accessor :disable_clear_and_finalize, :resources_path_names
       attr_accessor :default_url_options
-      attr_reader :env_key
+      attr_reader :env_key, :polymorphic_mappings
 
       alias :routes :set
 
@@ -422,6 +381,7 @@ module ActionDispatch
         @set    = Journey::Routes.new
         @router = Journey::Router.new @set
         @formatter = Journey::Formatter.new self
+        @polymorphic_mappings = {}
       end
 
       def eager_load!
@@ -483,6 +443,7 @@ module ActionDispatch
         named_routes.clear
         set.clear
         formatter.clear
+        @polymorphic_mappings.clear
         @prepend.each { |blk| eval_block(blk) }
       end
 
@@ -554,6 +515,14 @@ module ActionDispatch
               @_proxy.optimize_routes_generation?
             end
 
+            def polymorphic_url(record_or_hash_or_array, options = {})
+              @_proxy.polymorphic_url(record_or_hash_or_array, options)
+            end
+
+            def polymorphic_path(record_or_hash_or_array, options = {})
+              @_proxy.polymorphic_path(record_or_hash_or_array, options)
+            end
+
             def _routes; @_proxy._routes; end
             def url_options; {}; end
           end
@@ -629,10 +598,61 @@ module ActionDispatch
         route
       end
 
+      def add_polymorphic_mapping(options, &block)
+        defaults = options.dup
+        klass = defaults.delete(:class)
+        if klass.nil?
+          raise ArgumentError, "Missing :class key from polymorphic mapping options"
+        end
+
+        @polymorphic_mappings[klass] = DirectUrlHelper.new(klass, defaults, &block)
+      end
+
       def add_url_helper(name, options, &block)
         named_routes.add_url_helper(name, options, &block)
       end
 
+      class DirectUrlHelper
+        attr_reader :name, :defaults, :block
+
+        def initialize(name, defaults, &block)
+          @name = name
+          @defaults = defaults
+          @block = block
+        end
+
+        def call(t, args, options, outer_options = {})
+          url_options = eval_block(t, args, options)
+
+          case url_options
+          when String
+            t.url_for(url_options)
+          when Hash
+            t.url_for(url_options.merge(outer_options))
+          when ActionController::Parameters
+            if url_options.permitted?
+              t.url_for(url_options.to_h.merge(outer_options))
+            else
+              raise ArgumentError, "Generating an URL from non sanitized request parameters is insecure!"
+            end
+          when Array
+            opts = url_options.extract_options!
+            t.url_for(url_options.push(opts.merge(outer_options)))
+          else
+            t.url_for([url_options, outer_options])
+          end
+        end
+
+        private
+          def eval_block(t, args, options)
+            t.instance_exec(*args, merge_defaults(options), &block)
+          end
+
+          def merge_defaults(options)
+            defaults ? defaults.merge(options) : options
+          end
+      end
+
       class Generator
         PARAMETERIZE = lambda do |name, value|
           if name == :controller