initial sketch for polymorphic getivar. hangs on added test, not sure where

Author: XrXr

zjit/src/backend/arm64/mod.rs

@@ -1717,7 +1717,6 @@ mod tests {
 
     use super::*;
     use insta::assert_snapshot;
-    use crate::hir;
 
     static TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
 

zjit/src/backend/tests.rs

@@ -3,7 +3,6 @@ use crate::backend::lir::*;
 use crate::cruby::*;
 use crate::codegen::c_callable;
 use crate::options::rb_zjit_prepare_options;
-use crate::hir;
 
 #[test]
 fn test_add() {

zjit/src/distribution.rs

@@ -79,6 +79,8 @@ pub struct DistributionSummary<T: Copy + PartialEq + Default + std::fmt::Debug,
 const SKEW_THRESHOLD: f64 = 0.75;
 
 impl<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> DistributionSummary<T, N> {
+    const N: usize = N; // TODO: better API
+
     pub fn new(dist: &Distribution<T, N>) -> Self {
         #[cfg(debug_assertions)]
         {
@@ -134,6 +136,10 @@ impl<T: Copy + PartialEq + Default + std::fmt::Debug, const N: usize> Distributi
         assert!(idx < N, "index {idx} out of bounds for buckets[{N}]");
         self.buckets[idx]
     }
+
+    pub fn buckets(&self) -> &[T] {
+        &self.buckets
+    }
 }
 
 #[cfg(test)]

zjit/src/hir.rs

@@ -2638,6 +2638,22 @@ impl Function {
         }
     }
 
+    fn polymorphic_summary<'profile>(&self, profiles: &'profile ProfileOracle, recv: InsnId, insn_idx: usize) -> Option<&'profile TypeDistributionSummary> {
+        let Some(entries) = profiles.types.get(&insn_idx) else {
+            return None;
+        };
+        let recv = self.chase_insn(recv);
+        for (entry_insn, entry_type_summary) in entries {
+            if self.union_find.borrow().find_const(*entry_insn) == recv {
+                if entry_type_summary.is_polymorphic() {
+                    return Some(entry_type_summary);
+                }
+                return None;
+            }
+        }
+        None
+    }
+
     /// Prepare arguments for a direct send, handling keyword argument reordering and default synthesis.
     /// Returns the (state, processed_args, kw_bits) to use for the SendWithoutBlockDirect instruction,
     /// or Err with the fallback reason if direct send isn't possible.
@@ -5929,6 +5945,25 @@ fn invalidates_locals(opcode: u32, operands: *const VALUE) -> bool {
 /// The index of the self parameter in the HIR function
 pub const SELF_PARAM_IDX: usize = 0;
 
+// TODO place this better
+fn new_branch_block(
+    fun: &mut Function,
+    insn_idx: u32,
+    exit_state: &FrameState,
+    locals_count: usize,
+    stack_count: usize,
+) -> (BlockId, InsnId, FrameState, InsnId) {
+    let block = fun.new_block(insn_idx);
+    let self_param = fun.push_insn(block, Insn::Param);
+    let mut state = exit_state.clone();
+    state.locals.clear();
+    state.stack.clear();
+    state.locals.extend((0..locals_count).map(|_| fun.push_insn(block, Insn::Param)));
+    state.stack.extend((0..stack_count).map(|_| fun.push_insn(block, Insn::Param)));
+    let snapshot = fun.push_insn(block, Insn::Snapshot { state: state.clone() });
+    (block, self_param, state, snapshot)
+}
+
 /// Compile ISEQ into High-level IR
 pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
     if !ZJITState::can_compile_iseq(iseq) {
@@ -6456,24 +6491,6 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     state.stack_push(fun.push_insn(block, Insn::Const { val: Const::Value(unsafe { rb_block_param_proxy }) }));
                 }
                 YARVINSN_getblockparam => {
-                    fn new_branch_block(
-                        fun: &mut Function,
-                        insn_idx: u32,
-                        exit_state: &FrameState,
-                        locals_count: usize,
-                        stack_count: usize,
-                    ) -> (BlockId, InsnId, FrameState, InsnId) {
-                        let block = fun.new_block(insn_idx);
-                        let self_param = fun.push_insn(block, Insn::Param);
-                        let mut state = exit_state.clone();
-                        state.locals.clear();
-                        state.stack.clear();
-                        state.locals.extend((0..locals_count).map(|_| fun.push_insn(block, Insn::Param)));
-                        state.stack.extend((0..stack_count).map(|_| fun.push_insn(block, Insn::Param)));
-                        let snapshot = fun.push_insn(block, Insn::Snapshot { state: state.clone() });
-                        (block, self_param, state, snapshot)
-                    }
-
                     fn finish_getblockparam_branch(
                         fun: &mut Function,
                         block: BlockId,
@@ -6834,8 +6851,84 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledYARVInsn(opcode) });
                         break;  // End the block
                     }
