From cfb8ed6495c4c2eacf8689f436822b61b5a4c129 Mon Sep 17 00:00:00 2001
From: ParadoxV5 <paradox.ver5@gmail.com>
Date: Wed, 8 Jan 2025 00:05:39 -0700
Subject: [PATCH] Add type for `Enumerator::Chain`

---
 core/enumerator.rbs            | 24 ++++++++++++++++++++++--
 test/stdlib/Enumerator_test.rb | 19 +++++++++++++++++++
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/core/enumerator.rbs b/core/enumerator.rbs
index c87367155..c1f8ec394 100644
--- a/core/enumerator.rbs
+++ b/core/enumerator.rbs
@@ -613,7 +613,22 @@ end
 # This type of objects can be created by Enumerable#chain and Enumerator#+.
 #
 class Enumerator::Chain[out Elem] < Enumerator[Elem, void]
-  include Enumerable[Elem]
+  # * `Rubocop: `self` type is not allowed in this context
+  # * {Enumerator::Chain#each} without block doesn't return `self`, unlike {Enumerator#each}.
+  #include Enumerator::_Each[Enum, self] # Workaround: def each:
+
+  # <!--
+  #   rdoc-file=enumerator.c
+  #   - Enumerator::Chain.new(*enums) -> enum
+  # -->
+  # Generates a new enumerator object that iterates over the elements of given
+  # enumerable objects in sequence.
+  #
+  #     e = Enumerator::Chain.new(1..3, [4, 5])
+  #     e.to_a #=> [1, 2, 3, 4, 5]
+  #     e.size #=> 5
+  #
+  def initialize: (*_Each[Elem] enums) -> void
 
   # <!--
   #   rdoc-file=enumerator.c
@@ -626,5 +641,10 @@ class Enumerator::Chain[out Elem] < Enumerator[Elem, void]
   #
   # If no block is given, returns an enumerator.
   #
-  def each: () { (Elem) -> void } -> void
+  def each: () { (Enum) -> void } -> self
+          | () -> Enumerator[E, self]
+
+  # wrong argument type chain (expected enumerator) (TypeError)
+  def with_index: (?Integer) ?{ (?) -> untyped } -> bot
+  def each_with_index: () ?{ (?) -> untyped } -> bot
 end
diff --git a/test/stdlib/Enumerator_test.rb b/test/stdlib/Enumerator_test.rb
index 744bb4b63..f5c61d5fa 100644
--- a/test/stdlib/Enumerator_test.rb
+++ b/test/stdlib/Enumerator_test.rb
@@ -88,3 +88,22 @@ def test_to_proc
     end.next
   end
 end
+
+class EnumeratorChainTest < Test::Unit::TestCase
+  include TestHelper
+
+  testing "::Enumerator::Chain[::Integer]"
+
+  def test_class_new
+    assert_send_type "(*_Each[Integer] enums) -> Enumerator::Chain[Integer]",
+                     Enumerator::Chain, :new, 1..3, [4, 5]
+  end
+
+  def test_each
+    enum = Enumerator::Chain.new 1..3, [4, 5]
+    assert_send_type "() { (Integer) -> nil } -> Enumerator::Chain[Integer]",
+                     enum, :each do end
+    assert_send_type "() -> Enumerator[Integer, Enumerator::Chain[Integer]]",
+                     enum, :each
+  end
+end