From 78233e83529d7e3aee030cc6760f45104247fe51 Mon Sep 17 00:00:00 2001
From: Ian Candy <ipc103@github.com>
Date: Thu, 7 Sep 2023 23:15:24 -0400
Subject: [PATCH] Add `String#getbyte` YJIT implementation (#8397)

* Add getbyte JIT implementation

Adds an implementation for String#getbyte for YJIT, along with a
bootstrap test. This should be helpful for pure Ruby implementations
and to avoid unneeded allocations.

Co-authored-by: John Hawthorn <jhawthorn@github.com>

* Skip the getbyte test for RJIT for now

---------

Co-authored-by: John Hawthorn <jhawthorn@github.com>
Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
---
 bootstraptest/test_yjit.rb | 18 ++++++++++++++++++
 yjit/src/codegen.rs        | 29 +++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 3c53641f91..82fd1e1376 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1832,6 +1832,24 @@ assert_equal 'true', %q{
   jittable_method
 }
 
+# test getbyte on string class
+assert_equal '[97, :nil, 97, :nil, :raised]', %q{
+  def getbyte(s, i)
+   byte = begin
+    s.getbyte(i)
+   rescue TypeError
+    :raised
+   end
+
+   byte || :nil
+  end
+
+  getbyte("a", 0)
+  getbyte("a", 0)
+
+  [getbyte("a", 0), getbyte("a", 1), getbyte("a", -1), getbyte("a", -2), getbyte("a", "a")]
+} unless defined?(RubyVM::RJIT) && RubyVM::RJIT.enabled? # Not yet working on RJIT
+
 # Test << operator on string subclass
 assert_equal 'abab', %q{
   class MyString < String; end
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index e00c1788c3..69b1c52c9c 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4713,6 +4713,34 @@ fn jit_rb_str_bytesize(
     true
 }
 
+fn jit_rb_str_getbyte(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    _ocb: &mut OutlinedCb,
+    _ci: *const rb_callinfo,
+    _cme: *const rb_callable_method_entry_t,
+    _block: Option<BlockHandler>,
+    _argc: i32,
+    _known_recv_class: *const VALUE,
+) -> bool {
+    asm.comment("String#getbyte");
+    extern "C" {
+        fn rb_str_getbyte(str: VALUE, index: VALUE) -> VALUE;
+    }
+    // Raises when non-integers are passed in
+    jit_prepare_routine_call(jit, asm);
+
+    let index = asm.stack_pop(1);
+    let recv = asm.stack_pop(1);
+    let ret_opnd = asm.ccall(rb_str_getbyte as *const u8, vec![recv, index]);
+
+    // Can either return a FIXNUM or nil
+    let out_opnd = asm.stack_push(Type::UnknownImm);
+    asm.mov(out_opnd, ret_opnd);
+
+    true
+}
+
 // Codegen for rb_str_to_s()
 // When String#to_s is called on a String instance, the method returns self and
 // most of the overhead comes from setting up the method call. We observed that
@@ -8758,6 +8786,7 @@ impl CodegenGlobals {
             self.yjit_reg_method(rb_cString, "to_s", jit_rb_str_to_s);
             self.yjit_reg_method(rb_cString, "to_str", jit_rb_str_to_s);
             self.yjit_reg_method(rb_cString, "bytesize", jit_rb_str_bytesize);
+            self.yjit_reg_method(rb_cString, "getbyte", jit_rb_str_getbyte);
             self.yjit_reg_method(rb_cString, "<<", jit_rb_str_concat);
             self.yjit_reg_method(rb_cString, "+@", jit_rb_str_uplus);