Add newrange support to zjit

Author: st0012

test/ruby/test_zjit.rb

@@ -227,6 +227,27 @@ def test = [1,2,3]
     }
   end
 
+  def test_new_range_inclusive
+    assert_compiles '1..5', %q{
+      def test(a, b) = a..b
+      test(1, 5)
+    }
+  end
+
+  def test_new_range_exclusive
+    assert_compiles '1...5', %q{
+      def test(a, b) = a...b
+      test(1, 5)
+    }
+  end
+
+  def test_new_range_with_literal
+    assert_compiles '3..10', %q{
+      def test(n) = n..10
+      test(3)
+    }
+  end
+
   def test_if
     assert_compiles '[0, nil]', %q{
       def test(n)

zjit/src/codegen.rs

@@ -251,6 +251,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::PutSelf => gen_putself(),
         Insn::Const { val: Const::Value(val) } => gen_const(*val),
         Insn::NewArray { elements, state } => gen_new_array(jit, asm, elements, &function.frame_state(*state)),
+        Insn::NewRange { low, high, flag, state } => gen_new_range(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)),
         Insn::Param { idx } => unreachable!("block.insns should not have Insn::Param({idx})"),
         Insn::Snapshot { .. } => return Some(()), // we don't need to do anything for this instruction at the moment
@@ -550,6 +551,28 @@ fn gen_new_array(
     new_array
 }
 
+/// Compile a new range instruction
+fn gen_new_range(
+    asm: &mut Assembler,
+    low: lir::Opnd,
+    high: lir::Opnd,
+    flag: u32,
+    state: &FrameState,
+) -> lir::Opnd {
+    asm_comment!(asm, "call rb_range_new");
+
+    // Save PC
+    gen_save_pc(asm, state);
+
+    // Call rb_range_new(low, high, flag)
+    let new_range = asm.ccall(
+        rb_range_new as *const u8,
+        vec![low, high, lir::Opnd::Imm(flag as i64)],
+    );
+
+    new_range
+}
+
 /// Compile code that exits from JIT code with a return value
 fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
     // Pop the current frame (ec->cfp++)

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 }); }