-                    let result = fun.push_insn(block, Insn::GetIvar { self_val: self_param, id, ic, state: exit_id });
-                    state.stack_push(result);
+
+                    let mut t_object_varying_shapes = true;
+                    if let Some(summary) = fun.polymorphic_summary(&profiles, self_param, exit_state.insn_idx) {
+                        // Only split in cases that vary in shape for now -- all T_OBJECT
+                        for &profiled_type in summary.buckets() {
+                            if profiled_type.is_empty() { break; }
+                            if !profiled_type.flags().is_t_object() {
+                                t_object_varying_shapes = false;
+                            }
+                        }
+
+                        // Just a GetIvar if not splitting
+                        if !t_object_varying_shapes {
+                            let result = fun.push_insn(block, Insn::GetIvar { self_val: self_param, id, ic, state: exit_id });
+                            state.stack_push(result);
+                        } else {
+                            let self_val = fun.push_insn(block, Insn::GuardType { val: self_param, guard_type: types::HeapBasicObject, state: exit_id });
+                            let shape = fun.load_shape(block, self_val);
+                            let join_block = insn_idx_to_block.get(&insn_idx).copied().unwrap_or_else(|| fun.new_block(insn_idx));
+
+                            // For each profiled shape, make a block that handle it
+                            for &profiled_type in summary.buckets() {
+                                if profiled_type.is_empty() { break; }
+                                let mut ivar_index: u16 = 0;
+                                if ! unsafe { rb_shape_get_iv_index(profiled_type.shape().0, id, &mut ivar_index) } {
+                                    // TODO revisit this, can just return nil.
+                                    // If there is no IVAR index, then the ivar was undefined when we
+                                    // entered the compiler.  That means we can just return nil for this
+                                    // shape + iv name
+                                    continue;
+                                }
+                                let branch_insn_idx = exit_state.insn_idx as u32;
+                                let locals_count = state.locals.len();
+                                let stack_count = state.stack.len();
+                                let (get_field_block, field_block_self, mut get_field_state, _) = new_branch_block(&mut fun, branch_insn_idx, &exit_state, locals_count, stack_count);
+                                let expected_shape = fun.push_insn(block, Insn::Const { val: Const::CShape(profiled_type.shape()) });
+                                let test_id = fun.push_insn(block, Insn::IsBitEqual { left: shape, right: expected_shape });
+                                fun.push_insn(block, Insn::IfTrue {
+                                    val: test_id,
+                                    target: BranchEdge {
+                                        target: get_field_block,
+                                        args: state.as_args(self_param)
+                                    }
+                                });
+
+                                let ivar = if profiled_type.flags().is_embedded() {
+                                    // See ROBJECT_FIELDS() from include/ruby/internal/core/robject.h
+                                    let offset = ROBJECT_OFFSET_AS_ARY as i32 + (SIZEOF_VALUE * ivar_index.to_usize()) as i32;
+                                    fun.push_insn(get_field_block, Insn::LoadField { recv: self_val, id, offset, return_type: types::BasicObject })
+                                } else {
+                                    let as_heap = fun.push_insn(get_field_block, Insn::LoadField { recv: self_val, id: ID!(_as_heap), offset: ROBJECT_OFFSET_AS_HEAP_FIELDS as i32, return_type: types::CPtr });
+
+                                    let offset = SIZEOF_VALUE_I32 * ivar_index as i32;
+                                    fun.push_insn(get_field_block, Insn::LoadField { recv: as_heap, id, offset, return_type: types::BasicObject })
+                                };
+
+                                get_field_state.stack_push(ivar);
+                                fun.push_insn(get_field_block, Insn::Jump(BranchEdge {
+                                    target: join_block,
+                                    args: get_field_state.as_args(field_block_self),
+                                }));
+                            }
+
+                            // The fallback case, just getivar
+                            let ivar = fun.push_insn(block, Insn::GetIvar { self_val: self_param, id, ic, state: exit_id });
+                            state.stack_push(ivar);
+                            fun.push_insn(block, Insn::Jump(BranchEdge {
+                                target: join_block,
+                                args: state.as_args(self_param),
+                            }));
+                            queue.push_back((state, join_block, insn_idx, local_inval));
+                            break;  // End the block
+                        }
+                    } else {
+                        // no profile
+                        let result = fun.push_insn(block, Insn::GetIvar { self_val: self_param, id, ic, state: exit_id });
+                        state.stack_push(result);
+                    }
                 }
                 YARVINSN_setinstancevariable => {
                     let id = ID(get_arg(pc, 0).as_u64());

zjit/src/hir/opt_tests.rs

@@ -3968,6 +3968,37 @@ mod hir_opt_tests {
         ");
     }
 
+    // TODO better placement
+    #[test]
+    fn test_polymorphic_getinstancevariable() {
+        set_call_threshold(3);
+        eval("
+            module Tester
+              def test = @foo
+            end
+
+            class A
+              include Tester
+              def initialize
+                @a = 1
+                @foo = 50
+              end
+            end
+
+            class B
+              include Tester
+              def initialize = (@foo = 100)
+            end
+
+            a = A.new
+            b = B.new
+            a.test
+            b.test
+            a.test
+        ");
+        assert_snapshot!(hir_string_proc("Tester.instance_method(:test)"), @r"");
+    }
+
     #[test]
     fn test_setinstancevariable() {
         eval("