ZJIT: Return nil instead of side-exiting on OOB Array#[] / String#getbyte

Bounds failures on these inlines used to side-exit with no recompile
feedback, so OOB-heavy callers (e.g. ActionDispatch::Journey's route
tokenizer) ran interpreted forever. Both return nil out of bounds, so
emit ArrayArefOrNil / StringGetbyteOrNil which handle it in JIT code.

3.5% faster on railsbench.

Author: nmeans

zjit/src/codegen.rs

@@ -628,6 +628,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::ArrayDup { val, state } => gen_array_dup(jit, asm, function, *val, opnd!(val), &function.frame_state(*state)),
         Insn::AdjustBounds { index, length } => gen_adjust_bounds(asm, opnd!(index), opnd!(length)),
         Insn::ArrayAref { array, index, .. } => gen_array_aref(asm, opnd!(array), opnd!(index)),
+        Insn::ArrayArefOrNil { array, index } => gen_array_aref_or_nil(jit, asm, opnd!(array), opnd!(index)),
         Insn::ArrayAset { array, index, val } => {
             no_output!(gen_array_aset(asm, opnd!(array), opnd!(index), opnd!(val)))
         }
@@ -638,6 +639,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::StringCopy { val, chilled, state } => gen_string_copy(jit, asm, function, *val, opnd!(val), *chilled, &function.frame_state(*state)),
         Insn::StringConcat { strings, state } => gen_string_concat(jit, asm, function, opnds!(strings), &function.frame_state(*state)),
         &Insn::StringGetbyte { string, index } => gen_string_getbyte(asm, opnd!(string), opnd!(index)),
+        &Insn::StringGetbyteOrNil { string, index } => gen_string_getbyte_or_nil(jit, asm, opnd!(string), opnd!(index)),
         Insn::StringSetbyteFixnum { string, index, value } => gen_string_setbyte_fixnum(asm, opnd!(string), opnd!(index), opnd!(value)),
         Insn::StringAppend { recv, other, state } => gen_string_append(jit, asm, function, opnd!(recv), opnd!(other), &function.frame_state(*state)),
         Insn::StringAppendCodepoint { recv, other, state } => gen_string_append_codepoint(jit, asm, function, opnd!(recv), opnd!(other), &function.frame_state(*state)),
@@ -2217,6 +2219,69 @@ fn gen_array_aref(
     asm.load(Opnd::mem(VALUE_BITS, elem_ptr, 0))
 }
 
