diff --git a/lib/puppet/functions/map.rb b/lib/puppet/functions/map.rb
index 7acd89a2b..2141d1e81 100644
--- a/lib/puppet/functions/map.rb
+++ b/lib/puppet/functions/map.rb
@@ -1,78 +1,97 @@
 # Applies a parameterized block to each element in a sequence of entries from the first
 # argument and returns an array with the result of each invocation of the parameterized block.
 #
 # This function takes two mandatory arguments: the first should be an Array, Hash, or of Enumerable type
 # (integer, Integer range, or String), and the second a parameterized block as produced by the puppet syntax:
 #
 #       $a.map |$x| { ... }
 #       map($a) |$x| { ... }
 #
 # When the first argument `$a` is an Array or of enumerable type, the block is called with each entry in turn.
 # When the first argument is a hash the entry is an array with `[key, value]`.
 #
 # @example Using map with two arguments
 #
 #      # Turns hash into array of values
 #      $a.map |$x|{ $x[1] }
 #
 #      # Turns hash into array of keys
 #      $a.map |$x| { $x[0] }
 #
 # When using a block with 2 parameters, the element's index (starting from 0) for an array, and the key for a hash
 # is given to the block's first parameter, and the value is given to the block's second parameter.args.
 #
 # @example Using map with two arguments
 #
 #      # Turns hash into array of values
 #      $a.map |$key,$val|{ $val }
 #
 #      # Turns hash into array of keys
 #      $a.map |$key,$val|{ $key }
 #
 # @since 3.4 for Array and Hash
 # @since 3.5 for other enumerables, and support for blocks with 2 parameters
 # @note requires `parser = future`
 #
 Puppet::Functions.create_function(:map) do
-  dispatch :map_Hash do
+  dispatch :map_Hash_2 do
     param 'Hash[Any, Any]', :hash
-    required_block_param
+    required_block_param 'Callable[2,2]', :block
   end
 
-  dispatch :map_Enumerable do
+  dispatch :map_Hash_1 do
+    param 'Hash[Any, Any]', :hash
+    required_block_param 'Callable[1,1]', :block
+  end
+
+  dispatch :map_Enumerable_2 do
     param 'Any', :enumerable
-    required_block_param
+    required_block_param 'Callable[2,2]', :block
   end
 
-  require 'puppet/util/functions/iterative_support'
-  include Puppet::Util::Functions::IterativeSupport
+  dispatch :map_Enumerable_1 do
+    param 'Any', :enumerable
+    required_block_param 'Callable[1,1]', :block
+  end
 
-  def map_Hash(hash, pblock)
-    if asserted_serving_size(pblock, 'key') == 1
-      hash.map {|x, y| pblock.call(nil, [x, y]) }
-    else
+  def map_Hash_1(hash, pblock)
+    hash.map {|x, y| pblock.call(nil, [x, y]) }
+  end
+
+  def map_Hash_2(hash, pblock)
       hash.map {|x, y| pblock.call(nil, x, y) }
+  end
+
+  def map_Enumerable_1(enumerable, pblock)
+    result = []
+    index = 0
+    enum = asserted_enumerable(enumerable)
+    begin
+      loop { result << pblock.call(nil, enum.next) }
+    rescue StopIteration
     end
+    result
   end
 
-  def map_Enumerable(enumerable, pblock)
+  def map_Enumerable_2(enumerable, pblock)
     result = []
     index = 0
     enum = asserted_enumerable(enumerable)
-    if asserted_serving_size(pblock, 'index') == 1
-      begin
-        loop { result << pblock.call(nil, enum.next) }
-      rescue StopIteration
-      end
-    else
-      begin
-        loop do
-          result << pblock.call(nil, index, enum.next)
-          index = index +1
-        end
-      rescue StopIteration
+    begin
+      loop do
+        result << pblock.call(nil, index, enum.next)
+        index = index +1
       end
+    rescue StopIteration
     end
     result
   end
+
+  def asserted_enumerable(obj)
+    unless enum = Puppet::Pops::Types::Enumeration.enumerator(obj)
+      raise ArgumentError, ("#{self.class.name}(): wrong argument type (#{obj.class}; must be something enumerable.")
+    end
+    enum
+  end
+
 end
