[Patch] Array#shift(n)

Hello,

please have a look into following patch.
I've changed:
- move alloc behind checks in ary_new() (some obj alloc time could be 
spared when args are wrong)
- make Ary#shift take arg to behave as Ary#unshift
- correct doc for Ary#first(n) and Ary#last(n)
- change Ary#select to not accept any arg (as it is, but has been marked 
as -1)

Is that good?

        Michal

Index: array.c
===================================================================
RCS file: /var/cvs/src/ruby/array.c,v
retrieving revision 1.148
diff -u -u -p -r1.148 array.c
--- array.c	11 Jun 2004 13:33:47 -0000	1.148
+++ array.c	14 Jun 2004 11:05:03 -0000
@@ -33,7 +33,7 @@ rb_mem_clear(mem, size)
     }
 }
 
-static void
+static inline void
 memfill(mem, size, val)
     register VALUE *mem;
     register long size;
@@ -72,6 +72,7 @@ rb_ary_modify(ary)
 	RARRAY(ary)->ptr = ptr;
     }
 }
+
 VALUE
 rb_ary_freeze(ary)
     VALUE ary;
@@ -96,7 +97,6 @@ rb_ary_frozen_p(ary)
     return Qfalse;
 }
 
-static VALUE ary_alloc _((VALUE));
 static VALUE
 ary_alloc(klass)
     VALUE klass;
@@ -116,7 +116,7 @@ ary_new(klass, len)
     VALUE klass;
     long len;
 {
-    VALUE ary = ary_alloc(klass);
+    VALUE ary;
 
     if (len < 0) {
 	rb_raise(rb_eArgError, "negative array size (or size too big)");
@@ -125,6 +125,8 @@ ary_new(klass, len)
 	rb_raise(rb_eArgError, "array size too big");
     }
     if (len == 0) len++;
+    
+    ary = ary_alloc(klass);
     RARRAY(ary)->ptr = ALLOC_N(VALUE, len);
     RARRAY(ary)->aux.capa = len;
 
@@ -146,11 +148,11 @@ rb_ary_new()
 }
 
 #ifdef HAVE_STDARG_PROTOTYPES
-#include <stdarg.h>
-#define va_init_list(a,b) va_start(a,b)
+#  include <stdarg.h>
+#  define va_init_list(a,b) va_start(a,b)
 #else
-#include <varargs.h>
-#define va_init_list(a,b) va_start(a)
+#  include <varargs.h>
+#  define va_init_list(a,b) va_start(a)
 #endif
 
 VALUE
@@ -544,6 +546,22 @@ rb_ary_pop(ary)
     return RARRAY(ary)->ptr[--RARRAY(ary)->len];
 }
 
+VALUE
+rb_ary_shift(ary)
+    VALUE ary;
+{
+    VALUE top;
+
+    rb_ary_modify_check(ary);
+    if (RARRAY(ary)->len == 0) return Qnil;
+    top = RARRAY(ary)->ptr[0];
+    ary_make_shared(ary);
+    RARRAY(ary)->ptr++;		/* shift ptr */
+    RARRAY(ary)->len--;
+
+    return top;
+}
+
 /*
  *  call-seq:
  *     array.shift   =>   obj or nil
@@ -553,24 +571,56 @@ rb_ary_pop(ary)
  *  is empty.
  *     
  *     args = [ "-m", "-q", "filename" ]
- *     args.shift   #=> "-m"
- *     args         #=> ["-q", "filename"]
+ *     args.shift     #=> "-m"
+ *     args           #=> ["-q", "filename"]
+ *
+ *     args = [ "-m", "-q", "filename" ]
+ *     args.shift(2)  #=> ["-m", "-q"]
+ *     args           #=> ["filename"]
  */
 