+/// Compile array access (`array[index]`) where `index` is a raw (unboxed)
+/// C long that may be negative or out of bounds. Out-of-bounds reads yield
+/// nil instead of side-exiting, matching rb_ary_entry semantics.
+fn gen_array_aref_or_nil(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    array: Opnd,
+    index: Opnd,
+) -> lir::Opnd {
+    asm_comment!(asm, "ArrayArefOrNil");
+    let hir_block_id = asm.current_block().hir_block_id;
+    let rpo_idx = asm.current_block().rpo_index;
+
+    let load_block = asm.new_block(hir_block_id, false, rpo_idx);
+    let result_block = asm.new_block(hir_block_id, false, rpo_idx);
+    let result_edge = |v| Target::Block(Box::new(lir::BranchEdge {
+        target: result_block,
+        args: vec![v],
+    }));
+
+    let idx = asm.load_mem(index);
+    let array = asm.load_mem(array);
+    let length = gen_array_length(asm, array);
+    let array_ptr = gen_array_ptr(asm, array);
+
+    // Adjust a negative index by the length, as in gen_adjust_bounds
+    let adjusted = asm.add(idx, length);
+    asm.test(idx, idx);
+    let adjusted = asm.csel_l(adjusted, idx);
+
+    // In bounds iff (u64)adjusted < (u64)length; this single unsigned
+    // comparison covers both negative-out-of-range and past-the-end indices.
+    asm.cmp(adjusted, length);
+    asm.jb(Target::Block(Box::new(lir::BranchEdge {
+        target: load_block,
+        args: vec![array_ptr, adjusted],
+    })));
+
+    // Out of bounds: nil
+    asm.jmp(result_edge(Qnil.into()));
+
+    // In bounds: load the element
+    asm.set_current_block(load_block);
+    let label = jit.get_label(asm, load_block, hir_block_id);
+    asm.write_label(label);
+    let array_ptr = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(array_ptr);
+    let adjusted = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(adjusted);
+    let elem_offset = asm.lshift(adjusted, Opnd::UImm(SIZEOF_VALUE.trailing_zeros() as u64));
+    let elem_ptr = asm.add(array_ptr, elem_offset);
+    let elem = asm.load(Opnd::mem(VALUE_BITS, elem_ptr, 0));
+    asm.jmp(result_edge(elem));
+
+    // Join block
+    asm.set_current_block(result_block);
+    let label = jit.get_label(asm, result_block, hir_block_id);
+    asm.write_label(label);
+    let param = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(param);
+    param
+}
+
 fn gen_array_aset(
     asm: &mut Assembler,
     array: Opnd,
@@ -4068,6 +4133,77 @@ fn gen_string_getbyte(asm: &mut Assembler, string: Opnd, index: Opnd) -> Opnd {
     // TODO(max): Use SIB indexing here once the backend supports it
     let string_ptr = asm.add(string_ptr, index);
     let byte = asm.load(Opnd::mem(8, string_ptr, 0));
+    tag_byte(asm, byte)
+}
+
+/// Compile byte access (`string.getbyte(index)`) where `index` is a raw
+/// (unboxed) C long that may be negative or out of bounds. Out-of-bounds
+/// reads yield nil instead of side-exiting, matching rb_str_getbyte semantics.
+fn gen_string_getbyte_or_nil(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    string: Opnd,
+    index: Opnd,
+) -> lir::Opnd {
+    asm_comment!(asm, "StringGetbyteOrNil");
+    let hir_block_id = asm.current_block().hir_block_id;
+    let rpo_idx = asm.current_block().rpo_index;
+
+    let load_block = asm.new_block(hir_block_id, false, rpo_idx);
+    let result_block = asm.new_block(hir_block_id, false, rpo_idx);
+    let result_edge = |v| Target::Block(Box::new(lir::BranchEdge {
+        target: result_block,
+        args: vec![v],
+    }));
+
+    let idx = asm.load_mem(index);
+    let string_reg = asm.load_mem(string);
+    // struct RString stores `len` at a fixed offset for both embedded and
+    // heap strings, so this is a plain load (unlike RArray).
+    let length = asm.load(Opnd::mem(VALUE_BITS, string_reg, RUBY_OFFSET_RSTRING_LEN));
+    let string_ptr = get_string_ptr(asm, string);
+
+    // Adjust a negative index by the length, as in gen_adjust_bounds
+    let adjusted = asm.add(idx, length);
+    asm.test(idx, idx);
+    let adjusted = asm.csel_l(adjusted, idx);
+
+    // In bounds iff (u64)adjusted < (u64)length; this single unsigned
+    // comparison covers both negative-out-of-range and past-the-end indices.
+    asm.cmp(adjusted, length);
+    asm.jb(Target::Block(Box::new(lir::BranchEdge {
+        target: load_block,
+        args: vec![string_ptr, adjusted],
+    })));
+
+    // Out of bounds: nil
+    asm.jmp(result_edge(Qnil.into()));
+
+    // In bounds: load the byte
+    asm.set_current_block(load_block);
+    let label = jit.get_label(asm, load_block, hir_block_id);
+    asm.write_label(label);
+    let string_ptr = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(string_ptr);
+    let adjusted = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(adjusted);
+    // TODO(max): Use SIB indexing here once the backend supports it
+    let byte_ptr = asm.add(string_ptr, adjusted);
+    let byte = asm.load(Opnd::mem(8, byte_ptr, 0));
+    let byte = tag_byte(asm, byte);
+    asm.jmp(result_edge(byte));
+
+    // Join block
+    asm.set_current_block(result_block);
+    let label = jit.get_label(asm, result_block, hir_block_id);
+    asm.write_label(label);
+    let param = asm.new_block_param(VALUE_BITS);
+    asm.current_block().add_parameter(param);
+    param
+}
+
+/// Zero-extend a byte loaded from a string and tag it as a Fixnum
+fn tag_byte(asm: &mut Assembler, byte: Opnd) -> Opnd {
     // Zero-extend the byte to 64 bits
     let byte = byte.with_num_bits(64);
     let byte = asm.and(byte, 0xFF.into());

zjit/src/cruby_methods.rs

@@ -367,13 +367,9 @@ fn inline_array_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::In
             let recv = fun.coerce_to(block, recv, types::Array, state);
             let index = fun.coerce_to(block, index, types::Fixnum, state);
             let index = fun.push_insn(block, hir::Insn::UnboxFixnum { val: index });
-            let length = fun.push_insn(block, hir::Insn::ArrayLength { array: recv });
-            let index = fun.push_insn(block, hir::Insn::GuardLess { left: index, right: length, reason: Box::new(SideExitReason::GuardLess), state });
-            let index = fun.push_insn(block, hir::Insn::AdjustBounds { index, length });
-            let zero = fun.push_insn(block, hir::Insn::Const { val: hir::Const::CInt64(0) });
-            use crate::hir::SideExitReason;
-            let index = fun.push_insn(block, hir::Insn::GuardGreaterEq { left: index, right: zero, reason: Box::new(SideExitReason::GuardGreaterEq), state });
-            let result = fun.push_insn(block, hir::Insn::ArrayAref { array: recv, index });
+            // Out-of-bounds reads are part of Array#[] semantics (they return nil),
+            // so use a total instruction instead of bounds guards that side-exit.
+            let result = fun.push_insn(block, hir::Insn::ArrayArefOrNil { array: recv, index });
             return Some(result);
         }
     }
