Enumerator.produce accepts an optional size keyword argument

When not specified, the size is unknown (`nil`).  Previously, the size was always `Float::INFINITY` and not specifiable.

[Feature #21701]

NEWS.md

@@ -33,6 +33,30 @@ Note that each entry is kept to a minimum, see links for details.
 
 Note: We're only listing outstanding class updates.
 
+* Enumerator
+
+    * `Enumerator.produce` now accepts an optional `size` keyword argument
+      to specify the size of the enumerator.  It can be an integer,
+      `Float::INFINITY`, a callable object (such as a lambda), or `nil` to
+      indicate unknown size.  When not specified, the size is unknown (`nil`).
+      Previously, the size was always `Float::INFINITY` and not specifiable.
+
+        ```ruby
+        # Infinite enumerator
+        enum = Enumerator.produce(1, size: Float::INFINITY, &:succ)
+        enum.size  # => Float::INFINITY
+
+        # Finite enumerator with known/computable size
+        abs_dir = File.expand_path("./baz") # => "/foo/bar/baz"
+        traverser = Enumerator.produce(abs_dir, size: -> { abs_dir.count("/") + 1 }) {
+          raise StopIteration if it == "/"
+          File.dirname(it)
+        }
+        traverser.size  # => 4
+        ```
+
+      [[Feature #21701]]
+
 * Kernel
 
     * `Kernel#inspect` now checks for the existence of a `#instance_variables_to_inspect` method,
@@ -454,3 +478,4 @@ A lot of work has gone into making Ractors more stable, performant, and usable.
 [Feature #21550]: https://bugs.ruby-lang.org/issues/21550
 [Feature #21557]: https://bugs.ruby-lang.org/issues/21557
 [Bug #21654]:     https://bugs.ruby-lang.org/issues/21654
+[Feature #21701]: https://bugs.ruby-lang.org/issues/21701

enumerator.c

@@ -221,6 +221,7 @@ struct yielder {
 struct producer {
     VALUE init;
     VALUE proc;
+    VALUE size;
 };
 
 typedef struct MEMO *lazyenum_proc_func(VALUE, struct MEMO *, VALUE, long);
@@ -2876,6 +2877,7 @@ producer_mark_and_move(void *p)
     struct producer *ptr = p;
     rb_gc_mark_and_move(&ptr->init);
     rb_gc_mark_and_move(&ptr->proc);
+    rb_gc_mark_and_move(&ptr->size);
 }
 
 #define producer_free RUBY_TYPED_DEFAULT_FREE
@@ -2919,12 +2921,13 @@ producer_allocate(VALUE klass)
     obj = TypedData_Make_Struct(klass, struct producer, &producer_data_type, ptr);
     ptr->init = Qundef;
     ptr->proc = Qundef;
+    ptr->size = Qnil;
 
     return obj;
 }
 
 static VALUE
-producer_init(VALUE obj, VALUE init, VALUE proc)
+producer_init(VALUE obj, VALUE init, VALUE proc, VALUE size)
 {
     struct producer *ptr;
 
@@ -2936,6 +2939,7 @@ producer_init(VALUE obj, VALUE init, VALUE proc)
 
     RB_OBJ_WRITE(obj, &ptr->init, init);
     RB_OBJ_WRITE(obj, &ptr->proc, proc);
+    RB_OBJ_WRITE(obj, &ptr->size, size);
 
     return obj;
 }
@@ -2986,12 +2990,18 @@ producer_each(VALUE obj)
 static VALUE
 producer_size(VALUE obj, VALUE args, VALUE eobj)
 {
-    return DBL2NUM(HUGE_VAL);
+    struct producer *ptr = producer_ptr(obj);
+    VALUE size = ptr->size;
+
+    if (NIL_P(size)) return Qnil;
+    if (RB_INTEGER_TYPE_P(size) || RB_FLOAT_TYPE_P(size)) return size;
+
+    return rb_funcall(size, id_call, 0);
 }
 
 /*
  * call-seq:
- *    Enumerator.produce(initial = nil) { |prev| block } -> enumerator
+ *    Enumerator.produce(initial = nil, size: nil) { |prev| block } -> enumerator
  *
  * Creates an infinite enumerator from any block, just called over and
  * over.  The result of the previous iteration is passed to the next one.
@@ -3023,19 +3033,43 @@ producer_size(VALUE obj, VALUE args, VALUE eobj)
  *   PATTERN = %r{\d+|[-/+*]}
  *   Enumerator.produce { scanner.scan(PATTERN) }.slice_after { scanner.eos? }.first
  *   # => ["7", "+", "38", "/", "6"]
+ *
+ * The optional +size+ keyword argument specifies the size of the enumerator,
+ * which can be retrieved by Enumerator#size.  It can be an integer,
+ * +Float::INFINITY+, a callable object (such as a lambda), or +nil+ to
+ * indicate unknown size.  When not specified, the size is unknown (+nil+).
+ *
+ *   # Infinite enumerator
+ *   enum = Enumerator.produce(1, size: Float::INFINITY, &:succ)
+ *   enum.size  # => Float::INFINITY
+ *
+ *   # Finite enumerator with known/computable size
+ *   abs_dir = File.expand_path("./baz") # => "/foo/bar/baz"
+ *   traverser = Enumerator.produce(abs_dir, size: -> { abs_dir.count("/") + 1 }) {
+ *     raise StopIteration if it == "/"
+ *     File.dirname(it)
+ *   }
+ *   traverser.size  # => 4
  */
 static VALUE
 enumerator_s_produce(int argc, VALUE *argv, VALUE klass)
 {
-    VALUE init, producer;
+    VALUE init, producer, opts, size;
+    ID keyword_ids[1];
 
     if (!rb_block_given_p()) rb_raise(rb_eArgError, "no block given");
 
-    if (rb_scan_args(argc, argv, "01", &init) == 0) {
+    keyword_ids[0] = rb_intern("size");
+    rb_scan_args_kw(RB_SCAN_ARGS_LAST_HASH_KEYWORDS, argc, argv, "01:", &init, &opts);
+    rb_get_kwargs(opts, keyword_ids, 0, 1, &size);
+
+    size = UNDEF_P(size) ? Qnil : convert_to_feasible_size_value(size);
+
+    if (argc == 0 || (argc == 1 && !NIL_P(opts))) {
         init = Qundef;
     }
 
-    producer = producer_init(producer_allocate(rb_cEnumProducer), init, rb_block_proc());
+    producer = producer_init(producer_allocate(rb_cEnumProducer), init, rb_block_proc(), size);
 
     return rb_enumeratorize_with_size_kw(producer, sym_each, 0, 0, producer_size, RB_NO_KEYWORDS);
 }

test/ruby/test_enumerator.rb

@@ -886,12 +886,13 @@ class TestEnumerator < Test::Unit::TestCase
 
   def test_produce
     assert_raise(ArgumentError) { Enumerator.produce }
+    assert_raise(ArgumentError) { Enumerator.produce(a: 1, b: 1) {} }
 
     # Without initial object
     passed_args = []
     enum = Enumerator.produce { |obj| passed_args << obj; (obj || 0).succ }
     assert_instance_of(Enumerator, enum)
-    assert_equal Float::INFINITY, enum.size
+    assert_nil enum.size
     assert_equal [1, 2, 3], enum.take(3)
     assert_equal [nil, 1, 2], passed_args
 
@@ -899,22 +900,14 @@ class TestEnumerator < Test::Unit::TestCase
     passed_args = []
     enum = Enumerator.produce(1) { |obj| passed_args << obj; obj.succ }
     assert_instance_of(Enumerator, enum)
-    assert_equal Float::INFINITY, enum.size
+    assert_nil enum.size
     assert_equal [1, 2, 3], enum.take(3)
     assert_equal [1, 2], passed_args
 
-    # With initial keyword arguments
-    passed_args = []
-    enum = Enumerator.produce(a: 1, b: 1) { |obj| passed_args << obj; obj.shift if obj.respond_to?(:shift)}
-    assert_instance_of(Enumerator, enum)
-    assert_equal Float::INFINITY, enum.size
-    assert_equal [{b: 1}, [1], :a, nil], enum.take(4)
-    assert_equal [{b: 1}, [1], :a], passed_args
-
     # Raising StopIteration
     words = "The quick brown fox jumps over the lazy dog.".scan(/\w+/)
     enum = Enumerator.produce { words.shift or raise StopIteration }
-    assert_equal Float::INFINITY, enum.size
+    assert_nil enum.size
     assert_instance_of(Enumerator, enum)
     assert_equal %w[The quick brown fox jumps over the lazy dog], enum.to_a
 
@@ -924,7 +917,7 @@ class TestEnumerator < Test::Unit::TestCase
       obj.respond_to?(:first) or raise StopIteration
       obj.first
     }
-    assert_equal Float::INFINITY, enum.size
+    assert_nil enum.size
     assert_instance_of(Enumerator, enum)
     assert_nothing_raised {
       assert_equal [
@@ -935,6 +928,25 @@ class TestEnumerator < Test::Unit::TestCase
         "abc",
       ], enum.to_a
     }
+
+    # With size keyword argument
+    enum = Enumerator.produce(1, size: 10) { |obj| obj.succ }
+    assert_equal 10, enum.size
+    assert_equal [1, 2, 3], enum.take(3)
+
+    enum = Enumerator.produce(1, size: -> { 5 }) { |obj| obj.succ }
+    assert_equal 5, enum.size
+
+    enum = Enumerator.produce(1, size: nil) { |obj| obj.succ }
+    assert_equal nil, enum.size
+
+    enum = Enumerator.produce(1, size: Float::INFINITY) { |obj| obj.succ }
+    assert_equal Float::INFINITY, enum.size
+
+    # Without initial value but with size
+    enum = Enumerator.produce(size: 3) { |obj| (obj || 0).succ }
+    assert_equal 3, enum.size
+    assert_equal [1, 2, 3], enum.take(3)
   end
 
   def test_chain_each_lambda