From 9e0a56fb24ed741d2215a4c4f7e560fa7b88dc8b Mon Sep 17 00:00:00 2001
From: John Hawthorn <john@hawthorn.email>
Date: Thu, 24 Jun 2021 08:15:53 -0700
Subject: Add FLONUM detection

---
 bootstraptest/test_yjit.rb | 19 +++++++++++++++++++
 yjit_codegen.c             | 18 +++++++++++-------
 yjit_core.h                |  2 ++
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index db494b7911..4bf9bc4262 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1347,3 +1347,22 @@ assert_equal 'bar', %q{
   to_string((-"bar").to_sym)
   to_string((-"bar").to_sym)
 }
+
+# Call to flonum and heap float
+assert_equal '[nil, nil, nil, 1]', %q{
+  def is_inf(obj)
+    obj.infinite?
+  end
+
+  is_inf(0.0)
+  is_inf(0.0)
+  is_inf(1e256)
+  is_inf(1e256)
+
+  [
+    is_inf(0.0),
+    is_inf(1.0),
+    is_inf(1e256),
+    is_inf(1.0/0.0)
+  ]
+}
diff --git a/yjit_codegen.c b/yjit_codegen.c
index 0def2f9402..2470853dd1 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -2251,7 +2251,7 @@ jit_guard_known_klass(jitstate_t *jit, ctx_t *ctx, VALUE known_klass, insn_opnd_
         }
     }
     else if (known_klass == rb_cInteger && FIXNUM_P(sample_instance)) {
-        // We will guard FIXNUM and BIGNUM as though they were separate classes
+        // We will guard fixnum and bignum as though they were separate classes
         // BIGNUM can be handled by the general else case below
         if (val_type.type != ETYPE_FIXNUM || !val_type.is_imm) {
             ADD_COMMENT(cb, "guard object is fixnum");
@@ -2268,15 +2268,19 @@ jit_guard_known_klass(jitstate_t *jit, ctx_t *ctx, VALUE known_klass, insn_opnd_
             STATIC_ASSERT(special_shift_is_8, RUBY_SPECIAL_SHIFT == 8);
             cmp(cb, REG0_8, imm_opnd(RUBY_SYMBOL_FLAG));
             jit_chain_guard(JCC_JNE, jit, ctx, max_chain_depth, side_exit);
-
             ctx_set_opnd_type(ctx, insn_opnd, TYPE_STATIC_SYMBOL);
         }
     }
-    else if (known_klass == rb_cFloat) {
-        // Can't guard for for these classes because they can have both
-        // immediate (special const) instances and instances on the heap. Can
-        // remove this by adding appropriate dynamic checks.
-        return false;
+    else if (known_klass == rb_cFloat && FLONUM_P(sample_instance)) {
+        if (val_type.type != ETYPE_FLONUM || !val_type.is_imm) {
+            // We will guard flonum vs heap float as though they were separate classes
+            ADD_COMMENT(cb, "guard object is flonum");
+            mov(cb, REG1, REG0);
+            and(cb, REG1, imm_opnd(RUBY_FLONUM_MASK));
+            cmp(cb, REG1, imm_opnd(RUBY_FLONUM_FLAG));
+            jit_chain_guard(JCC_JNE, jit, ctx, max_chain_depth, side_exit);
+            ctx_set_opnd_type(ctx, insn_opnd, TYPE_FLONUM);
+        }
     }
     else if (FL_TEST(known_klass, FL_SINGLETON) && sample_instance == rb_attr_get(known_klass, id__attached__)) {
         // Singleton classes are attached to one specific object, so we can
diff --git a/yjit_core.h b/yjit_core.h
index 5939f5b74b..45b7c18557 100644
--- a/yjit_core.h
+++ b/yjit_core.h
@@ -35,6 +35,7 @@ enum yjit_type_enum
     ETYPE_TRUE,
     ETYPE_FALSE,
     ETYPE_FIXNUM,
+    ETYPE_FLONUM,
     ETYPE_ARRAY,
     ETYPE_HASH,
     ETYPE_SYMBOL,
@@ -69,6 +70,7 @@ STATIC_ASSERT(val_type_size, sizeof(val_type_t) == 1);
 #define TYPE_TRUE ( (val_type_t){ .is_imm = 1, .type = ETYPE_TRUE } )
 #define TYPE_FALSE ( (val_type_t){ .is_imm = 1, .type = ETYPE_FALSE } )
 #define TYPE_FIXNUM ( (val_type_t){ .is_imm = 1, .type = ETYPE_FIXNUM } )
+#define TYPE_FLONUM ( (val_type_t){ .is_imm = 1, .type = ETYPE_FLONUM } )
 #define TYPE_STATIC_SYMBOL ( (val_type_t){ .is_imm = 1, .type = ETYPE_SYMBOL } )
 #define TYPE_ARRAY ( (val_type_t){ .is_heap = 1, .type = ETYPE_ARRAY } )
 #define TYPE_HASH ( (val_type_t){ .is_heap = 1, .type = ETYPE_HASH } )
-- 
cgit v1.2.3