YJIT: Simplify Kernel#send guards and admit more cases (#9956)

Previously, our compile time check rejected dynamic symbols (e.g. what
String#to_sym could return) even though we could handle them just fine.
The runtime guards for the type of method name was also overly
restrictive and didn't accept dynamic symbols.

Fold the type check into the rb_get_symbol_id() and take advantage of
the guard already checking for 0. This also avoids generating the same
call twice in case the same method name is presented as different
types.

4 files changed, 24 insertions, 41 deletions

diff --git a/symbol.c b/symbol.c
index ab917f9f53..0e220b81a8 100644
--- a/symbol.c
+++ b/symbol.c
@@ -1140,10 +1140,12 @@ rb_get_symbol_id(VALUE name)
             return 0;
         }
     }
-    else {
-        RUBY_ASSERT_ALWAYS(RB_TYPE_P(name, T_STRING));
+    else if (RB_TYPE_P(name, T_STRING)) {
         return lookup_str_id(name);
     }
+    else {
+        return 0;
+    }
 }
 
 
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 33cb7b16cf..f1b32226ec 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1560,6 +1560,19 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_send_polymorphic_method_name
+    assert_compiles(<<~'RUBY', result: %i[ok ok], no_send_fallbacks: true)
+      mid = "dynamic_mid_#{rand(100..200)}"
+      mid_dsym = mid.to_sym
+
+      define_method(mid) { :ok }
+
+      define_method(:send_site) { send(_1) }
+
+      [send_site(mid), send_site(mid_dsym)]
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 5658fa5950..92bae47c72 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -7994,13 +7994,9 @@ fn gen_send_general(
 
                         let compile_time_name = jit.peek_at_stack(&asm.ctx, argc as isize);
 
-                        if !compile_time_name.string_p() && !compile_time_name.static_sym_p()  {
-                            gen_counter_incr(asm, Counter::send_send_chain_not_string_or_sym);
-                            return None;
-                        }
-
                         mid = unsafe { rb_get_symbol_id(compile_time_name) };
                         if mid == 0 {
+                            // This also rejects method names that need convserion
                             gen_counter_incr(asm, Counter::send_send_null_mid);
                             return None;
                         }
@@ -8015,40 +8011,15 @@ fn gen_send_general(
 
                         jit.assume_method_lookup_stable(asm, ocb, cme);
 
-                        let (known_class, type_mismatch_counter) = {
-                            if compile_time_name.string_p() {
-                                (
-                                    unsafe { rb_cString },
-                                    Counter::guard_send_send_chain_not_string,
-                                )
-                            } else {
-                                (
-                                    unsafe { rb_cSymbol },
-                                    Counter::guard_send_send_chain_not_sym,
-                                )
-                            }
-                        };
-
-                        let name_opnd = asm.stack_opnd(argc);
-                        jit_guard_known_klass(
-                            jit,
+                        asm_comment!(
                             asm,
-                            ocb,
-                            known_class,
-                            name_opnd,
-                            name_opnd.into(),
-                            compile_time_name,
-                            2, // We have string or symbol, so max depth is 2
-                            type_mismatch_counter
+                            "guard sending method name \'{}\'",
+                            unsafe { cstr_to_rust_string(rb_id2name(mid)) }.unwrap_or_else(|| "<unknown>".to_owned()),
                         );
 
-                        // Need to do this here so we don't have too many live
-                        // values for the register allocator.
-                        let name_opnd = asm.load(name_opnd);
-
+                        let name_opnd = asm.stack_opnd(argc);
                         let symbol_id_opnd = asm.ccall(rb_get_symbol_id as *const u8, vec![name_opnd]);
 
-                        asm_comment!(asm, "chain_guard_send");
                         asm.cmp(symbol_id_opnd, mid.into());
                         jit_chain_guard(
                             JCC_JNE,
@@ -8056,7 +8027,7 @@ fn gen_send_general(
                             asm,
                             ocb,
                             SEND_MAX_DEPTH,
-                            Counter::guard_send_send_chain,
+                            Counter::guard_send_send_name_chain,
                         );
 
                         // We have changed the argc, flags, mid, and cme, so we need to re-enter the match
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 46626060bb..5a76c306d9 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -407,7 +407,6 @@ make_counters! {
     send_send_null_cme,
     send_send_nested,
     send_send_chain_string,
-    send_send_chain_not_string_or_sym,
     send_send_getter,
     send_send_builtin,
     send_iseq_has_rest_and_captured,
@@ -450,9 +449,7 @@ make_counters! {
     guard_send_splatarray_length_not_equal,
     guard_send_splatarray_last_ruby_2_keywords,
     guard_send_splat_not_array,
-    guard_send_send_chain,
-    guard_send_send_chain_not_string,
-    guard_send_send_chain_not_sym,
+    guard_send_send_name_chain,
     guard_send_iseq_has_rest_and_splat_too_few,
     guard_send_is_a_class_mismatch,
     guard_send_instance_of_class_mismatch,