From 11e7ab79deb5935573d15ba4f33fc41e5c2b7522 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Fri, 11 Oct 2024 10:22:44 -0400
Subject: [PATCH] Remove 1 allocation in Enumerable#each_with_index (#11868)

* Remove 1 allocation in Enumerable#each_with_index

Previously, each call to Enumerable#each_with_index allocates 2
objects, one for the counting index, the other an imemo_ifunc passed
to `self.each` as a block.

Use `struct vm_ifunc::data` to hold the counting index directly to
remove 1 allocation.

* [DOC] Brief summary for usages of `struct vm_ifunc`
---
 enum.c           | 14 +++++---------
 internal/imemo.h |  7 ++++++-
 internal/vm.h    |  1 +
 vm_eval.c        | 11 +++++++++++
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/enum.c b/enum.c
index 6aec34d850..d8a7cb73f3 100644
--- a/enum.c
+++ b/enum.c
@@ -2984,13 +2984,12 @@ enum_member(VALUE obj, VALUE val)
 }
 
 static VALUE
-each_with_index_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, memo))
+each_with_index_i(RB_BLOCK_CALL_FUNC_ARGLIST(_, index))
 {
-    struct MEMO *m = MEMO_CAST(memo);
-    VALUE n = imemo_count_value(m);
+    struct vm_ifunc *ifunc = rb_current_ifunc();
+    ifunc->data = (const void *)rb_int_succ(index);
 
-    imemo_count_up(m);
-    return rb_yield_values(2, rb_enum_values_pack(argc, argv), n);
+    return rb_yield_values(2, rb_enum_values_pack(argc, argv), index);
 }
 
 /*
@@ -3024,12 +3023,9 @@ each_with_index_i(RB_BLOCK_CALL_FUNC_ARGLIST(i, memo))
 static VALUE
 enum_each_with_index(int argc, VALUE *argv, VALUE obj)
 {
-    struct MEMO *memo;
-
     RETURN_SIZED_ENUMERATOR(obj, argc, argv, enum_size);
 
-    memo = MEMO_NEW(0, 0, 0);
-    rb_block_call(obj, id_each, argc, argv, each_with_index_i, (VALUE)memo);
+    rb_block_call(obj, id_each, argc, argv, each_with_index_i, INT2FIX(0));
     return obj;
 }
 
diff --git a/internal/imemo.h b/internal/imemo.h
index 7420909a14..cb0a528518 100644
--- a/internal/imemo.h
+++ b/internal/imemo.h
@@ -79,7 +79,12 @@ struct vm_ifunc_argc {
 #endif
 };
 
-/*! IFUNC (Internal FUNCtion) */
+/*! IFUNC (Internal FUNCtion)
+ *
+ * Bookkeeping for converting a C function and some closed-over data into a
+ * block passable to methods. Like Ruby Proc, but not directly accessible at
+ * Ruby level since this is an imemo. See rb_block_call() and friends.
+ */
 struct vm_ifunc {
     VALUE flags;
     VALUE *svar_lep;
diff --git a/internal/vm.h b/internal/vm.h
index 0b6b92c279..c30e26a8b4 100644
--- a/internal/vm.h
+++ b/internal/vm.h
@@ -77,6 +77,7 @@ VALUE rb_lambda_call(VALUE obj, ID mid, int argc, const VALUE *argv,
 void rb_check_stack_overflow(void);
 #define RB_BLOCK_NO_USE_PACKED_ARGS 2
 VALUE rb_block_call2(VALUE obj, ID mid, int argc, const VALUE *argv, rb_block_call_func_t bl_proc, VALUE data2, long flags);
+struct vm_ifunc *rb_current_ifunc(void);
 
 #if USE_YJIT
 /* vm_exec.c */
diff --git a/vm_eval.c b/vm_eval.c
index ff8e2e6780..469b89c03a 100644
--- a/vm_eval.c
+++ b/vm_eval.c
@@ -2702,6 +2702,17 @@ rb_current_realfilepath(void)
     return Qnil;
 }
 
+// Assert that an internal function is running and return
+// the imemo object that represents it.
+struct vm_ifunc *
+rb_current_ifunc(void)
+{
+    // Search VM_FRAME_MAGIC_IFUNC to see ifunc imemos put on the iseq field.
+    VALUE ifunc = (VALUE)GET_EC()->cfp->iseq;
+    RUBY_ASSERT_ALWAYS(imemo_type_p(ifunc, imemo_ifunc));
+    return (struct vm_ifunc *)ifunc;
+}
+
 void
 Init_vm_eval(void)
 {