Array#first and Array#last in Ruby

ruby · Mar 23, 2023 · 0112a5b · 0112a5b
1 parent c9fd81b
commit 0112a5b
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 17 deletions.
diff --git a/array.c b/array.c
@@ -1428,20 +1428,11 @@ enum ary_take_pos_flags
 };
 
 static VALUE
-ary_take_first_or_last(int argc, const VALUE *argv, VALUE ary, enum ary_take_pos_flags last)
+ary_take_first_or_last_n(VALUE ary, long n, enum ary_take_pos_flags last)
 {
-    long n;
-    long len;
+    long len = RARRAY_LEN(ary);
     long offset = 0;
 
-    argc = rb_check_arity(argc, 0, 1);
-    /* the case optional argument is omitted should be handled in
-     * callers of this function.  if another arity case is added,
-     * this arity check needs to rewrite. */
-    RUBY_ASSERT_ALWAYS(argc == 1);
-
-    n = NUM2LONG(argv[0]);
-    len = RARRAY_LEN(ary);
     if (n > len) {
         n = len;
     }
@@ -1454,6 +1445,17 @@ ary_take_first_or_last(int argc, const VALUE *argv, VALUE ary, enum ary_take_pos
     return ary_make_partial(ary, rb_cArray, offset, n);
 }
 
+static VALUE
+ary_take_first_or_last(int argc, const VALUE *argv, VALUE ary, enum ary_take_pos_flags last)
+{
+    argc = rb_check_arity(argc, 0, 1);
+    /* the case optional argument is omitted should be handled in
+     * callers of this function.  if another arity case is added,
+     * this arity check needs to rewrite. */
+    RUBY_ASSERT_ALWAYS(argc == 1);
+    return ary_take_first_or_last_n(ary, NUM2LONG(argv[0]), last);
+}
+
 /*
  *  call-seq:
  *    array << object -> self
@@ -2027,6 +2029,7 @@ rb_ary_at(VALUE ary, VALUE pos)
     return rb_ary_entry(ary, NUM2LONG(pos));
 }
 
+#if 0
 /*
  *  call-seq:
  *    array.first -> object or nil
@@ -2071,6 +2074,20 @@ rb_ary_first(int argc, VALUE *argv, VALUE ary)
         return ary_take_first_or_last(argc, argv, ary, ARY_TAKE_FIRST);
     }
 }
+#endif
+
+static VALUE
+ary_first(VALUE self)
+{
+    return (RARRAY_LEN(self) == 0) ? Qnil : RARRAY_AREF(self, 0);
+}
+
+static VALUE
+ary_last(VALUE self)
+{
+    long len = RARRAY_LEN(self);
+    return (len == 0) ? Qnil : RARRAY_AREF(self, len-1);
+}
 
 /*
  *  call-seq:
@@ -2107,12 +2124,10 @@ rb_ary_first(int argc, VALUE *argv, VALUE ary)
  */
 
 VALUE
-rb_ary_last(int argc, const VALUE *argv, VALUE ary)
+rb_ary_last(int argc, const VALUE *argv, VALUE ary) // used by parse.y
 {
     if (argc == 0) {
-        long len = RARRAY_LEN(ary);
-        if (len == 0) return Qnil;
-        return RARRAY_AREF(ary, len-1);
+        return ary_last(ary);
     }
     else {
         return ary_take_first_or_last(argc, argv, ary, ARY_TAKE_LAST);
@@ -8809,8 +8824,6 @@ Init_Array(void)
     rb_define_method(rb_cArray, "[]=", rb_ary_aset, -1);
     rb_define_method(rb_cArray, "at", rb_ary_at, 1);
     rb_define_method(rb_cArray, "fetch", rb_ary_fetch, -1);
-    rb_define_method(rb_cArray, "first", rb_ary_first, -1);
-    rb_define_method(rb_cArray, "last", rb_ary_last, -1);
     rb_define_method(rb_cArray, "concat", rb_ary_concat_multi, -1);
     rb_define_method(rb_cArray, "union", rb_ary_union_multi, -1);
     rb_define_method(rb_cArray, "difference", rb_ary_difference_multi, -1);

diff --git a/array.rb b/array.rb
@@ -66,4 +66,30 @@ def sample(n = (ary = false), random: Random)
       Primitive.ary_sample(random, n, ary)
     end
   end
+
+  def first n = unspecified = true
+    if Primitive.mandatory_only?
+      Primitive.attr! :leaf
+      Primitive.cexpr! %q{ ary_first(self) }
+    else
+      if unspecified
+        Primitive.cexpr! %q{ ary_first(self) }
+      else
+        Primitive.cexpr! %q{  ary_take_first_or_last_n(self, NUM2LONG(n), ARY_TAKE_FIRST) }
+      end
+    end
+  end
+
+  def last n = unspecified = true
+    if Primitive.mandatory_only?
+      Primitive.attr! :leaf
+      Primitive.cexpr! %q{ ary_last(self) }
+    else
+      if unspecified
+        Primitive.cexpr! %q{ ary_last(self) }
+      else
+        Primitive.cexpr! %q{ ary_take_first_or_last_n(self, NUM2LONG(n), ARY_TAKE_LAST) }
+      end
+    end
+  end
 end