mrbgems/mruby-array-ext/src/array.c

   1 #include <mruby.h>
   2 #include <mruby/value.h>
   3 #include <mruby/array.h>
   4 #include <mruby/range.h>
   5 #include <mruby/hash.h>
   6 #include <mruby/internal.h>
   7 #include <mruby/presym.h>
   8
   9 /*
  10  *  call-seq:
  11  *     ary.assoc(obj)   -> new_ary  or  nil
  12  *
  13  *  Searches through an array whose elements are also arrays
  14  *  comparing _obj_ with the first element of each contained array
  15  *  using obj.==.
  16  *  Returns the first contained array that matches (that
  17  *  is, the first associated array),
  18  *  or +nil+ if no match is found.
  19  *  See also <code>Array#rassoc</code>.
  20  *
  21  *     s1 = [ "colors", "red", "blue", "green" ]
  22  *     s2 = [ "letters", "a", "b", "c" ]
  23  *     s3 = "foo"
  24  *     a  = [ s1, s2, s3 ]
  25  *     a.assoc("letters")  #=> [ "letters", "a", "b", "c" ]
  26  *     a.assoc("foo")      #=> nil
  27  */
  28
  29 static mrb_value
  30 mrb_ary_assoc(mrb_state *mrb, mrb_value ary)
  31 {
  32   mrb_int i;
  33   mrb_value v;
  34   mrb_value k = mrb_get_arg1(mrb);
  35
  36   for (i = 0; i < RARRAY_LEN(ary); ++i) {
  37     v = mrb_check_array_type(mrb, RARRAY_PTR(ary)[i]);
  38     if (!mrb_nil_p(v) && RARRAY_LEN(v) > 0 &&
  39         mrb_equal(mrb, RARRAY_PTR(v)[0], k))
  40       return v;
  41   }
  42   return mrb_nil_value();
  43 }
  44
  45 /*
  46  *  call-seq:
  47  *     ary.rassoc(obj) -> new_ary or nil
  48  *
  49  *  Searches through the array whose elements are also arrays. Compares
  50  *  _obj_ with the second element of each contained array using
  51  *  <code>==</code>. Returns the first contained array that matches. See
  52  *  also <code>Array#assoc</code>.
  53  *
  54  *     a = [ [ 1, "one"], [2, "two"], [3, "three"], ["ii", "two"] ]
  55  *     a.rassoc("two")    #=> [2, "two"]
  56  *     a.rassoc("four")   #=> nil
  57  */
  58
  59 static mrb_value
  60 mrb_ary_rassoc(mrb_state *mrb, mrb_value ary)
  61 {
  62   mrb_int i;
  63   mrb_value v;
  64   mrb_value value = mrb_get_arg1(mrb);
  65
  66   for (i = 0; i < RARRAY_LEN(ary); ++i) {
  67     v = RARRAY_PTR(ary)[i];
  68     if (mrb_array_p(v) &&
  69         RARRAY_LEN(v) > 1 &&
  70         mrb_equal(mrb, RARRAY_PTR(v)[1], value))
  71       return v;
  72   }
  73   return mrb_nil_value();
  74 }
  75
  76 /*
  77  *  call-seq:
  78  *     ary.at(index)   ->   obj  or nil
  79  *
  80  *  Returns the element at _index_. A
  81  *  negative index counts from the end of +self+.  Returns +nil+
  82  *  if the index is out of range. See also <code>Array#[]</code>.
  83  *
  84  *     a = [ "a", "b", "c", "d", "e" ]
  85  *     a.at(0)     #=> "a"
  86  *     a.at(-1)    #=> "e"
  87  */
  88
  89 static mrb_value
  90 mrb_ary_at(mrb_state *mrb, mrb_value ary)
  91 {
  92   mrb_int pos;
  93   mrb_get_args(mrb, "i", &pos);
  94
  95   return mrb_ary_entry(ary, pos);
  96 }
  97
  98 static mrb_value
  99 ary_ref(mrb_state *mrb, mrb_value ary, mrb_int n)
 100 {
 101   return mrb_ary_entry(ary, n);
 102 }
 103
 104 static mrb_value
 105 mrb_ary_values_at(mrb_state *mrb, mrb_value self)
 106 {
 107   mrb_int argc = mrb_get_argc(mrb);
 108   const mrb_value *argv = mrb_get_argv(mrb);
 109
 110   return mrb_get_values_at(mrb, self, RARRAY_LEN(self), argc, argv, ary_ref);
 111 }
 112
 113 mrb_value mrb_ary_delete_at(mrb_state *mrb, mrb_value self);
 114
 115 /*
 116  *  call-seq:
 117  *     ary.slice!(index)         -> obj or nil
 118  *     ary.slice!(start, length) -> new_ary or nil
 119  *     ary.slice!(range)         -> new_ary or nil
 120  *
 121  *  Deletes the element(s) given by an +index+ (optionally up to +length+
 122  *  elements) or by a +range+.
 123  *
 124  *  Returns the deleted object (or objects), or +nil+ if the +index+ is out of
 125  *  range.
 126  *
 127  *     a = [ "a", "b", "c" ]
 128  *     a.slice!(1)     #=> "b"
 129  *     a               #=> ["a", "c"]
 130  *     a.slice!(-1)    #=> "c"
 131  *     a               #=> ["a"]
 132  *     a.slice!(100)   #=> nil
 133  *     a               #=> ["a"]
 134  */
 135
 136 static mrb_value
 137 mrb_ary_slice_bang(mrb_state *mrb, mrb_value self)
 138 {
 139   struct RArray *a = mrb_ary_ptr(self);
 140   mrb_int i, j, len, alen;
 141   mrb_value *ptr;
 142   mrb_value ary;
 143
 144   mrb_ary_modify(mrb, a);
 145
 146   if (mrb_get_argc(mrb) == 1) {
 147     mrb_value index = mrb_get_arg1(mrb);
 148
 149     if (mrb_type(index) == MRB_TT_RANGE) {
 150       if (mrb_range_beg_len(mrb, index, &i, &len, ARY_LEN(a), TRUE) == MRB_RANGE_OK) {
 151         goto delete_pos_len;
 152       }
 153       return mrb_nil_value();
 154     }
 155     return mrb_ary_delete_at(mrb, self);
 156   }
 157
 158   mrb_get_args(mrb, "ii", &i, &len);
 159  delete_pos_len:
 160   alen = ARY_LEN(a);
 161   if (i < 0) i += alen;
 162   if (i < 0 || alen < i) return mrb_nil_value();
 163   if (len < 0) return mrb_nil_value();
 164   if (alen == i) return mrb_ary_new(mrb);
 165   if (len > alen - i) len = alen - i;
 166
 167   ptr = ARY_PTR(a) + i;
 168   ary = mrb_ary_new_from_values(mrb, len, ptr);
 169
 170   for (j = i; j < alen - len; ++j) {
 171     *ptr = *(ptr+len);
 172     ++ptr;
 173   }
 174
 175   mrb_ary_resize(mrb, self, alen - len);
 176   return ary;
 177 }
 178
 179 /*
 180  * call-seq:
 181  *    ary.compact     -> new_ary
 182  *
 183  * Returns a copy of +self+ with all +nil+ elements removed.
 184  *
 185  *   [ "a", nil, "b", nil, "c", nil ].compact
 186  *                      #=> [ "a", "b", "c" ]
 187  */
 188
 189 static mrb_value
 190 mrb_ary_compact(mrb_state *mrb, mrb_value self)
 191 {
 192   mrb_value ary = mrb_ary_new(mrb);
 193   mrb_int len = RARRAY_LEN(self);
 194   mrb_value *p = RARRAY_PTR(self);
 195
 196   for (mrb_int i = 0; i < len; ++i) {
 197     if (!mrb_nil_p(p[i])) {
 198       mrb_ary_push(mrb, ary, p[i]);
 199     }
 200   }
 201   return ary;
 202 }
 203
 204 /*
 205  * call-seq:
 206  *    ary.compact!    -> ary  or  nil
 207  *
 208  * Removes +nil+ elements from the array.
 209  * Returns +nil+ if no changes were made, otherwise returns
 210  * <i>ary</i>.
 211  *
 212  *    [ "a", nil, "b", nil, "c" ].compact! #=> [ "a", "b", "c" ]
 213  *    [ "a", "b", "c" ].compact!           #=> nil
 214  */
 215 static mrb_value
 216 mrb_ary_compact_bang(mrb_state *mrb, mrb_value self)
 217 {
 218   struct RArray *a = mrb_ary_ptr(self);
 219   mrb_int i, j = 0;
 220   mrb_int len = ARY_LEN(a);
 221   mrb_value *p = ARY_PTR(a);
 222
 223   mrb_ary_modify(mrb, a);
 224   for (i = 0; i < len; ++i) {
 225     if (!mrb_nil_p(p[i])) {
 226       if (i != j) p[j] = p[i];
 227       j++;
 228     }
 229   }
 230   if (i == j) return mrb_nil_value();
 231   if (j < len) ARY_SET_LEN(RARRAY(self), j);
 232   return self;
 233 }
 234
 235
 236 /*
 237  *  call-seq:
 238  *     ary.rotate(count=1)    -> new_ary
 239  *
 240  *  Returns a new array by rotating +self+ so that the element at +count+ is
 241  *  the first element of the new array.
 242  *
 243  *  If +count+ is negative then it rotates in the opposite direction, starting
 244  *  from the end of +self+ where +-1+ is the last element.
 245  *
 246  *     a = [ "a", "b", "c", "d" ]
 247  *     a.rotate         #=> ["b", "c", "d", "a"]
 248  *     a                #=> ["a", "b", "c", "d"]
 249  *     a.rotate(2)      #=> ["c", "d", "a", "b"]
 250  *     a.rotate(-3)     #=> ["b", "c", "d", "a"]
 251  */
 252 static mrb_value
 253 mrb_ary_rotate(mrb_state *mrb, mrb_value self)
 254 {
 255   mrb_int count=1;
 256   mrb_get_args(mrb, "|i", &count);
 257
 258   mrb_value ary = mrb_ary_new(mrb);
 259   mrb_int len = RARRAY_LEN(self);
 260   mrb_value *p = RARRAY_PTR(self);
 261   mrb_int idx;
 262
 263   if (len <= 0) return ary;
 264   if (count < 0) {
 265     idx = len - (~count % len) - 1;
 266   }
 267   else {
 268     idx = count % len;
 269   }
 270   for (mrb_int i = 0; i<len; i++) {
 271     mrb_ary_push(mrb, ary, p[idx++]);
 272     if (idx == len) idx = 0;
 273   }
 274   return ary;
 275 }
 276
 277 static void
 278 rev(mrb_value *p, mrb_int beg, mrb_int end)
 279 {
 280   for (mrb_int i=beg,j=end-1; i<j; i++,j--) {
 281     mrb_value v = p[i];
 282     p[i] = p[j];
 283     p[j] = v;
 284   }
 285 }
 286
 287 /*
 288  *  call-seq:
 289  *     ary.rotate!(count=1)   -> ary
 290  *
 291  *  Rotates +self+ in place so that the element at +count+ comes first, and
 292  *  returns +self+.
 293  *
 294  *  If +count+ is negative then it rotates in the opposite direction, starting
 295  *  from the end of the array where +-1+ is the last element.
 296  *
 297  *     a = [ "a", "b", "c", "d" ]
 298  *     a.rotate!        #=> ["b", "c", "d", "a"]
 299  *     a                #=> ["b", "c", "d", "a"]
 300  *     a.rotate!(2)     #=> ["d", "a", "b", "c"]
 301  *     a.rotate!(-3)    #=> ["a", "b", "c", "d"]
 302  */
 303 static mrb_value
 304 mrb_ary_rotate_bang(mrb_state *mrb, mrb_value self)
 305 {
 306   mrb_int count=1;
 307   mrb_get_args(mrb, "|i", &count);
 308
 309   struct RArray *a = mrb_ary_ptr(self);
 310   mrb_int len = ARY_LEN(a);
 311   mrb_value *p = ARY_PTR(a);
 312   mrb_int idx;
 313
 314   mrb_ary_modify(mrb, a);
 315   if (len == 0 || count == 0) return self;
 316   if (count == 1) {
 317     mrb_value v = p[0];
 318     for (mrb_int i=1; i<len; i++) {
 319       p[i-1] = p[i];
 320     }
 321     p[len-1] = v;
 322     return self;
 323   }
 324   if (count < 0) {
 325     idx = len - (~count % len) - 1;
 326   }
 327   else {
 328     idx = count % len;
 329   }
 330   /* e.g. [1,2,3,4,5].rotate!(2) -> [3,4,5,1,2] */
 331   /* first, reverse the whole array */
 332   /* [1,2,3,4,5] -> [5,4,3,2,1] */
 333   rev(p, 0, len);
 334   /* then, re-reverse part before idx */
 335   /* [5,4,3,2,1] -> [3,4,5,2,1] */
 336   /*        ^idx     ~~~~~      */
 337   rev(p, 0, len-idx);
 338   /* finally, re-reverse part after idx */
 339   /* [3,4,5,2,1] -> [3,4,5,1,2] */
 340   /*        ^idx           ~~~  */
 341   rev(p, len-idx, len);
 342   return self;
 343 }
 344
 345 void
 346 mrb_mruby_array_ext_gem_init(mrb_state* mrb)
 347 {
 348   struct RClass * a = mrb->array_class;
 349
 350   mrb_define_method(mrb, a, "assoc",  mrb_ary_assoc,  MRB_ARGS_REQ(1));
 351   mrb_define_method(mrb, a, "at",     mrb_ary_at,     MRB_ARGS_REQ(1));
 352   mrb_define_method(mrb, a, "rassoc", mrb_ary_rassoc, MRB_ARGS_REQ(1));
 353   mrb_define_method(mrb, a, "values_at", mrb_ary_values_at, MRB_ARGS_ANY());
 354   mrb_define_method(mrb, a, "slice!", mrb_ary_slice_bang, MRB_ARGS_ARG(1,1));
 355   mrb_define_method(mrb, a, "compact", mrb_ary_compact, MRB_ARGS_NONE());
 356   mrb_define_method(mrb, a, "compact!", mrb_ary_compact_bang, MRB_ARGS_NONE());
 357   mrb_define_method(mrb, a, "rotate", mrb_ary_rotate, MRB_ARGS_OPT(1));
 358   mrb_define_method(mrb, a, "rotate!", mrb_ary_rotate_bang, MRB_ARGS_OPT(1));
 359 }
 360
 361 void
 362 mrb_mruby_array_ext_gem_final(mrb_state* mrb)
 363 {
 364 }