diff --git a/spec/unit/functions/map_spec.rb b/spec/unit/functions/map_spec.rb
index 6f54648d9..9077cce17 100644
--- a/spec/unit/functions/map_spec.rb
+++ b/spec/unit/functions/map_spec.rb
@@ -1,209 +1,191 @@
 require 'puppet'
 require 'spec_helper'
 require 'puppet_spec/compiler'
 require 'matchers/resource'
 
 require 'unit/functions/shared'
 
 describe 'the map method' do
   include PuppetSpec::Compiler
   include Matchers::Resource
 
   before :each do
     Puppet[:parser] = "future"
   end
 
   context "using future parser" do
     it 'map on an array (multiplying each value by 2)' do
       catalog = compile_to_catalog(<<-MANIFEST)
         $a = [1,2,3]
         $a.map |$x|{ $x*2}.each |$v|{
           file { "/file_$v": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_2]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_4]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_6]").with_parameter(:ensure, 'present')
     end
 
     it 'map on an enumerable type (multiplying each value by 2)' do
       catalog = compile_to_catalog(<<-MANIFEST)
         $a = Integer[1,3]
         $a.map |$x|{ $x*2}.each |$v|{
           file { "/file_$v": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_2]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_4]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_6]").with_parameter(:ensure, 'present')
     end
 
     it 'map on an integer (multiply each by 3)' do
       catalog = compile_to_catalog(<<-MANIFEST)
         3.map |$x|{ $x*3}.each |$v|{
           file { "/file_$v": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_0]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_3]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_6]").with_parameter(:ensure, 'present')
     end
 
     it 'map on a string' do
       catalog = compile_to_catalog(<<-MANIFEST)
         $a = {a=>x, b=>y}
         "ab".map |$x|{$a[$x]}.each |$v|{
           file { "/file_$v": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_x]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_y]").with_parameter(:ensure, 'present')
     end
 
     it 'map on an array (multiplying value by 10 in even index position)' do
       catalog = compile_to_catalog(<<-MANIFEST)
         $a = [1,2,3]
         $a.map |$i, $x|{ if $i % 2 == 0 {$x} else {$x*10}}.each |$v|{
           file { "/file_$v": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_1]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_20]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_3]").with_parameter(:ensure, 'present')
     end
 
     it 'map on a hash selecting keys' do
       catalog = compile_to_catalog(<<-MANIFEST)
       $a = {'a'=>1,'b'=>2,'c'=>3}
       $a.map |$x|{ $x[0]}.each |$k|{
           file { "/file_$k": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_a]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_b]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_c]").with_parameter(:ensure, 'present')
     end
 
     it 'map on a hash selecting keys - using two block parameters' do
       catalog = compile_to_catalog(<<-MANIFEST)
       $a = {'a'=>1,'b'=>2,'c'=>3}
       $a.map |$k,$v|{ file { "/file_$k": ensure => present }
         }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_a]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_b]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_c]").with_parameter(:ensure, 'present')
     end
 
     it 'map on a hash using captures-last parameter' do
       catalog = compile_to_catalog(<<-MANIFEST)
       $a = {'a'=>present,'b'=>absent,'c'=>present}
       $a.map |*$kv|{ file { "/file_${kv[0]}": ensure => $kv[1] } }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_a]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_b]").with_parameter(:ensure, 'absent')
       expect(catalog).to have_resource("File[/file_c]").with_parameter(:ensure, 'present')
     end
 
     it 'each on a hash selecting value' do
       catalog = compile_to_catalog(<<-MANIFEST)
       $a = {'a'=>1,'b'=>2,'c'=>3}
       $a.map |$x|{ $x[1]}.each |$k|{ file { "/file_$k": ensure => present } }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_1]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_2]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_3]").with_parameter(:ensure, 'present')
     end
 
     it 'each on a hash selecting value - using two block parameters' do
       catalog = compile_to_catalog(<<-MANIFEST)
       $a = {'a'=>1,'b'=>2,'c'=>3}
       $a.map |$k,$v|{ file { "/file_$v": ensure => present } }
       MANIFEST
 
       expect(catalog).to have_resource("File[/file_1]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_2]").with_parameter(:ensure, 'present')
       expect(catalog).to have_resource("File[/file_3]").with_parameter(:ensure, 'present')
     end
 
     context "handles data type corner cases" do
       it "map gets values that are false" do
         catalog = compile_to_catalog(<<-MANIFEST)
           $a = [false,false]
           $a.map |$x| { $x }.each |$i, $v| {
             file { "/file_$i.$v": ensure => present }
           }
         MANIFEST
 
         expect(catalog).to have_resource("File[/file_0.false]").with_parameter(:ensure, 'present')
         expect(catalog).to have_resource("File[/file_1.false]").with_parameter(:ensure, 'present')
       end
 
       it "map gets values that are nil" do
         Puppet::Parser::Functions.newfunction(:nil_array, :type => :rvalue) do |args|
           [nil]
         end
         catalog = compile_to_catalog(<<-MANIFEST)
           $a = nil_array()
           $a.map |$x| { $x }.each |$i, $v| {
             file { "/file_$i.$v": ensure => present }
           }
         MANIFEST
 
         expect(catalog).to have_resource("File[/file_0.]").with_parameter(:ensure, 'present')
       end
 
       it "map gets values that are undef" do
         pending "Test is flawed, but has good intentions - should be rewritten when map has moved to new func API"
         # The test is broken because:
         # - a bug caused the given value to always be overridden by a given lambda default
         # - non existing variable results in nil / undef, which is transformed to empty string in the 3x func API
         # - when lambda is called, it gets an empty string, and it is then expected to use the default value
         #
         # This is not the semantics we want (only missing argument should trigger the default value).
         # Finally, it is not possible to test missing arguments with the map function since the call adapts itself
         # to the number of lambda parameters. (There is testing of this elsewhere).
         #
         # TODO: Rewrite map function, then test that undef / nil values are passed correctly to the lambda
         #
         catalog = compile_to_catalog(<<-MANIFEST)
           $a = [$does_not_exist]
           $a.map |$x = "something"| { $x }.each |$i, $v| {
             file { "/file_$i.$v": ensure => present }
           }
         MANIFEST
         catalog.resource(:file, "/file_0.something")['ensure'].should == 'present'
       end
     end
 
-  context 'map checks arguments and' do
-    it 'raises an error when block has more than 2 argument' do
-      expect do
-        compile_to_catalog(<<-MANIFEST)
-          [1].map |$index, $x, $yikes|{  }
-        MANIFEST
-      end.to raise_error(Puppet::Error, /block must define at most two parameters/)
-    end
-
-    it 'raises an error when block has fewer than 1 argument' do
-      expect do
-        compile_to_catalog(<<-MANIFEST)
-          [1].map || {  }
-        MANIFEST
-      end.to raise_error(Puppet::Error, /block must define at least one parameter/)
-    end
-  end
-
   it_should_behave_like 'all iterative functions argument checks', 'map'
   it_should_behave_like 'all iterative functions hash handling', 'map'
   end
 end