From 17c0aff0d6d8a2629b535a52a93e348904e27b97 Mon Sep 17 00:00:00 2001
From: marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Tue, 6 Nov 2012 17:13:48 +0000
Subject: [PATCH] * hash.c: Support for enumerators created by ENV:          
 each, each_value, ...   [Feature #6636]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37512 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 hash.c                       | 34 ++++++++++++++++++++++++++--------
 test/ruby/test_enumerator.rb |  6 ++++++
 2 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/hash.c b/hash.c
index 126fd41dc8..b8acbc8491 100644
--- a/hash.c
+++ b/hash.c
@@ -2548,13 +2548,31 @@ env_keys(void)
  *
  * An Enumerator is returned if no block is given.
  */
+static VALUE
+rb_env_size(VALUE ehash)
+{
+    char **env;
+    long cnt = 0;
+
+    rb_secure(4);
+
+    env = GET_ENVIRON(environ);
+    for (; *env ; ++env) {
+	if (strchr(*env, '=')) {
+	    cnt++;
+	}
+    }
+    FREE_ENVIRON(environ);
+    return LONG2FIX(cnt);
+}
+
 static VALUE
 env_each_key(VALUE ehash)
 {
     VALUE keys;
     long i;
 
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     keys = env_keys();	/* rb_secure(4); */
     for (i=0; i<RARRAY_LEN(keys); i++) {
 	rb_yield(RARRAY_PTR(keys)[i]);
@@ -2603,7 +2621,7 @@ env_each_value(VALUE ehash)
     VALUE values;
     long i;
 
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     values = env_values();	/* rb_secure(4); */
     for (i=0; i<RARRAY_LEN(values); i++) {
 	rb_yield(RARRAY_PTR(values)[i]);
@@ -2629,7 +2647,7 @@ env_each_pair(VALUE ehash)
     VALUE ary;
     long i;
 
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
 
     rb_secure(4);
     ary = rb_ary_new();
@@ -2666,7 +2684,7 @@ env_reject_bang(VALUE ehash)
     long i;
     int del = 0;
 
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     keys = env_keys();	/* rb_secure(4); */
     for (i=0; i<RARRAY_LEN(keys); i++) {
 	VALUE val = rb_f_getenv(Qnil, RARRAY_PTR(keys)[i]);
@@ -2694,7 +2712,7 @@ env_reject_bang(VALUE ehash)
 static VALUE
 env_delete_if(VALUE ehash)
 {
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     env_reject_bang(ehash);
     return envtbl;
 }
@@ -2735,7 +2753,7 @@ env_select(VALUE ehash)
     VALUE result;
     char **env;
 
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     rb_secure(4);
     result = rb_hash_new();
     env = GET_ENVIRON(environ);
@@ -2769,7 +2787,7 @@ env_select_bang(VALUE ehash)
     long i;
     int del = 0;
 
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     keys = env_keys();	/* rb_secure(4); */
     for (i=0; i<RARRAY_LEN(keys); i++) {
 	VALUE val = rb_f_getenv(Qnil, RARRAY_PTR(keys)[i]);
@@ -2797,7 +2815,7 @@ env_select_bang(VALUE ehash)
 static VALUE
 env_keep_if(VALUE ehash)
 {
-    RETURN_ENUMERATOR(ehash, 0, 0);
+    RETURN_SIZED_ENUMERATOR(ehash, 0, 0, rb_env_size);
     env_select_bang(ehash);
     return envtbl;
 }
diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb
index af41341227..884f8c6b27 100644
--- a/test/ruby/test_enumerator.rb
+++ b/test/ruby/test_enumerator.rb
@@ -450,6 +450,12 @@ class TestEnumerator < Test::Unit::TestCase
     end
   end
 
+  def test_size_for_enum_created_from_env
+    %i[each_pair reject! delete_if select select! keep_if].each do |method|
+      assert_equal ENV.size, ENV.send(method).size
+    end
+  end
+
   def check_consistency_for_combinatorics(method)
     [ [], [:a, :b, :c, :d, :e] ].product([-2, 0, 2, 5, 6]) do |array, arg|
       assert_equal array.send(method, arg).to_a.size, array.send(method, arg).size,