@@ -478,17 +474,9 @@ fn inline_string_getbyte(fun: &mut hir::Function, block: hir::BlockId, recv: hir
         // when converting the index to a C integer.
         let index = fun.coerce_to(block, index, types::Fixnum, state);
         let unboxed_index = fun.push_insn(block, hir::Insn::UnboxFixnum { val: index });
-        let len = fun.load_string_length(block, recv);
-        // TODO(max): Find a way to mark these guards as not needed for correctness... as in, once
-        // the data dependency is gone (say, the StringGetbyte is elided), they can also be elided.
-        //
-        // This is unlike most other guards.
-        let unboxed_index = fun.push_insn(block, hir::Insn::GuardLess { left: unboxed_index, right: len, reason: Box::new(SideExitReason::GuardLess), state });
-        let unboxed_index = fun.push_insn(block, hir::Insn::AdjustBounds { index: unboxed_index, length: len });
-        let zero = fun.push_insn(block, hir::Insn::Const { val: hir::Const::CInt64(0) });
-        use crate::hir::SideExitReason;
-        let _ = fun.push_insn(block, hir::Insn::GuardGreaterEq { left: unboxed_index, right: zero, reason: Box::new(SideExitReason::GuardGreaterEq), state });
-        let result = fun.push_insn(block, hir::Insn::StringGetbyte { string: recv, index: unboxed_index });
+        // Out-of-bounds reads are part of String#getbyte semantics (they return nil),
+        // so use a total instruction instead of bounds guards that side-exit.
+        let result = fun.push_insn(block, hir::Insn::StringGetbyteOrNil { string: recv, index: unboxed_index });
         return Some(result);
     }
     None

zjit/src/hir.rs

