ZJIT: Annotate & method for true and false

Author: st0012

test/ruby/test_zjit.rb

@@ -323,6 +323,82 @@ def test(x, y) = x & y
     RUBY
   end
 
+  def test_opt_and_true_true
+    assert_compiles 'true', %q{
+      def test(a, b) = a & b
+      test(true, true)
+      test(true, true)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_true_false
+    assert_compiles 'false', %q{
+      def test(a, b) = a & b
+      test(true, false)
+      test(true, false)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_false_true
+    assert_compiles 'false', %q{
+      def test(a, b) = a & b
+      test(false, true)
+      test(false, true)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_false_false
+    assert_compiles 'false', %q{
+      def test(a, b) = a & b
+      test(false, false)
+      test(false, false)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_true_with_truthy
+    assert_compiles 'true', %q{
+      def test(a, b) = a & b
+      test(true, true)
+      test(true, 1)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_true_with_falsey
+    assert_compiles 'false', %q{
+      def test(a, b) = a & b
+      test(true, true)
+      test(true, nil)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_false_with_truthy
+    assert_compiles 'false', %q{
+      def test(a, b) = a & b
+      test(false, true)
+      test(false, "hi")
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_bool_fallthrough_to_integer
+    assert_compiles '1', %q{
+      def test(a, b) = a & b
+      test(true, false)
+      test(true, true)
+      test(3, 1)
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
+  def test_opt_and_mixed_types_side_exit
+    assert_compiles '[false, true, 1]', %q{
+      def test(a, b) = a & b
+      test(true, true)
+      r1 = test(true, false)
+      r2 = test(true, 1)
+      r3 = test(3, 1)
+      [r1, r2, r3]
+    }, call_threshold: 2, insns: [:opt_and]
+  end
+
   def test_opt_or
     assert_compiles('[11, [3, 2, 1]]', <<~RUBY, insns: [:opt_or])
       def test(x, y) = x | y

zjit/src/cruby_methods.rs

@@ -82,6 +82,8 @@ pub fn init() -> Annotations {
     annotate!(rb_cArray, "size", types::Fixnum, no_gc, leaf, elidable);
     annotate!(rb_cNilClass, "nil?", types::TrueClassExact, no_gc, leaf, elidable);
     annotate!(rb_mKernel, "nil?", types::FalseClassExact, no_gc, leaf, elidable);
+    annotate!(rb_cTrueClass, "&", types::BoolExact, no_gc, leaf, elidable);
+    annotate!(rb_cFalseClass, "&", types::FalseClassExact, no_gc, leaf, elidable);
 
     Annotations {
         cfuncs: std::mem::take(cfuncs)

zjit/src/hir.rs

@@ -6787,4 +6787,72 @@ mod opt_tests {
               Return v8
         "#]]);
     }
+
+    #[test]
+    fn test_guard_true_and_true() {
+        eval("
+            def test(a, b) = a & b
+
+            test(true, true)
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+              PatchPoint MethodRedefined(TrueClass@0x1000, &@0x1008)
+              v8:TrueClassExact = GuardType v1, TrueClassExact
+              v9:BoolExact = CCall &@0x1010, v8, v2
+              Return v9
+        "#]]);
+    }
+
+    #[test]
+    fn test_guard_true_and_false() {
+        eval("
+            def test(a, b) = a & b
+
+            test(true, false)
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+              PatchPoint MethodRedefined(TrueClass@0x1000, &@0x1008)
+              v8:TrueClassExact = GuardType v1, TrueClassExact
+              v9:BoolExact = CCall &@0x1010, v8, v2
+              Return v9
+        "#]]);
+    }
+
+    #[test]
+    fn test_guard_false_and_true() {
+        eval("
+            def test(a, b) = a & b
+
+            test(false, true)
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+              PatchPoint MethodRedefined(FalseClass@0x1000, &@0x1008)
+              v8:FalseClassExact = GuardType v1, FalseClassExact
+              v9:FalseClassExact = CCall &@0x1010, v8, v2
+              Return v9
+        "#]]);
+    }
+
+    #[test]
+    fn test_guard_false_and_false() {
+        eval("
+            def test(a, b) = a & b
+
+            test(false, false)
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+              PatchPoint MethodRedefined(FalseClass@0x1000, &@0x1008)
+              v8:FalseClassExact = GuardType v1, FalseClassExact
+              v9:FalseClassExact = CCall &@0x1010, v8, v2
+              Return v9
+        "#]]);
+    }
 }