ZJIT: NewRangeFixnum instruction (#14409)

* Failing optimization tests for NewRangeFixnum

* NewRangeFixnum general idea

* Use gen_prepare_call_with_gc on gen_new_range_fixnum; add additional hir tests

* Remove unused NewRange rewrite trigger when neither range is Fixnum literal

* Remove misleading 'profiled' name in range optimization tests

* Adjustments as per review comments

* Include new_range_fixnum tests

* remove non-ASCII character from comments as per PR review

* remove non-ASCII character from comments as per PR review

3 files changed, 233 insertions, 1 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 8e450458dd..e79e80fb44 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -807,6 +807,60 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_new_range_fixnum_both_literals_inclusive
+    assert_compiles '1..2', %q{
+      def test()
+        (1..2)
+      end
+      test; test
+    }, call_threshold: 2
+  end
+
+  def test_new_range_fixnum_both_literals_exclusive
+    assert_compiles '1...2', %q{
+      def test()
+        (1...2)
+      end
+      test; test
+    }, call_threshold: 2
+  end
+
+  def test_new_range_fixnum_low_literal_inclusive
+    assert_compiles '1..3', %q{
+      def test(a)
+        (1..a)
+      end
+      test(2); test(3)
+    }, call_threshold: 2
+  end
+
+  def test_new_range_fixnum_low_literal_exclusive
+    assert_compiles '1...3', %q{
+      def test(a)
+        (1...a)
+      end
+      test(2); test(3)
+    }, call_threshold: 2
+  end
+
+  def test_new_range_fixnum_high_literal_inclusive
+    assert_compiles '3..10', %q{
+      def test(a)
+        (a..10)
+      end
+      test(2); test(3)
+    }, call_threshold: 2
+  end
+
+  def test_new_range_fixnum_high_literal_exclusive
+    assert_compiles '3...10', %q{
+      def test(a)
+        (a...10)
+      end
+      test(2); test(3)
+    }, call_threshold: 2
+  end
+
   def test_if
     assert_compiles '[0, nil]', %q{
       def test(n)
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 17d8b817bd..851f8b4aff 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -342,6 +342,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::NewArray { elements, state } => gen_new_array(asm, opnds!(elements), &function.frame_state(*state)),
         Insn::NewHash { elements, state } => gen_new_hash(jit, asm, elements, &function.frame_state(*state)),
         Insn::NewRange { low, high, flag, state } => gen_new_range(jit, asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
+        Insn::NewRangeFixnum { low, high, flag, state } => gen_new_range_fixnum(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::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)),
         // concatstrings shouldn't have 0 strings
@@ -1136,6 +1137,17 @@ fn gen_new_range(
     asm_ccall!(asm, rb_range_new, low, high, (flag as i64).into())
 }
 
+fn gen_new_range_fixnum(
+    asm: &mut Assembler,
+    low: lir::Opnd,
+    high: lir::Opnd,
+    flag: RangeType,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_prepare_call_with_gc(asm, state);
+    asm_ccall!(asm, rb_range_new, low, high, (flag as i64).into())
+}
+
 /// Compile code that exits from JIT code with a return value
 fn gen_return(asm: &mut Assembler, val: lir::Opnd) {
     // Pop the current frame (ec->cfp++)
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 4a07a1aa74..bd712ade04 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -501,6 +501,7 @@ pub enum Insn {
     /// NewHash contains a vec of (key, value) pairs
     NewHash { elements: Vec<(InsnId,InsnId)>, state: InsnId },
     NewRange { low: InsnId, high: InsnId, flag: RangeType, state: InsnId },
+    NewRangeFixnum { low: InsnId, high: InsnId, flag: RangeType, state: InsnId },
     ArrayDup { val: InsnId, state: InsnId },
     ArrayMax { elements: Vec<InsnId>, state: InsnId },
     /// Extend `left` with the elements from `right`. `left` and `right` must both be `Array`.
@@ -689,6 +690,7 @@ impl Insn {
             // TODO: NewRange is effects free if we can prove the two ends to be Fixnum,
             // but we don't have type information here in `impl Insn`. See rb_range_new().
             Insn::NewRange { .. } => true,
+            Insn::NewRangeFixnum { .. } => false,
             _ => true,
         }
     }
@@ -734,6 +736,9 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::NewRange { low, high, flag, .. } => {
                 write!(f, "NewRange {low} {flag} {high}")
             }
+            Insn::NewRangeFixnum { low, high, flag, .. } => {
+                write!(f, "NewRangeFixnum {low} {flag} {high}")
+            }
             Insn::ArrayMax { elements, .. } => {
                 write!(f, "ArrayMax")?;
                 let mut prefix = " ";
@@ -1313,6 +1318,7 @@ impl Function {
                 NewHash { elements: found_elements, state: find!(state) }
             }
             &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state: find!(state) },
+            &NewRangeFixnum { low, high, flag, state } => NewRangeFixnum { low: find!(low), high: find!(high), flag, state: find!(state) },
             &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
             &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
@@ -1383,6 +1389,7 @@ impl Function {
             Insn::NewHash { .. } => types::HashExact,
             Insn::HashDup { .. } => types::HashExact,
             Insn::NewRange { .. } => types::RangeExact,
+            Insn::NewRangeFixnum { .. } => 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)),
@@ -1940,6 +1947,52 @@ impl Function {
             .unwrap_or(insn_id)
     }
 
+    fn optimize_ranges(&mut self) {
+        for block in self.rpo() {
+            let old_insns = std::mem::take(&mut self.blocks[block.0].insns);
+            assert!(self.blocks[block.0].insns.is_empty());
+
+            for insn_id in old_insns {
+                match self.find(insn_id) {
+                    Insn::NewRange { low, high, flag, state } => {
+
+                        // The NewRange rewrite triggers mostly on literals because that is the
+                        // case we can easily prove Fixnum statically and cheaply guard the other
+                        // side.
+                        let low_is_fix  = self.is_a(low,  types::Fixnum);
+                        let high_is_fix = self.is_a(high, types::Fixnum);
+
+                        if low_is_fix && high_is_fix {
+                            // Both statically fixnum => specialize directly
+                            let repl = self.push_insn(block, Insn::NewRangeFixnum { low, high, flag, state });
+                            self.make_equal_to(insn_id, repl);
+                            self.insn_types[repl.0] = self.infer_type(repl);
+                        } else if low_is_fix {
+                            // Only left is fixnum => guard right
+                            let high_fix = self.coerce_to_fixnum(block, high, state);
+                            let repl = self.push_insn(block, Insn::NewRangeFixnum { low, high: high_fix, flag, state });
+                            self.make_equal_to(insn_id, repl);
+                            self.insn_types[repl.0] = self.infer_type(repl);
+                        } else if high_is_fix {
+                            // Only right is fixnum => guard left
+                            let low_fix = self.coerce_to_fixnum(block, low, state);
+                            let repl = self.push_insn(block, Insn::NewRangeFixnum { low: low_fix, high, flag, state });
+                            self.make_equal_to(insn_id, repl);
+                            self.insn_types[repl.0] = self.infer_type(repl);
+                        } else {
+                            // Keep generic op
+                            self.push_insn_id(block, insn_id);
+                        }
+                    }
+                    _ => {
+                        self.push_insn_id(block, insn_id);
+                    }
+                }
+            }
+        }
+        self.infer_types();
+    }
+
     /// Use type information left by `infer_types` to fold away operations that can be evaluated at compile-time.
     ///
     /// It can fold fixnum math, truthiness tests, and branches with constant conditionals.
@@ -2077,7 +2130,8 @@ impl Function {
                 }
                 worklist.push_back(state);
             }
-            &Insn::NewRange { low, high, state, .. } => {
+            &Insn::NewRange { low, high, state, .. }
+            | &Insn::NewRangeFixnum { low, high, state, .. } => {
                 worklist.push_back(low);
                 worklist.push_back(high);
                 worklist.push_back(state);
@@ -2331,6 +2385,8 @@ impl Function {
         #[cfg(debug_assertions)] self.assert_validates();
         self.optimize_c_calls();
         #[cfg(debug_assertions)] self.assert_validates();
+        self.optimize_ranges();
+        #[cfg(debug_assertions)] self.assert_validates();
         self.fold_constants();
         #[cfg(debug_assertions)] self.assert_validates();
         self.clean_cfg();
@@ -6549,6 +6605,116 @@ mod opt_tests {
     }
 
     #[test]
+    fn test_optimize_range_fixnum_inclusive_literals() {
+        eval("
+            def test()
+              (1..2)
+            end
+            test; test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject):
+          v2:RangeExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          CheckInterrupts
+          Return v2
+        ");
+    }
+
+    #[test]
+    fn test_optimize_range_fixnum_exclusive_literals() {
+        eval("
+            def test()
+              (1...2)
+            end
+            test; test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject):
+          v2:RangeExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          CheckInterrupts
+          Return v2
+        ");
+    }
+
+    #[test]
+    fn test_optimize_range_fixnum_inclusive_high_guarded() {
+        eval("
+            def test(a)
+              (1..a)
+            end
+            test(2); test(3)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v3:Fixnum[1] = Const Value(1)
+          v9:Fixnum = GuardType v1, Fixnum
+          v10:RangeExact = NewRangeFixnum v3 NewRangeInclusive v9
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_optimize_range_fixnum_exclusive_high_guarded() {
+        eval("
+            def test(a)
+              (1...a)
+            end
+            test(2); test(3)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v3:Fixnum[1] = Const Value(1)
+          v9:Fixnum = GuardType v1, Fixnum
+          v10:RangeExact = NewRangeFixnum v3 NewRangeExclusive v9
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_optimize_range_fixnum_inclusive_low_guarded() {
+        eval("
+            def test(a)
+              (a..10)
+            end
+            test(2); test(3)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v3:Fixnum[10] = Const Value(10)
+          v9:Fixnum = GuardType v1, Fixnum
+          v10:RangeExact = NewRangeFixnum v9 NewRangeInclusive v3
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
+    fn test_optimize_range_fixnum_exclusive_low_guarded() {
+        eval("
+            def test(a)
+              (a...10)
+            end
+            test(2); test(3)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v3:Fixnum[10] = Const Value(10)
+          v9:Fixnum = GuardType v1, Fixnum
+          v10:RangeExact = NewRangeFixnum v9 NewRangeExclusive v3
+          CheckInterrupts
+          Return v10
+        ");
+    }
+
+    #[test]
     fn test_eliminate_new_array() {
         eval("
             def test()