-VALUE
-rb_ary_shift(ary)
+static VALUE
+rb_ary_shift_m(argc, argv, ary)
+    int argc;
+    VALUE *argv;
     VALUE ary;
 {
-    VALUE top;
+    VALUE num, klass, ary2;
+    long shift = 1;
+
+    if (argc == 0) {
+	return rb_ary_shift(ary);
+    }
 
     rb_ary_modify_check(ary);
+
+    if (rb_scan_args(argc, argv, "01", &num)) {
+	shift = NUM2LONG(num);
+    }
     if (RARRAY(ary)->len == 0) return Qnil;
-    top = RARRAY(ary)->ptr[0];
+
+    if (shift < 0) {
+	rb_raise(rb_eArgError, "negative number of arguments");
+    }
+    if (shift == 0) return Qnil;
+    if (shift > RARRAY(ary)->len) {
+	shift = RARRAY(ary)->len;
+    }
+
+    klass = rb_obj_class(ary);
+
     ary_make_shared(ary);
-    RARRAY(ary)->ptr++;		/* shift ptr */
-    RARRAY(ary)->len--;
 
-    return top;
+    ary2 = ary_alloc(klass);
+    RARRAY(ary2)->ptr = RARRAY(ary)->ptr;
+    RARRAY(ary2)->len = shift;
+    RARRAY(ary2)->aux.shared = RARRAY(ary)->aux.shared;
+    FL_SET(ary2, ELTS_SHARED);
+
+    RARRAY(ary)->ptr += shift;		/* shift ptr */
+    RARRAY(ary)->len -= shift;
+
+    return ary2;
 }
 
 VALUE
@@ -777,13 +827,15 @@ rb_ary_at(ary, pos)
 
 /*
  *  call-seq:
- *     array.first   =>   obj or nil
+ *     array.first     =>   obj or nil
+ *     array.first(n)  =>   an_array
  *  
  *  Returns the first element of the array. If the array is empty,
  *  returns <code>nil</code>.
  *     
  *     a = [ "q", "r", "s", "t" ]
- *     a.first   #=> "q"
+ *     a.first     #=> "q"
+ *     a.first(2)  #=> ["q", "r"]
  */
 
 static VALUE
@@ -819,7 +871,9 @@ rb_ary_first(argc, argv, ary)
  *  Returns the last element(s) of <i>self</i>. If the array is empty,
  *  the first form returns <code>nil</code>.
  *     
- *     [ "w", "x", "y", "z" ].last   #=> "z"
+ *     a = [ "w", "x", "y", "z" ]
+ *     a.last     #=> "z"
+ *     a.last(2)  #=> ["y", "z"]
  */
 
 static VALUE
@@ -1802,17 +1856,12 @@ rb_ary_values_at(argc, argv, ary)
  */
 
 static VALUE
-rb_ary_select(argc, argv, ary)
-    int argc;
-    VALUE *argv;
+rb_ary_select(ary)
     VALUE ary;
 {
     VALUE result;
     long i;
 
-    if (argc > 0) {
-	rb_raise(rb_eArgError, "wrong number of arguments (%d for 0)", argc);
-    }
     result = rb_ary_new2(RARRAY(ary)->len);
     for (i = 0; i < RARRAY(ary)->len; i++) {
 	if (RTEST(rb_yield(RARRAY(ary)->ptr[i]))) {
@@ -2993,7 +3042,7 @@ Init_Array()
     rb_define_method(rb_cArray, "<<", rb_ary_push, 1);
     rb_define_method(rb_cArray, "push", rb_ary_push_m, -1);
     rb_define_method(rb_cArray, "pop", rb_ary_pop, 0);
-    rb_define_method(rb_cArray, "shift", rb_ary_shift, 0);
+    rb_define_method(rb_cArray, "shift", rb_ary_shift_m, -1);
     rb_define_method(rb_cArray, "unshift", rb_ary_unshift_m, -1);
     rb_define_method(rb_cArray, "insert", rb_ary_insert, -1);
     rb_define_method(rb_cArray, "each", rb_ary_each, 0);
@@ -3013,7 +3062,7 @@ Init_Array()
     rb_define_method(rb_cArray, "collect!", rb_ary_collect_bang, 0);
     rb_define_method(rb_cArray, "map", rb_ary_collect, 0);
     rb_define_method(rb_cArray, "map!", rb_ary_collect_bang, 0);
-    rb_define_method(rb_cArray, "select", rb_ary_select, -1);
+    rb_define_method(rb_cArray, "select", rb_ary_select, 0);
     rb_define_method(rb_cArray, "values_at", rb_ary_values_at, -1);
     rb_define_method(rb_cArray, "delete", rb_ary_delete, 1);
     rb_define_method(rb_cArray, "delete_at", rb_ary_delete_at_m, 1);