Enumerable::Lazy#select added, exceptions when no block given, cleanup

ruby · gregolsen · Mar 3, 2012 · Mar 3, 2012 · Mar 3, 2012 · Mar 3, 2012
commit e00738d81308de3e4c849c9e1b24aacfaa97ffb6
diff --git a/enumerator.c b/enumerator.c
@@ -102,7 +102,7 @@
  */
 VALUE rb_cEnumerator;
 VALUE rb_cLazy;
-static ID id_rewind, id_each;
+static ID id_rewind, id_each, id_new, id_initialize, id_yield, id_call;
 static VALUE sym_each;
 
 VALUE rb_eStopIteration;
@@ -1192,46 +1192,73 @@ lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv)
     if (rb_block_given_p()) {
         return rb_yield(rb_ary_new3(2, m, val));
     } else {
-        return rb_funcall(m, rb_intern("yield"), 1, val);
+        return rb_funcall(m, id_yield, 1, val);
     }
 }
 
 static VALUE
 lazy_init_block(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    return rb_block_call(m, rb_intern("each"), 0, 0, lazy_init_iterator, val);
+    return rb_block_call(m, id_each, 0, 0, lazy_init_iterator, val);
 }
 
 static VALUE
 lazy_initialize(int argc, VALUE *argv, VALUE obj)
 {
     VALUE generator = generator_allocate(rb_cGenerator);
-    rb_block_call(generator, rb_intern("initialize"), 0, 0, lazy_init_block, argv[0]);
+
+    rb_block_call(generator, id_initialize, 0, 0, lazy_init_block, argv[0]);
     enumerator_init(obj, generator, sym_each, 0, 0);
+
     return obj;
 }
 
 static VALUE
 enumerable_lazy(VALUE obj)
 {
-    return rb_funcall(rb_cLazy, rb_intern("new"), 1, obj);
+    return rb_funcall(rb_cLazy, id_new, 1, obj);
 }
 
 static VALUE
 lazy_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
-    VALUE result;
+    VALUE result = rb_funcall(rb_block_proc(), id_call, 1, rb_ary_entry(val, 1));
 
-    result = rb_funcall(m, rb_intern("call"), 1, rb_ary_entry(val, 1));
-
-    return rb_funcall(rb_ary_entry(val, 0), rb_intern("yield"), 1, result);
+    return rb_funcall(rb_ary_entry(val, 0), id_yield, 1, result);
 }
 
 static VALUE
 lazy_map(VALUE obj)
 {
-    VALUE map_proc = rb_block_proc();
-    return rb_block_call(rb_cLazy, rb_intern("new"), 1, &obj, lazy_map_func, map_proc);
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "tried to call lazy map without a block");
+    }
+
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_map_func, 0);
+}
+
+
+static VALUE
+lazy_select_func(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    VALUE element = rb_ary_entry(val, 1);
+    VALUE result = rb_funcall(rb_block_proc(), id_call, 1, element);
+
+    if (result) {
+        return rb_funcall(rb_ary_entry(val, 0), id_yield, 1, element);
+    } else {
+        return result;
+    }
+}
+
+static VALUE
+lazy_select(VALUE obj)
+{
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "tried to call lazy select without a block");
+    }
+
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_select_func, 0);
 }
 
 static VALUE
@@ -1270,7 +1297,10 @@ Init_Enumerator(void)
     rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0);
     rb_define_method(rb_cLazy, "initialize", lazy_initialize, -1);
     rb_define_method(rb_cLazy, "map", lazy_map, 0);
+    rb_define_method(rb_cLazy, "select", lazy_select, 0);
+
     rb_define_alias(rb_cLazy, "collect", "map");
+    rb_define_alias(rb_cLazy, "find_all", "select");
 
     rb_eStopIteration = rb_define_class("StopIteration", rb_eIndexError);
     rb_define_method(rb_eStopIteration, "result", stop_result, 0);
@@ -1292,6 +1322,10 @@ Init_Enumerator(void)
 
     id_rewind = rb_intern("rewind");
     id_each = rb_intern("each");
+    id_call = rb_intern("call");
+    id_yield = rb_intern("yield");
+    id_new = rb_intern("new");
+    id_initialize = rb_intern("initialize");
     sym_each = ID2SYM(id_each);
 
     rb_provide("enumerator.so");	/* for backward compatibility */