require 'enumerator'

module Enumerable

  # Like <tt>#map</tt>/<tt>#collect</tt>, but generates a Hash.  The block
  # is expected to return two values: the key and the value for the new hash.
  #
  #   numbers  = (1..3)
  #   squares  = numbers.collate { |n| [n, n*n] }   # { 1=>1, 2=>4, 3=>9 }
  #   sq_roots = numbers.collate { |n| [n*n, n] }   # { 1=>1, 4=>2, 9=>3 }
  #
  # CREDIT: Trans
  # CREDIT: Andrew Dudzik (adudzik)

  def collate(&yld)
    if yld
      inject({}) do |h, *kv| # Used to be inject({}) do |h,kv|
        r = *yld[*kv]        # The *-op works differnt from to_a on single element hash!!!
        nk, nv = *r          # Used to be nk, nv = *yld[*kv].to_a.flatten
        h[nk] = nv
        h
      end
    else
      Enumerator.new(self,:collate)  # Used to be Hash[*self.to_a] or Hash[*self.to_a.flatten]
    end
  end

  alias :graph :collate

  # Returns a new hash built by iterating through each key,value
  # pair and updating the new hash.
  #def collate  #:yield:
  #  newhash = {}
  #  each_pair{ |k,v| newhash.update(yield(k,v)) }
  #  newhash
  #end

end


class Hash

  # In place version of #collate.
  #
  #   NOTE Hash#collate! is only useful for Hash. It is not generally
  #        applicable to Enumerable.

  def collate!(&yld)
    replace(collate(&yld))
  end

  # Alias for #collate! (THIS WILL BE DEPRECATED)

  alias :graph! :collate!

end


#  _____         _
# |_   _|__  ___| |_
#   | |/ _ \/ __| __|
#   | |  __/\__ \ |_
#   |_|\___||___/\__|
#
=begin test

  require 'test/unit'

  class TestEnumerableCollate < Test::Unit::TestCase

    def test_collate_01
      a = { :a => 1, :b => 2, :c => 3 }
      e = { :a => 2, :b => 3, :c => 4 }
      assert_equal( e, a.collate{ |k,v| { k => v+1 } } )
    end

    def test_collate_02
      numbers  = (1..3)
      squares  = numbers.collate{ |n| [n, n*n] }
      assert_equal( {1=>1, 2=>4, 3=>9}, squares )
    end

    def test_collate_03
      numbers  = (1..3)
      sq_roots = numbers.collate{ |n| [n*n, n] }
      assert_equal( {1=>1, 4=>2, 9=>3}, sq_roots )
    end

    def test_collate_01!
      a = { :a => 1, :b => 2, :c => 3 }
      e = { :a => 2, :b => 3, :c => 4 }
      a.collate!{ |k,v| { k => v+1 } }
      assert_equal( e, a )
    end

    def test_collate_02!
      h  = {:a=>1,:b=>2,:c=>3}
      h.collate!{ |k,v| [v, v*v] }
      assert_equal( {1=>1, 2=>4, 3=>9}, h )
    end

  end

=end