From 509da028c2249cd386190f2856e91549cc9e6c23 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Sun, 25 Dec 2022 21:46:29 -0800
Subject: [PATCH] Rewrite Kernel#loop in Ruby (#6983)

* Rewrite Kernel#loop in Ruby

* Use enum_for(:loop) { Float::INFINITY }

Co-authored-by: Ufuk Kayserilioglu <ufuk@paralaus.com>

* Limit the scope to rescue StopIteration

Co-authored-by: Ufuk Kayserilioglu <ufuk@paralaus.com>
---
 benchmark/loop_generator.rb    |  2 +-
 kernel.rb                      | 41 +++++++++++++++++++++++
 test/ruby/test_settracefunc.rb |  2 +-
 vm_eval.c                      | 60 ----------------------------------
 4 files changed, 43 insertions(+), 62 deletions(-)

diff --git a/benchmark/loop_generator.rb b/benchmark/loop_generator.rb
index d3375c744c..6a3194b670 100644
--- a/benchmark/loop_generator.rb
+++ b/benchmark/loop_generator.rb
@@ -1,4 +1,4 @@
-max = 600000
+max = 6000000
 
 if defined? Fiber
   gen = (1..max).each
diff --git a/kernel.rb b/kernel.rb
index 9cc58bc1d9..32e7dac42f 100644
--- a/kernel.rb
+++ b/kernel.rb
@@ -150,6 +150,47 @@ module Kernel
 
   module_function
 
+  # call-seq:
+  #    loop { block }
+  #    loop            -> an_enumerator
+  #
+  # Repeatedly executes the block.
+  #
+  # If no block is given, an enumerator is returned instead.
+  #
+  #    loop do
+  #      print "Input: "
+  #      line = gets
+  #      break if !line or line =~ /^q/i
+  #      # ...
+  #    end
+  #
+  # StopIteration raised in the block breaks the loop.  In this case,
+  # loop returns the "result" value stored in the exception.
+  #
+  #    enum = Enumerator.new { |y|
+  #      y << "one"
+  #      y << "two"
+  #      :ok
+  #    }
+  #
+  #    result = loop {
+  #      puts enum.next
+  #    } #=> :ok
+  def loop
+    unless Primitive.block_given_p
+      return enum_for(:loop) { Float::INFINITY }
+    end
+
+    begin
+      while true
+        yield
+      end
+    rescue StopIteration => e
+      e.result
+    end
+  end
+
   #
   #  call-seq:
   #     Float(arg, exception: true)    -> float or nil
diff --git a/test/ruby/test_settracefunc.rb b/test/ruby/test_settracefunc.rb
index 6aa6d5911d..05d774016b 100644
--- a/test/ruby/test_settracefunc.rb
+++ b/test/ruby/test_settracefunc.rb
@@ -1668,7 +1668,7 @@ CODE
       Bug10724.new
     }
 
-    assert_equal([:call, :return], evs)
+    assert_equal([:call, :call, :return, :return], evs)
   end
 
   require 'fiber'
diff --git a/vm_eval.c b/vm_eval.c
index 2e1a9b80a6..f219e7037b 100644
--- a/vm_eval.c
+++ b/vm_eval.c
@@ -1440,64 +1440,6 @@ rb_yield_block(RB_BLOCK_CALL_FUNC_ARGLIST(val, arg))
                                rb_keyword_given_p());
 }
 
-static VALUE
-loop_i(VALUE _)
-{
-    for (;;) {
-        rb_yield_0(0, 0);
-    }
-    return Qnil;
-}
-
-static VALUE
-loop_stop(VALUE dummy, VALUE exc)
-{
-    return rb_attr_get(exc, id_result);
-}
-
-static VALUE
-rb_f_loop_size(VALUE self, VALUE args, VALUE eobj)
-{
-    return DBL2NUM(HUGE_VAL);
-}
-
-/*
- *  call-seq:
- *     loop { block }
- *     loop            -> an_enumerator
- *
- *  Repeatedly executes the block.
- *
- *  If no block is given, an enumerator is returned instead.
- *
- *     loop do
- *       print "Input: "
- *       line = gets
- *       break if !line or line =~ /^q/i
- *       # ...
- *     end
- *
- *  StopIteration raised in the block breaks the loop.  In this case,
- *  loop returns the "result" value stored in the exception.
- *
- *     enum = Enumerator.new { |y|
- *       y << "one"
- *       y << "two"
- *       :ok
- *     }
- *
- *     result = loop {
- *       puts enum.next
- *     } #=> :ok
- */
-
-static VALUE
-rb_f_loop(VALUE self)
-{
-    RETURN_SIZED_ENUMERATOR(self, 0, 0, rb_f_loop_size);
-    return rb_rescue2(loop_i, (VALUE)0, loop_stop, (VALUE)0, rb_eStopIteration, (VALUE)0);
-}
-
 #if VMDEBUG
 static const char *
 vm_frametype_name(const rb_control_frame_t *cfp);
@@ -2580,8 +2522,6 @@ Init_vm_eval(void)
     rb_define_global_function("catch", rb_f_catch, -1);
     rb_define_global_function("throw", rb_f_throw, -1);
 
-    rb_define_global_function("loop", rb_f_loop, 0);
-
     rb_define_method(rb_cBasicObject, "instance_eval", rb_obj_instance_eval_internal, -1);
     rb_define_method(rb_cBasicObject, "instance_exec", rb_obj_instance_exec_internal, -1);
     rb_define_private_method(rb_cBasicObject, "method_missing", rb_method_missing, -1);