Add newrange support to zjit

Author: st0012

zjit/src/cruby_bindings.inc.rs

@@ -334,6 +334,7 @@ pub enum Insn {
     NewArray { elements: Vec<InsnId>, state: InsnId },
     /// NewHash contains a vec of (key, value) pairs
     NewHash { elements: Vec<(InsnId,InsnId)>, state: InsnId },
+    NewRange { low: InsnId, high: InsnId, flag: u32, state: InsnId },
     ArraySet { array: InsnId, idx: usize, val: InsnId },
     ArrayDup { val: InsnId, state: InsnId },
     ArrayMax { elements: Vec<InsnId>, state: InsnId },
@@ -443,6 +444,7 @@ impl Insn {
             Insn::StringCopy { .. } => false,
             Insn::NewArray { .. } => false,
             Insn::NewHash { .. } => false,
+            Insn::NewRange { .. } => false,
             Insn::ArrayDup { .. } => false,
             Insn::HashDup { .. } => false,
             Insn::Test { .. } => false,
@@ -494,6 +496,14 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::NewRange { low, high, flag, .. } => {
+                let flag_str = match *flag {
+                    0 => "..",
+                    1 => "...",
+                    _ => &format!("flag({})", flag),
+                };
+                write!(f, "NewRange {low} {flag_str} {high}")
+            }
             Insn::ArrayMax { elements, .. } => {
                 write!(f, "ArrayMax")?;
                 let mut prefix = " ";
@@ -916,6 +926,7 @@ impl Function {
                 }
                 NewHash { elements: found_elements, state: find!(state) }
             }
+            &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state },
             ArrayMax { elements, state } => ArrayMax { elements: find_vec!(*elements), state: find!(*state) },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
             &SetIvar { self_val, id, val, state } => SetIvar { self_val: find!(self_val), id, val, state },
@@ -976,6 +987,7 @@ impl Function {
             Insn::ArrayDup { .. } => types::ArrayExact,
             Insn::NewHash { .. } => types::HashExact,
             Insn::HashDup { .. } => types::HashExact,
+            Insn::NewRange { .. } => types::RangeExact,
             Insn::CCall { return_type, .. } => *return_type,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
             Insn::GuardBitEquals { val, expected, .. } => self.type_of(*val).intersection(Type::from_value(*expected)),
@@ -1490,6 +1502,11 @@ impl Function {
                     }
                     worklist.push_back(state);
                 }
+                Insn::NewRange { low, high, state, .. } => {
+                    worklist.push_back(low);
+                    worklist.push_back(high);
+                    worklist.push_back(state);
+                }
                 Insn::StringCopy { val }
                 | Insn::StringIntern { val }
                 | Insn::Return { val }
@@ -2316,6 +2333,14 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let val = state.stack_pop()?;
                     fun.push_insn(block, Insn::SetIvar { self_val, id, val, state: exit_id });
                 }
