From 44374034268913246d517ff25365e0bccb453e02 Mon Sep 17 00:00:00 2001
From: marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Tue, 6 Nov 2012 17:15:45 +0000
Subject: * string.c: Support for
 String#{each_byte,each_char,each_codepoint}.size   [Feature #6636]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37520 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 string.c                     | 25 ++++++++++++++++++++++---
 test/ruby/test_enumerator.rb |  6 ++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/string.c b/string.c
index d674e8903c..224b018374 100644
--- a/string.c
+++ b/string.c
@@ -6222,6 +6222,11 @@ rb_str_each_line(int argc, VALUE *argv, VALUE str)
     return orig;
 }
 
+static VALUE
+rb_str_each_byte_size(VALUE str, VALUE args)
+{
+    return LONG2FIX(RSTRING_LEN(str));
+}
 
 /*
  *  call-seq:
@@ -6246,13 +6251,27 @@ rb_str_each_byte(VALUE str)
 {
     long i;
 
-    RETURN_ENUMERATOR(str, 0, 0);
+    RETURN_SIZED_ENUMERATOR(str, 0, 0, rb_str_each_byte_size);
     for (i=0; i<RSTRING_LEN(str); i++) {
 	rb_yield(INT2FIX(RSTRING_PTR(str)[i] & 0xff));
     }
     return str;
 }
 
+static VALUE
+rb_str_each_char_size(VALUE str)
+{
+    long len = RSTRING_LEN(str);
+    if (!single_byte_optimizable(str)) {
+	const char *ptr = RSTRING_PTR(str);
+	rb_encoding *enc = rb_enc_get(str);
+	const char *end_ptr = ptr + len;
+	for (len = 0; ptr < end_ptr; ++len) {
+	    ptr += rb_enc_mbclen(ptr, end_ptr, enc);
+	}
+    }
+    return LONG2FIX(len);
+}
 
 /*
  *  call-seq:
@@ -6280,7 +6299,7 @@ rb_str_each_char(VALUE str)
     const char *ptr;
     rb_encoding *enc;
 
-    RETURN_ENUMERATOR(str, 0, 0);
+    RETURN_SIZED_ENUMERATOR(str, 0, 0, rb_str_each_char_size);
     str = rb_str_new4(str);
     ptr = RSTRING_PTR(str);
     len = RSTRING_LEN(str);
@@ -6333,7 +6352,7 @@ rb_str_each_codepoint(VALUE str)
     rb_encoding *enc;
 
     if (single_byte_optimizable(str)) return rb_str_each_byte(str);
-    RETURN_ENUMERATOR(str, 0, 0);
+    RETURN_SIZED_ENUMERATOR(str, 0, 0, rb_str_each_char_size);
     str = rb_str_new4(str);
     ptr = RSTRING_PTR(str);
     end = RSTRING_END(str);
diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb
index 2c5158e726..896bf3f3f2 100644
--- a/test/ruby/test_enumerator.rb
+++ b/test/ruby/test_enumerator.rb
@@ -550,5 +550,11 @@ class TestEnumerator < Test::Unit::TestCase
     assert_equal 102, 1.downto(-100).size
     assert_equal Float::INFINITY, 42.upto(Float::INFINITY).size
   end
+
+  def test_size_for_string
+    assert_equal 5, 'hello'.each_byte.size
+    assert_equal 5, 'hello'.each_char.size
+    assert_equal 5, 'hello'.each_codepoint.size
+  end
 end
 
-- 
cgit v1.2.3