From 7a31096255cd226d16c918fa859e3e4a654e64b9 Mon Sep 17 00:00:00 2001
From: marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Tue, 6 Nov 2012 17:16:44 +0000
Subject: * enumerator.c: Add support for lazy.drop.size   [Feature #6636]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37524 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 enumerator.c | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

(limited to 'enumerator.c')

diff --git a/enumerator.c b/enumerator.c
index e5dd80a841..0a51709ac1 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -1673,6 +1673,19 @@ lazy_take_while(VALUE obj)
 			   Qnil, 0);
 }
 
+static VALUE
+lazy_drop_size(VALUE lazy) {
+    long len = NUM2LONG(RARRAY_PTR(rb_ivar_get(lazy, id_arguments))[0]);
+    VALUE receiver = lazy_receiver_size(lazy);
+    if (NIL_P(receiver))
+	return receiver;
+    if (FIXNUM_P(receiver)) {
+	len = FIX2LONG(receiver) - len;
+	return LONG2FIX(len < 0 ? 0 : len);
+    }
+    return rb_funcall(receiver, '-', 1, LONG2NUM(len));
+}
+
 static VALUE
 lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv)
 {
@@ -1699,7 +1712,7 @@ lazy_drop(VALUE obj, VALUE n)
     memo = NEW_MEMO(0, 0, len);
     return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
 					 lazy_drop_func, (VALUE) memo),
-			   rb_ary_new3(1, n), 0);
+			   rb_ary_new3(1, n), lazy_drop_size);
 }
 
 static VALUE
-- 
cgit v1.2.3