+                YARVINSN_newrange => {
+                    let flag = get_arg(pc, 0).as_u32();
+                    let high = state.stack_pop()?;
+                    let low = state.stack_pop()?;
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let insn_id = fun.push_insn(block, Insn::NewRange { low, high, flag, state: exit_id });
+                    state.stack_push(insn_id);
+                }
                 _ => {
                     // Unknown opcode; side-exit into the interpreter
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
@@ -2709,6 +2734,52 @@ mod tests {
         "#]]);
     }
 
+    #[test]
+    fn test_new_range_with_one_element() {
+        eval("def test(a) = (a..10)");
+        assert_method_hir_with_opcode("test", YARVINSN_newrange, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:Fixnum[10] = Const Value(10)
+              v4:RangeExact = NewRange v0 .. v2
+              Return v4
+        "#]]);
+    }
+
+    #[test]
+    fn test_new_range_with_two_elements() {
+        eval("def test(a, b) = (a..b)");
+        assert_method_hir_with_opcode("test", YARVINSN_newrange, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject):
+              v4:RangeExact = NewRange v0 .. v1
+              Return v4
+        "#]]);
+    }
+
+    #[test]
+    fn test_new_range_exclude_with_one_element() {
+        eval("def test(a) = (a...10)");
+        assert_method_hir_with_opcode("test", YARVINSN_newrange, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:Fixnum[10] = Const Value(10)
+              v4:RangeExact = NewRange v0 ... v2
+              Return v4
+        "#]]);
+    }
+
+    #[test]
+    fn test_new_range_exclude_with_two_elements() {
+        eval("def test(a, b) = (a...b)");
+        assert_method_hir_with_opcode("test", YARVINSN_newrange, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject):
+              v4:RangeExact = NewRange v0 ... v1
+              Return v4
+        "#]]);
+    }
+
     #[test]
     fn test_array_dup() {
         eval("def test = [1, 2, 3]");
@@ -4233,6 +4304,23 @@ mod opt_tests {
         "#]]);
     }
 
+    #[test]
+    fn test_eliminate_new_range() {
+        eval("
+            def test()
+              c = (1..2)
+              5
+            end
+            test; test
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0():
+              v3:Fixnum[5] = Const Value(5)
+              Return v3
+        "#]]);
+    }
+
     #[test]
     fn test_eliminate_new_array_with_elements() {
         eval("

zjit/src/hir.rs

@@ -71,6 +71,7 @@ def base_type name
 base_type "String"
 base_type "Array"
 base_type "Hash"
+base_type "Range"
 
 (integer, integer_exact) = base_type "Integer"
 # CRuby partitions Integer into immediate and non-immediate variants.

zjit/src/hir_type/gen_hir_type.rb

@@ -1,6 +1,6 @@
 #![allow(non_upper_case_globals)]
-use crate::cruby::{Qfalse, Qnil, Qtrue, VALUE, RUBY_T_ARRAY, RUBY_T_STRING, RUBY_T_HASH};
-use crate::cruby::{rb_cInteger, rb_cFloat, rb_cArray, rb_cHash, rb_cString, rb_cSymbol, rb_cObject, rb_cTrueClass, rb_cFalseClass, rb_cNilClass};
+use crate::cruby::{Qfalse, Qnil, Qtrue, VALUE, RUBY_T_ARRAY, RUBY_T_STRING, RUBY_T_HASH, RUBY_T_RANGE};
+use crate::cruby::{rb_cInteger, rb_cFloat, rb_cArray, rb_cHash, rb_cString, rb_cSymbol, rb_cObject, rb_cTrueClass, rb_cFalseClass, rb_cNilClass, rb_cRange};
 use crate::cruby::ClassRelationship;
 use crate::cruby::get_class_name;
 use crate::hir::PtrPrintMap;
@@ -137,6 +137,10 @@ fn is_hash_exact(val: VALUE) -> bool {
     val.class_of() == unsafe { rb_cHash } || (val.class_of() == VALUE(0) && val.builtin_type() == RUBY_T_HASH)
 }
 
+fn is_range_exact(val: VALUE) -> bool {
+    val.class_of() == unsafe { rb_cRange } || (val.class_of() == VALUE(0) && val.builtin_type() == RUBY_T_RANGE)
+}
+
 impl Type {
     /// Create a `Type` from the given integer.
     pub const fn fixnum(val: i64) -> Type {
@@ -183,6 +187,9 @@ impl Type {
         else if is_hash_exact(val) {
             Type { bits: bits::HashExact, spec: Specialization::Object(val) }
         }
+        else if is_range_exact(val) {
+            Type { bits: bits::RangeExact, spec: Specialization::Object(val) }
+        }
         else if is_string_exact(val) {
             Type { bits: bits::StringExact, spec: Specialization::Object(val) }
         }
@@ -277,6 +284,7 @@ impl Type {
         if class == unsafe { rb_cInteger } { return true; }
         if class == unsafe { rb_cNilClass } { return true; }
         if class == unsafe { rb_cObject } { return true; }
+        if class == unsafe { rb_cRange } { return true; }
         if class == unsafe { rb_cString } { return true; }
         if class == unsafe { rb_cSymbol } { return true; }
         if class == unsafe { rb_cTrueClass } { return true; }
@@ -383,6 +391,7 @@ impl Type {
         if self.is_subtype(types::IntegerExact) { return Some(unsafe { rb_cInteger }); }
         if self.is_subtype(types::NilClassExact) { return Some(unsafe { rb_cNilClass }); }
         if self.is_subtype(types::ObjectExact) { return Some(unsafe { rb_cObject }); }
+        if self.is_subtype(types::RangeExact) { return Some(unsafe { rb_cRange }); }
         if self.is_subtype(types::StringExact) { return Some(unsafe { rb_cString }); }
         if self.is_subtype(types::SymbolExact) { return Some(unsafe { rb_cSymbol }); }
         if self.is_subtype(types::TrueClassExact) { return Some(unsafe { rb_cTrueClass }); }