@@ -980,6 +980,10 @@ pub enum Insn {
     StringConcat { strings: Vec<InsnId>, state: InsnId },
     /// Call rb_str_getbyte with known-Fixnum index
     StringGetbyte { string: InsnId, index: InsnId },
+    /// Load the byte at `index` (a raw, unboxed C `long`, possibly negative or
+    /// out of bounds). Returns the byte as a tagged Fixnum, or `nil` when out
+    /// of bounds. Unlike [`Insn::StringGetbyte`], never side-exits.
+    StringGetbyteOrNil { string: InsnId, index: InsnId },
     StringSetbyteFixnum { string: InsnId, index: InsnId, value: InsnId },
     StringAppend { recv: InsnId, other: InsnId, state: InsnId },
     StringAppendCodepoint { recv: InsnId, other: InsnId, state: InsnId },
@@ -1013,6 +1017,10 @@ pub enum Insn {
     /// Push `val` onto `array`, where `array` is already `Array`.
     ArrayPush { array: InsnId, val: InsnId, state: InsnId },
     ArrayAref { array: InsnId, index: InsnId },
+    /// Array access where `index` is a raw (unboxed) C `long`, possibly negative
+    /// or out of bounds. Returns the element, or `nil` when out of bounds.
+    /// Unlike [`Insn::ArrayAref`], never side-exits.
+    ArrayArefOrNil { array: InsnId, index: InsnId },
     ArrayAset { array: InsnId, index: InsnId, val: InsnId },
     ArrayPop { array: InsnId, state: InsnId },
     /// Return the length of the array as a C `long` ([`types::CInt64`])
@@ -1393,7 +1401,7 @@ macro_rules! for_each_operand_impl {
                 $visit_many!(strings);
                 $visit_one!(*state);
             }
-            Insn::StringGetbyte { string, index } => {
+            Insn::StringGetbyte { string, index } | Insn::StringGetbyteOrNil { string, index } => {
                 $visit_one!(*string);
                 $visit_one!(*index);
             }
@@ -1506,7 +1514,7 @@ macro_rules! for_each_operand_impl {
                 $visit_one!(*val);
                 $visit_one!(*state);
             }
-            Insn::ArrayAref { array, index } => {
+            Insn::ArrayAref { array, index } | Insn::ArrayArefOrNil { array, index } => {
                 $visit_one!(*array);
                 $visit_one!(*index);
             }
@@ -1723,6 +1731,7 @@ impl Insn {
             Insn::StringIntern { .. } => effects::Any,
             Insn::StringConcat { .. } => effects::Any,
             Insn::StringGetbyte { .. } => Effect::read_write(abstract_heaps::Other, abstract_heaps::Empty),
+            Insn::StringGetbyteOrNil { .. } => Effect::read_write(abstract_heaps::Other, abstract_heaps::Empty),
             Insn::StringSetbyteFixnum { .. } => effects::Any,
             Insn::StringAppend { .. } => effects::Any,
             Insn::StringAppendCodepoint { .. } => effects::Any,
@@ -1754,6 +1763,7 @@ impl Insn {
             Insn::ArrayExtend { .. } => effects::Any,
             Insn::ArrayPush { .. } => effects::Any,
             Insn::ArrayAref { ..  } => effects::Any,
+            Insn::ArrayArefOrNil { .. } => effects::Any,
             Insn::ArrayAset { .. } => effects::Any,
             Insn::ArrayPop { ..  } => effects::Any,
             Insn::ArrayLength { .. } => Effect::write(abstract_heaps::Empty),
@@ -2037,6 +2047,9 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::ArrayAref { array, index, .. } => {
                 write!(f, "ArrayAref {array}, {index}")
             }
+            Insn::ArrayArefOrNil { array, index, .. } => {
+                write!(f, "ArrayArefOrNil {array}, {index}")
+            }
             Insn::ArrayAset { array, index, val, ..} => {
                 write!(f, "ArrayAset {array}, {index}, {val}")
             }
@@ -2118,6 +2131,9 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::StringGetbyte { string, index, .. } => {
                 write!(f, "StringGetbyte {string}, {index}")
             }
+            Insn::StringGetbyteOrNil { string, index, .. } => {
+                write!(f, "StringGetbyteOrNil {string}, {index}")
+            }
             Insn::StringSetbyteFixnum { string, index, value, .. } => {
                 write!(f, "StringSetbyteFixnum {string}, {index}, {value}")
             }
@@ -3386,6 +3402,7 @@ impl Function {
             Insn::StringIntern { .. } => types::Symbol,
             Insn::StringConcat { .. } => types::StringExact,
             Insn::StringGetbyte { .. } => types::Fixnum,
+            Insn::StringGetbyteOrNil { .. } => types::Fixnum.union(types::NilClass),
             Insn::StringSetbyteFixnum { .. } => types::Fixnum,
             Insn::StringAppend { .. } => types::StringExact,
             Insn::StringAppendCodepoint { .. } => types::StringExact,
@@ -3394,6 +3411,7 @@ impl Function {
             Insn::NewArray { .. } => types::ArrayExact,
             Insn::ArrayDup { .. } => types::ArrayExact,
             Insn::ArrayAref { .. } => types::BasicObject,
+            Insn::ArrayArefOrNil { .. } => types::BasicObject,
             Insn::ArrayPop { .. } => types::BasicObject,
             Insn::ArrayLength { .. } => types::CInt64,
             Insn::AdjustBounds { .. } => types::CInt64,
@@ -6326,7 +6344,7 @@ impl Function {
                             _ => None,
                         })
                     }
-                    &Insn::ArrayAref { array, index }
+                    &Insn::ArrayAref { array, index } | &Insn::ArrayArefOrNil { array, index }
                         if self.type_of(array).ruby_object_known()
                             && self.type_of(index).is_subtype(types::CInt64) => {
                         let array_obj = self.type_of(array).ruby_object().unwrap();
@@ -7157,7 +7175,7 @@ impl Function {
             | Insn::ArrayLength { array, .. } => {
                 self.assert_subtype(insn_id, array, types::Array)
             }
-            Insn::ArrayAref { array, index } => {
+            Insn::ArrayAref { array, index } | Insn::ArrayArefOrNil { array, index } => {
                 self.assert_subtype(insn_id, array, types::Array)?;
                 self.assert_subtype(insn_id, index, types::CInt64)
             }
@@ -7300,7 +7318,7 @@ impl Function {
                 self.assert_subtype(insn_id, left, types::CInt64)?;
                 self.assert_subtype(insn_id, right, types::CInt64)
             },
-            Insn::StringGetbyte { string, index } => {
+            Insn::StringGetbyte { string, index } | Insn::StringGetbyteOrNil { string, index } => {
                 self.assert_subtype(insn_id, string, types::String)?;
                 self.assert_subtype(insn_id, index, types::CInt64)
             },