Generator method not working as expected in Ruby native extension - ruby

Question

I work in simulation and often need multiple data generators which are structurally identical but have different parameterizations. I'm trying to write a gem using C extensions to implement a factory method, which determines and creates an appropriate type of enumerator based on the parameterization provided. The following is not my actual code, but is intended as a minimal reproducible example to illustrate the behavior I find confusing:
#include "ruby.h"
VALUE rb_mTest = Qnil;
VALUE rb_cTest = Qnil;
static VALUE super_initialize(VALUE self) {
return self;
}
static VALUE rand_enum(VALUE self) {
RETURN_SIZED_ENUMERATOR(self, 0, 0, 0);
VALUE ary[1] = {LL2NUM(42)};
for(;;) {
rb_yield(rb_funcallv(rb_mKernel, rb_intern("rand"), 1, ary));
}
return Qnil;
}
static VALUE enum_by_2(int32_t argc, VALUE* argv , VALUE self) {
rb_check_arity(argc, 1, 1);
int64_t value = NUM2LL(argv[0]);
RETURN_SIZED_ENUMERATOR(self, 1, argv, (1ll << 63) - value);
for(;;) {
rb_yield(LL2NUM(value));
value += 2;
}
return Qnil;
}
static VALUE enum_factory(int32_t argc, VALUE* argv, VALUE self) {
VALUE argument;
rb_scan_args(argc, argv, "01", &argument);
switch (TYPE(argument)) {
case T_NIL:
printf("It's a nil\n");
return rand_enum(self);
case T_FIXNUM:
printf("It's a FIXNUM\n");
return enum_by_2(1, argv, self);
default:
printf("Unrecognized argument type\n");
return Qnil;
}
}
void Init_test(void) {
rb_mTest = rb_define_module("Test");
rb_cTest = rb_define_class_under(rb_mTest, "Tester", rb_cObject);
rb_define_method(rb_cTest, "initialize", super_initialize, 0);
rb_define_method(rb_cTest, "enum_factory", enum_factory, -1);
}
This code works, but in a way that I didn't expect. Here's an irb session from a test run after compiling the code above:
irb(main):001:0> require_relative 'lib/test'
=> true
irb(main):002:0> t = Test::Tester.new
=> #<Test::Tester:0x0000000104b23c70>
irb(main):003:0> g1 = t.enum_factory
It's a nil
=> #<Enumerator: ...>
irb(main):004:0> g2 = t.enum_factory(10)
It's a FIXNUM
=> #<Enumerator: ...>
irb(main):005:0> g1.take(5).to_a
It's a nil
=> [35, 21, 30, 20, 0]
irb(main):006:0> g2.take(5).to_a
It's a FIXNUM
=> [10, 12, 14, 16, 18]
As you can see, the enum_factory switches to return different enumerators based on the parameterization, and the printf statements and results indicate that the correct parameterization is being applied. My confusion stems from the printf output which shows that subsequent method calls applied directly to g1 and g2 still seem to be going through the factory method. My work often generates samples into the millions or even tens of millions, so my intentions were to maximize speed by writing this in C and to avoid parsing the parameters after doing it once in the factory, but clearly that's not what's happening. I'm probably misreading the (sparse) C API documentation and missing something trivial or obvious. I'd appreciate any pointers to where my misunderstanding lies.

Answer 1

You've only tagged ruby, I know ruby, but my c knowledge is 0, but maybe this will help.
I've added some extra output, because it was messing with my head, when things are actually executed in c vs ruby.
static VALUE enum_by_2(int32_t argc, VALUE* argv , VALUE self) {
rb_check_arity(argc, 1, 1);
int64_t value = NUM2LL(argv[0]);
printf("before RETURN_SIZED_ENUMERATOR\n");
RETURN_SIZED_ENUMERATOR(self, 1, argv, (1ll << 63) - value);
printf("after RETURN_SIZED_ENUMERATOR\n");
for(;;) {
printf("for\n");
rb_yield(LL2NUM(value));
printf("after rb_yield\n");
value += 2;
}
return Qnil;
}
Best I figured, this is how RETURN_SIZED_ENUMERATOR function behaves. It returns a proc from the ruby method that was called, enum_factory. When you chain another method to it, proc is called again:
>> require_relative "test"; t = Test::Tester.new; g2 = t.enum_factory(1)
# It's a FIXNUM
# before RETURN_SIZED_ENUMERATOR
=> #<Enumerator: ...>
# ^
# NOTE: not very helpful
>> g2.inspect
=> "#<Enumerator: #<Test::Tester:0x00007f0ddf6cecd0>:enum_factory(1)>"
# ^
# NOTE: a little more context -----------------------'
So g2 now holds a proc type of a thing that calls enum_factory.
>> g2.take(1)
# It's a FIXNUM
# before RETURN_SIZED_ENUMERATOR
# after RETURN_SIZED_ENUMERATOR
# for
=> [1]
# NOTE: take(2) actually does a second for loop.
# yes, take(3) does it three times. it makes so much sense now.
>> g2.take(2)
# It's a FIXNUM
# before RETURN_SIZED_ENUMERATOR
# after RETURN_SIZED_ENUMERATOR
# for
# after rb_yield
# for
=> [1, 3]
But it also depends on the implementation of the method. For example, next works differently:
>> g2.next
# It's a FIXNUM
# before RETURN_SIZED_ENUMERATOR
# after RETURN_SIZED_ENUMERATOR
# for
=> 1
>> g2.next
# after rb_yield
# for
=> 3
From ruby perspective, this is analogous to each and anything else that returns an enumerator:
>> g3 = [1].each
=> #<Enumerator: ...>
>> g3.inspect
=> "#<Enumerator: [1]:each>"
# ^
# NOTE: looks like enum_factory
That's about as much as I could figure out.
Update
Ok, I figured out a little more. This call ID2SYM(rb_frame_this_func()) is what determines the method/proc that is returned in ruby. So just copy pasting rb_enumeratorize_with_size and fixing up some arguments, does the job:
static VALUE enum_by_2(int32_t argc, VALUE* argv , VALUE self) {
rb_check_arity(argc, 1, 1);
int64_t value = NUM2LL(argv[0]);
/* RETURN_SIZED_ENUMERATOR(self, 1, argv, (1ll << 63) - value); */
if (!rb_block_given_p()){
return rb_enumeratorize_with_size(
(self), ID2SYM(rb_intern("enum_by_2")),
(argc), (argv), ((1ll << 63) - value)
);
}
/* there is also `while(0)` in RETURN_SIZED_ENUMERATOR,
I don't know what that's about, so skipped */
for(;;) {
rb_yield(LL2NUM(value));
value += 2;
}
return Qnil;
}
and define this method for ruby in Init_test:
rb_define_method(rb_cTest, "enum_by_2", enum_by_2, -1);
>> g2 = t.enum_factory(1)
# It's a FIXNUM
=> #<Enumerator: ...>
# NOTE: no FIXNUM the second time
>> g2.take(1)
=> [1]
>> g2.inspect
=> "#<Enumerator: #<Test::Tester:0x00007f6b11ba4158>:enum_by_2(1)>"
# ^
# NOTE: the final `enum_by_2` is returned in ruby ---'
This roughly resembles return enum_for(__callee__) unless block_given?.
...
#define RETURN_SIZED_ENUMERATOR(obj, argc, argv, size_fn) do {
...
https://github.com/ruby/ruby/blob/v3_1_2/include/ruby/internal/intern/enumerator.h#L198