1: /**********************************************************************
    2: 
    3:   enum.c -
    4: 
    5:   $Author: matz $
    6:   $Date: 2007-12-25 15:22:27 +0900 (Tue, 25 Dec 2007) $
    7:   created at: Fri Oct  1 15:15:19 JST 1993
    8: 
    9:   Copyright (C) 1993-2007 Yukihiro Matsumoto
   10: 
   11: **********************************************************************/
   12: 
   13: #include "ruby/ruby.h"
   14: #include "ruby/node.h"
   15: #include "ruby/util.h"
   16: 
   17: VALUE rb_mEnumerable;
   18: static ID id_each, id_eqq, id_cmp, id_next;
   19: 
   20: static VALUE
   21: enum_values_pack(int argc, VALUE *argv)
   22: {
   23:     if (argc == 0) return Qnil;
   24:     if (argc == 1) return argv[0];
   25:     return rb_ary_new4(argc, argv);
   26: }
   27: 
   28: static VALUE
   29: enum_yield(int argc, VALUE *argv)
   30: {
   31:     return rb_yield(enum_values_pack(argc, argv));
   32: }
   33: 
   34: static VALUE
   35: grep_i(VALUE i, VALUE *arg, int argc, VALUE *argv)
   36: {
   37:     if (RTEST(rb_funcall(arg[0], id_eqq, 1, i))) {
   38:         rb_ary_push(arg[1], enum_values_pack(argc, argv));
   39:     }
   40:     return Qnil;
   41: }
   42: 
   43: static VALUE
   44: grep_iter_i(VALUE i, VALUE *arg, int argc, VALUE *argv)
   45: {
   46:     if (RTEST(rb_funcall(arg[0], id_eqq, 1, i))) {
   47:         rb_ary_push(arg[1], enum_yield(argc, argv));
   48:     }
   49:     return Qnil;
   50: }
   51: 
   52: /*
   53:  *  call-seq:
   54:  *     enum.grep(pattern)                   => array
   55:  *     enum.grep(pattern) {| obj | block }  => array
   56:  *  
   57:  *  Returns an array of every element in <i>enum</i> for which
   58:  *  <code>Pattern === element</code>. If the optional <em>block</em> is
   59:  *  supplied, each matching element is passed to it, and the block's
   60:  *  result is stored in the output array.
   61:  *     
   62:  *     (1..100).grep 38..44   #=> [38, 39, 40, 41, 42, 43, 44]
   63:  *     c = IO.constants
   64:  *     c.grep(/SEEK/)         #=> ["SEEK_END", "SEEK_SET", "SEEK_CUR"]
   65:  *     res = c.grep(/SEEK/) {|v| IO.const_get(v) }
   66:  *     res                    #=> [2, 0, 1]
   67:  *     
   68:  */
   69: 
   70: static VALUE
   71: enum_grep(VALUE obj, VALUE pat)
   72: {
   73:     VALUE ary = rb_ary_new();
   74:     VALUE arg[2];
   75: 
   76:     arg[0] = pat;
   77:     arg[1] = ary;
   78: 
   79:     rb_block_call(obj, id_each, 0, 0, rb_block_given_p() ? grep_iter_i : grep_i, (VALUE)arg);
   80:     
   81:     return ary;
   82: }
   83: 
   84: static VALUE
   85: count_i(VALUE i, VALUE *arg)
   86: {
   87:     if (rb_equal(i, arg[0])) {
   88:         arg[1]++;
   89:     }
   90:     return Qnil;
   91: }
   92: 
   93: static VALUE
   94: count_iter_i(VALUE i, long *n, int argc, VALUE *argv)
   95: {
   96:     if (RTEST(enum_yield(argc, argv))) {
   97:         (*n)++;
   98:     }
   99:     return Qnil;
  100: }
  101: 
  102: /*
  103:  *  call-seq:
  104:  *     enum.count(item)             => int
  105:  *     enum.count {| obj | block }  => int
  106:  *  
  107:  *  Returns the number of items in <i>enum</i> for which equals to <i>item</i>.
  108:  *  If a block is given, counts the number of elements yielding a true value.
  109:  *     
  110:  *     ary = [1, 2, 4, 2]
  111:  *     ary.count(2)          # => 2
  112:  *     ary.count{|x|x%2==0}  # => 3
  113:  *     
  114:  */
  115: 
  116: static VALUE
  117: enum_count(int argc, VALUE *argv, VALUE obj)
  118: {
  119:     if (argc == 1) {
  120:         VALUE item, args[2];
  121: 
  122:         if (rb_block_given_p()) {
  123:             rb_warn("given block not used");
  124:         }
  125:         rb_scan_args(argc, argv, "1", &item);
  126:         args[0] = item;
  127:         args[1] = 0;
  128:         rb_block_call(obj, id_each, 0, 0, count_i, (VALUE)&args);
  129:         return INT2NUM(args[1]);
  130:     }
  131:     else if (argc == 0) {
  132:         long n;
  133: 
  134:         RETURN_ENUMERATOR(obj, 0, 0);
  135:         n = 0;
  136:         rb_block_call(obj, id_each, 0, 0, count_iter_i, (VALUE)&n);
  137:         return INT2NUM(n);
  138:     }
  139:     else {
  140:         VALUE v;
  141:         rb_scan_args(argc, argv, "1", &v);
  142:         return Qnil; /* not reached */
  143:     }
  144: }
  145: 
  146: static VALUE
  147: find_i(VALUE i, VALUE *memo, int argc, VALUE *argv)
  148: {
  149:     if (RTEST(enum_yield(argc, argv))) {
  150:         *memo = i;
  151:         rb_iter_break();
  152:     }
  153:     return Qnil;
  154: }
  155: 
  156: /*
  157:  *  call-seq:
  158:  *     enum.detect(ifnone = nil) {| obj | block }  => obj or nil
  159:  *     enum.find(ifnone = nil)   {| obj | block }  => obj or nil
  160:  *  
  161:  *  Passes each entry in <i>enum</i> to <em>block</em>. Returns the
  162:  *  first for which <em>block</em> is not <code>false</code>.  If no
  163:  *  object matches, calls <i>ifnone</i> and returns its result when it
  164:  *  is specified, or returns <code>nil</code>
  165:  *     
  166:  *     (1..10).detect  {|i| i % 5 == 0 and i % 7 == 0 }   #=> nil
  167:  *     (1..100).detect {|i| i % 5 == 0 and i % 7 == 0 }   #=> 35
  168:  *     
  169:  */
  170: 
  171: static VALUE
  172: enum_find(int argc, VALUE *argv, VALUE obj)
  173: {
  174:     VALUE memo = Qundef;
  175:     VALUE if_none;
  176: 
  177:     rb_scan_args(argc, argv, "01", &if_none);
  178:     RETURN_ENUMERATOR(obj, argc, argv);
  179:     rb_block_call(obj, id_each, 0, 0, find_i, (VALUE)&memo);
  180:     if (memo != Qundef) {
  181:         return memo;
  182:     }
  183:     if (!NIL_P(if_none)) {
  184:         return rb_funcall(if_none, rb_intern("call"), 0, 0);
  185:     }
  186:     return Qnil;
  187: }
  188: 
  189: static VALUE
  190: find_index_i(VALUE i, VALUE *memo, int argc, VALUE *argv)
  191: {
  192:     if (RTEST(enum_yield(argc, argv))) {
  193:         memo[0] = UINT2NUM(memo[1]);
  194:         rb_iter_break();
  195:     }
  196:     memo[1]++;
  197:     return Qnil;
  198: }
  199: 
  200: /*
  201:  *  call-seq:
  202:  *     enum.find_index()   {| obj | block }  => int
  203:  *  
  204:  *  Passes each entry in <i>enum</i> to <em>block</em>. Returns the
  205:  *  index for the first for which <em>block</em> is not <code>false</code>.
  206:  *  If no object matches, returns <code>nil</code>
  207:  *     
  208:  *     (1..10).find_index  {|i| i % 5 == 0 and i % 7 == 0 }   #=> nil
  209:  *     (1..100).find_index {|i| i % 5 == 0 and i % 7 == 0 }   #=> 34
  210:  *     
  211:  */
  212: 
  213: static VALUE
  214: enum_find_index(VALUE obj)
  215: {
  216:     VALUE memo[2];
  217: 
  218:     RETURN_ENUMERATOR(obj, 0, 0);
  219:     memo[0] = Qundef;
  220:     memo[1] = 0;
  221:     rb_block_call(obj, id_each, 0, 0, find_index_i, (VALUE)memo);
  222:     if (memo[0] != Qundef) {
  223:         return memo[0];
  224:     }
  225:     return Qnil;
  226: }
  227: 
  228: static VALUE
  229: find_all_i(VALUE i, VALUE ary, int argc, VALUE *argv)
  230: {
  231:     if (RTEST(enum_yield(argc, argv))) {
  232:         rb_ary_push(ary, i);
  233:     }
  234:     return Qnil;
  235: }
  236: 
  237: /*
  238:  *  call-seq:
  239:  *     enum.find_all {| obj | block }  => array
  240:  *     enum.select   {| obj | block }  => array
  241:  *  
  242:  *  Returns an array containing all elements of <i>enum</i> for which
  243:  *  <em>block</em> is not <code>false</code> (see also
  244:  *  <code>Enumerable#reject</code>).
  245:  *     
  246:  *     (1..10).find_all {|i|  i % 3 == 0 }   #=> [3, 6, 9]
  247:  *     
  248:  */
  249: 
  250: static VALUE
  251: enum_find_all(VALUE obj)
  252: {
  253:     VALUE ary;
  254:     
  255:     RETURN_ENUMERATOR(obj, 0, 0);
  256: 
  257:     ary = rb_ary_new();
  258:     rb_block_call(obj, id_each, 0, 0, find_all_i, ary);
  259: 
  260:     return ary;
  261: }
  262: 
  263: static VALUE
  264: reject_i(VALUE i, VALUE ary, int argc, VALUE *argv)
  265: {
  266:     if (!RTEST(enum_yield(argc, argv))) {
  267:         rb_ary_push(ary, i);
  268:     }
  269:     return Qnil;
  270: }
  271: 
  272: /*
  273:  *  call-seq:
  274:  *     enum.reject {| obj | block }  => array
  275:  *  
  276:  *  Returns an array for all elements of <i>enum</i> for which
  277:  *  <em>block</em> is false (see also <code>Enumerable#find_all</code>).
  278:  *     
  279:  *     (1..10).reject {|i|  i % 3 == 0 }   #=> [1, 2, 4, 5, 7, 8, 10]
  280:  *     
  281:  */
  282: 
  283: static VALUE
  284: enum_reject(VALUE obj)
  285: {
  286:     VALUE ary;
  287:     
  288:     RETURN_ENUMERATOR(obj, 0, 0);
  289: 
  290:     ary = rb_ary_new();
  291:     rb_block_call(obj, id_each, 0, 0, reject_i, ary);
  292: 
  293:     return ary;
  294: }
  295: 
  296: static VALUE
  297: collect_i(VALUE i, VALUE ary, int argc, VALUE *argv)
  298: {
  299:     rb_ary_push(ary, enum_yield(argc, argv));
  300: 
  301:     return Qnil;
  302: }
  303: 
  304: static VALUE
  305: collect_all(VALUE i, VALUE ary, int argc, VALUE *argv)
  306: {
  307:     rb_ary_push(ary, enum_values_pack(argc, argv));
  308: 
  309:     return Qnil;
  310: }
  311: 
  312: /*
  313:  *  call-seq:
  314:  *     enum.collect {| obj | block }  => array
  315:  *     enum.map     {| obj | block }  => array
  316:  *  
  317:  *  Returns a new array with the results of running <em>block</em> once
  318:  *  for every element in <i>enum</i>.
  319:  *     
  320:  *     (1..4).collect {|i| i*i }   #=> [1, 4, 9, 16]
  321:  *     (1..4).collect { "cat"  }   #=> ["cat", "cat", "cat", "cat"]
  322:  *     
  323:  */
  324: 
  325: static VALUE
  326: enum_collect(VALUE obj)
  327: {
  328:     VALUE ary;
  329: 
  330:     RETURN_ENUMERATOR(obj, 0, 0);
  331: 
  332:     ary = rb_ary_new();
  333:     rb_block_call(obj, id_each, 0, 0, collect_i, ary);
  334: 
  335:     return ary;
  336: }
  337: 
  338: /*
  339:  *  call-seq:
  340:  *     enum.to_a      =>    array
  341:  *     enum.entries   =>    array
  342:  *  
  343:  *  Returns an array containing the items in <i>enum</i>.
  344:  *     
  345:  *     (1..7).to_a                       #=> [1, 2, 3, 4, 5, 6, 7]
  346:  *     { 'a'=>1, 'b'=>2, 'c'=>3 }.to_a   #=> [["a", 1], ["b", 2], ["c", 3]]
  347:  */
  348: static VALUE
  349: enum_to_a(VALUE obj)
  350: {
  351:     VALUE ary = rb_ary_new();
  352: 
  353:     rb_block_call(obj, id_each, 0, 0, collect_all, ary);
  354: 
  355:     return ary;
  356: }
  357: 
  358: static VALUE
  359: inject_i(VALUE i, VALUE p, int argc, VALUE *argv)
  360: {
  361:     VALUE *memo = (VALUE *)p;
  362:     if (memo[0] == Qundef) {
  363:         memo[0] = i;
  364:     }
  365:     else {
  366:         memo[0] = rb_yield_values(2, memo[0], enum_values_pack(argc, argv));
  367:     }
  368:     return Qnil;
  369: }
  370: 
  371: static VALUE
  372: inject_op_i(VALUE i, VALUE p, int argc, VALUE *argv)
  373: {
  374:     VALUE *memo = (VALUE *)p;
  375: 
  376:     if (memo[0] == Qundef) {
  377:         memo[0] = enum_values_pack(argc, argv);
  378:     }
  379:     else {
  380:         memo[0] = rb_funcall(memo[0], (ID)memo[1], 1, i);
  381:     }
  382:     return Qnil;
  383: }
  384: 
  385: /*
  386:  *  call-seq:
  387:  *     enum.inject(initial, sym) => obj
  388:  *     enum.inject(sym)          => obj
  389:  *     enum.inject(initial) {| memo, obj | block }  => obj
  390:  *     enum.inject          {| memo, obj | block }  => obj
  391:  *
  392:  *     enum.reduce(initial, sym) => obj
  393:  *     enum.reduce(sym)          => obj
  394:  *     enum.reduce(initial) {| memo, obj | block }  => obj
  395:  *     enum.reduce          {| memo, obj | block }  => obj
  396:  *  
  397:  *  Combines all elements of <i>enum</i> by applying a binary
  398:  *  operation, specified by a block or a symbol that names a
  399:  *  method or operator.
  400:  *
  401:  *  If you specify a block, then for each element in <i>enum<i>
  402:  *  the block is passed an accumulator value (<i>memo</i>) and the element.
  403:  *  If you specify a symbol instead, then each element in the collection
  404:  *  will be passed to the named method of <i>memo</i>.
  405:  *  In either case, the result becomes the new value for <i>memo</i>. 
  406:  *  At the end of the iteration, the final value of <i>memo</i> is the
  407:  *  return value fo the method.
  408:  *
  409:  *  If you do not explicitly specify an <i>initial</i> value for <i>memo</i>,
  410:  *  then uses the first element of collection is used as the initial value
  411:  *  of <i>memo</i>.
  412:  *
  413:  *  Examples:   
  414:  *
  415:  *     # Sum some numbers
  416:  *     (5..10).reduce(:+)                            #=> 45
  417:  *     # Same using a block and inject
  418:  *     (5..10).inject {|sum, n| sum + n }            #=> 45
  419:  *     # Multiply some numbers
  420:  *     (5..10).reduce(1, :*)                         #=> 151200
  421:  *     # Same using a block
  422:  *     (5..10).inject(1) {|product, n| product * n } #=> 151200
  423:  *     # find the longest word
  424:  *     longest = %w{ cat sheep bear }.inject do |memo,word|
  425:  *        memo.length > word.length ? memo : word
  426:  *     end
  427:  *     longest                                       #=> "sheep"
  428:  *     
  429:  */
  430: static VALUE
  431: enum_inject(int argc, VALUE *argv, VALUE obj)
  432: {
  433:     VALUE memo[2];
  434:     VALUE (*iter)(VALUE, VALUE, int, VALUE*) = inject_i;
  435: 
  436:     switch (rb_scan_args(argc, argv, "02", &memo[0], &memo[1])) {
  437:       case 0:
  438:         memo[0] = Qundef;
  439:         break;
  440:       case 1:
  441:         if (rb_block_given_p()) {
  442:             break;
  443:         }
  444:         memo[1] = (VALUE)rb_to_id(memo[0]);
  445:         memo[0] = Qundef;
  446:         iter = inject_op_i;
  447:         break;
  448:       case 2:
  449:         if (rb_block_given_p()) {
  450:             rb_warning("given block not used");
  451:         }
  452:         memo[1] = (VALUE)rb_to_id(memo[1]);
  453:         iter = inject_op_i;
  454:         break;
  455:     }
  456:     rb_block_call(obj, id_each, 0, 0, iter, (VALUE)memo);
  457:     if (memo[0] == Qundef) return Qnil;
  458:     return memo[0];
  459: }
  460: 
  461: static VALUE
  462: partition_i(VALUE i, VALUE *ary, int argc, VALUE *argv)
  463: {
  464:     if (RTEST(enum_yield(argc, argv))) {
  465:         rb_ary_push(ary[0], i);
  466:     }
  467:     else {
  468:         rb_ary_push(ary[1], i);
  469:     }
  470:     return Qnil;
  471: }
  472: 
  473: /*
  474:  *  call-seq:
  475:  *     enum.partition {| obj | block }  => [ true_array, false_array ]
  476:  *  
  477:  *  Returns two arrays, the first containing the elements of
  478:  *  <i>enum</i> for which the block evaluates to true, the second
  479:  *  containing the rest.
  480:  *     
  481:  *     (1..6).partition {|i| (i&1).zero?}   #=> [[2, 4, 6], [1, 3, 5]]
  482:  *     
  483:  */
  484: 
  485: static VALUE
  486: enum_partition(VALUE obj)
  487: {
  488:     VALUE ary[2];
  489: 
  490:     RETURN_ENUMERATOR(obj, 0, 0);
  491: 
  492:     ary[0] = rb_ary_new();
  493:     ary[1] = rb_ary_new();
  494:     rb_block_call(obj, id_each, 0, 0, partition_i, (VALUE)ary);
  495: 
  496:     return rb_assoc_new(ary[0], ary[1]);
  497: }
  498: 
  499: static VALUE
  500: group_by_i(VALUE i, VALUE hash, int argc, VALUE *argv)
  501: {
  502:     VALUE group = enum_yield(argc, argv);
  503:     VALUE values;
  504: 
  505:     values = rb_hash_aref(hash, group);
  506:     if (NIL_P(values)) {
  507:         values = rb_ary_new3(1, i);
  508:         rb_hash_aset(hash, group, values);
  509:     }
  510:     else {
  511:         rb_ary_push(values, i);
  512:     }
  513:     return Qnil;
  514: }
  515: 
  516: /*
  517:  *  call-seq:
  518:  *     enum.group_by {| obj | block }  => a_hash
  519:  *  
  520:  *  Returns a hash, which keys are evaluated result from the
  521:  *  block, and values are arrays of elements in <i>enum</i>
  522:  *  corresponding to the key.
  523:  *     
  524:  *     (1..6).group_by {|i| i%3}   #=> {0=>[3, 6], 1=>[1, 4], 2=>[2, 5]}
  525:  *     
  526:  */
  527: 
  528: static VALUE
  529: enum_group_by(VALUE obj)
  530: {
  531:     VALUE hash;
  532: 
  533:     RETURN_ENUMERATOR(obj, 0, 0);
  534: 
  535:     hash = rb_hash_new();
  536:     rb_block_call(obj, id_each, 0, 0, group_by_i, hash);
  537: 
  538:     return hash;
  539: }
  540: 
  541: static VALUE
  542: first_i(VALUE i, VALUE *ary)
  543: {
  544:     if (NIL_P(ary[0])) {
  545:         ary[1] = i;
  546:         rb_iter_break();
  547:     }
  548:     else {
  549:         long n = NUM2LONG(ary[0]);
  550: 
  551:         if (n <= 0) {
  552:             rb_iter_break();
  553:         }
  554:         rb_ary_push(ary[1], i);
  555:         n--;
  556:         ary[0] = INT2NUM(n);
  557:     }
  558:     return Qnil;
  559: }
  560: 
  561: /*
  562:  *  call-seq:
  563:  *     enum.first      -> obj or nil
  564:  *     enum.first(n)   -> an_array
  565:  *  
  566:  *  Returns the first element, or the first +n+ elements, of the enumerable.
  567:  *  If the enumerable is empty, the first form returns <code>nil</code>, and the
  568:  *  second form returns an empty array.
  569:  *     
  570:  */
  571: 
  572: static VALUE
  573: enum_first(int argc, VALUE *argv, VALUE obj)
  574: {
  575:     VALUE n, ary[2];
  576:     
  577:     rb_scan_args(argc, argv, "01", &n);
  578: 
  579:     if (argc == 0) {
  580:         ary[0] = ary[1] = Qnil;
  581:     }
  582:     else {
  583:         ary[0] = n;
  584:         ary[1] = rb_ary_new2(NUM2LONG(n));
  585:     }
  586:     rb_block_call(obj, id_each, 0, 0, first_i, (VALUE)ary);
  587: 
  588:     return ary[1];
  589: }
  590: 
  591: 
  592: /*
  593:  *  call-seq:
  594:  *     enum.sort                     => array
  595:  *     enum.sort {| a, b | block }   => array
  596:  *  
  597:  *  Returns an array containing the items in <i>enum</i> sorted,
  598:  *  either according to their own <code><=></code> method, or by using
  599:  *  the results of the supplied block. The block should return -1, 0, or
  600:  *  +1 depending on the comparison between <i>a</i> and <i>b</i>. As of
  601:  *  Ruby 1.8, the method <code>Enumerable#sort_by</code> implements a
  602:  *  built-in Schwartzian Transform, useful when key computation or
  603:  *  comparison is expensive..
  604:  *     
  605:  *     %w(rhea kea flea).sort         #=> ["flea", "kea", "rhea"]
  606:  *     (1..10).sort {|a,b| b <=> a}   #=> [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]
  607:  */
  608: 
  609: static VALUE
  610: enum_sort(VALUE obj)
  611: {
  612:     return rb_ary_sort(enum_to_a(obj));
  613: }
  614: 
  615: static VALUE
  616: sort_by_i(VALUE i, VALUE ary, int argc, VALUE *argv)
  617: {
  618:     VALUE v;
  619:     NODE *memo;
  620: 
  621:     v = enum_yield(argc, argv);
  622:     if (RBASIC(ary)->klass) {
  623:         rb_raise(rb_eRuntimeError, "sort_by reentered");
  624:     }
  625:     memo = rb_node_newnode(NODE_MEMO, v, i, 0);
  626:     rb_ary_push(ary, (VALUE)memo);
  627:     return Qnil;
  628: }
  629: 
  630: static int
  631: sort_by_cmp(const void *ap, const void *bp, void *data)
  632: {
  633:     VALUE a = (*(NODE *const *)ap)->u1.value;
  634:     VALUE b = (*(NODE *const *)bp)->u1.value;
  635:     VALUE ary = (VALUE)data;
  636: 
  637:     if (RBASIC(ary)->klass) {
  638:         rb_raise(rb_eRuntimeError, "sort_by reentered");
  639:     }
  640:     return rb_cmpint(rb_funcall(a, id_cmp, 1, b), a, b);
  641: }
  642: 
  643: /*
  644:  *  call-seq:
  645:  *     enum.sort_by {| obj | block }    => array
  646:  *  
  647:  *  Sorts <i>enum</i> using a set of keys generated by mapping the
  648:  *  values in <i>enum</i> through the given block.
  649:  *     
  650:  *     %w{ apple pear fig }.sort_by {|word| word.length}
  651:                      #=> ["fig", "pear", "apple"]
  652: