RustPython
diff --git a/‎crates/jit/src/instructions.rs‎
Lines changed: 104 additions & 73 deletions b/‎crates/jit/src/instructions.rs‎
Lines changed: 104 additions & 73 deletions
@@ -1,6 +1,8 @@
+// instructions.rs
 // spell-checker: disable
 use super::{JitCompileError, JitSig, JitType};
 use alloc::collections::BTreeSet;
+use alloc::rc::Rc;
 use cranelift::codegen::ir::FuncRef;
 use cranelift::prelude::*;
 use num_traits::cast::ToPrimitive;
@@ -15,26 +17,37 @@ enum CustomTrapCode {
     /// Raised when shifting by a negative number
     NegativeShiftCount = 1,
 }
-
+#[derive(Debug)]
+enum ObjectKind {
+    Opaque,
+    Tuple(Rc<TupleShape>),
+}
+#[derive(Debug)]
+enum ElementShape {
+    Scalar(JitType),
+    Object(Rc<ObjectKind>),
+}
+#[derive(Debug)]
+struct TupleShape(Vec<ElementShape>);
 #[derive(Clone)]
 enum Local {
     Scalar {
         var: Variable,
         ty: JitType,
     },
-    Tuple {
-        elements: Vec<Local>,
-    },
+    Object {
+        var: Variable,
+        kind: Rc<ObjectKind>,
+    }
 }
-
 #[derive(Debug)]
 enum JitValue {
     Int(Value),
     Float(Value),
     Bool(Value),
     None,
     Null,
-    Tuple(Vec<Self>),
+    Object(Value, Rc<ObjectKind>),
     FuncRef(FuncRef),
 }
 
@@ -44,6 +57,7 @@ impl JitValue {
             JitType::Int => Self::Int(val),
             JitType::Float => Self::Float(val),
             JitType::Bool => Self::Bool(val),
+            JitType::Object => Self::Object(val, Rc::new(ObjectKind::Opaque)),
         }
     }
 
@@ -52,14 +66,16 @@ impl JitValue {
             Self::Int(_) => Some(JitType::Int),
             Self::Float(_) => Some(JitType::Float),
             Self::Bool(_) => Some(JitType::Bool),
-            Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None,
+            Self::Object(_, _) => Some(JitType::Object),
+            Self::None | Self::Null | Self::FuncRef(_) => None,
         }
     }
 
     fn into_value(self) -> Option<Value> {
         match self {
             Self::Int(val) | Self::Float(val) | Self::Bool(val) => Some(val),
-            Self::None | Self::Null | Self::Tuple(_) | Self::FuncRef(_) => None,
+            Self::Object(val, _) => Some(val),
+            Self::None | Self::Null | Self::FuncRef(_) => None,
         }
     }
 }
@@ -75,17 +91,17 @@ pub(crate) struct FunctionCompiler<'a, 'b> {
     stack: Vec<JitValue>,
     variables: Box<[Option<Local>]>,
     label_to_block: HashMap<Label, Block>,
-    tuple_pool: Vec<Vec<JitValue>>,
+    alloc_func: FuncRef,
     pub(crate) sig: JitSig,
 }
-
 impl<'a, 'b> FunctionCompiler<'a, 'b> {
     pub(crate) fn new(
         builder: &'a mut FunctionBuilder<'b>,
         num_variables: usize,
         arg_types: &[JitType],
         ret_type: Option<JitType>,
         entry_block: Block,
+        alloc_func: FuncRef,
     ) -> Self {
         let mut compiler = Self {
             builder,
@@ -96,7 +112,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 args: arg_types.to_vec(),
                 ret: ret_type,
             },
-            tuple_pool: Vec::new(),
+            alloc_func,
         };
         let params = compiler.builder.func.dfg.block_params(entry_block).to_vec();
         for (i, (ty, val)) in arg_types.iter().zip(params).enumerate() {
@@ -120,12 +136,10 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         value: JitValue,
     ) -> Result<Local, JitCompileError> {
         match value {
-            JitValue::Tuple(elements) => {
-                let mut locals = Vec::with_capacity(elements.len());
-                for element in elements {
-                    locals.push(Self::local_from_value(builder, element)?);
-                }
-                Ok(Local::Tuple { elements: locals })
+            JitValue::Object(ptr, kind) => {
+                let var = builder.declare_var(types::I64);
+                builder.def_var(var, ptr);
+                Ok(Local::Object { var, kind })
             }
 
             scalar => {
@@ -144,26 +158,16 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 ty.clone(),
                 builder.use_var(*var),
             )),
-            Local::Tuple { elements } => {
-                let mut values = Vec::with_capacity(elements.len());
-                for element in elements {
-                    values.push(Self::local_to_value(builder, element)?);
-                }
-                Ok(JitValue::Tuple(values))
-            }
+            Local::Object { var, kind } => Ok(JitValue::Object(builder.use_var(*var), kind.clone())),
         }
     }
-
     fn store_variable(&mut self, idx: oparg::VarNum, val: JitValue) -> Result<(), JitCompileError> {
         #[expect(clippy::mut_mut, reason = "This seems like a false positive")]
         let builder = &mut self.builder;
-
         let local = Self::local_from_value(builder, val)?;
         self.variables[idx] = Some(local);
-
         Ok(())
     }
-
     fn boolean_val(&mut self, val: JitValue) -> Result<Value, JitCompileError> {
         match val {
             JitValue::Float(val) => {
@@ -178,8 +182,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             }
             JitValue::Bool(val) => Ok(val),
             JitValue::None => Ok(self.builder.ins().iconst(types::I8, 0)),
-            JitValue::Null | JitValue::Tuple(_) | JitValue::FuncRef(_) => {
-                Err(JitCompileError::NotSupported)
+            JitValue::Null | JitValue::FuncRef(_) | JitValue::Object(_, _) => {
+                Err(JitCompileError::NotSupported) 
             }
         }
     }
@@ -279,7 +283,6 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             if label_targets.contains(&label) {
                 // Create or get the block for this label:
                 let target_block = self.get_or_create_block(label);
-
                 // If the current block isn't terminated, add a fallthrough jump
                 if let Some(cur) = self.builder.current_block()
                     && cur != target_block
@@ -360,41 +363,75 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             }
             BorrowedConstant::Tuple { elements } => {
                 let mut vals = Vec::new();
-
                 for el in elements {
                     let b = el.borrow_constant();
                     let res = self.prepare_const(b)?;
                     vals.push(res);
                 }
-
-                JitValue::Tuple(vals)
+                let (ptr, shape) = self.build_heap_tuple(vals)?;
+                JitValue::Object(ptr, Rc::new(ObjectKind::Tuple(shape)))
             }
             BorrowedConstant::None => JitValue::None,
             _ => return Err(JitCompileError::NotSupported),
         };
         Ok(value)
     }
-
-    fn tuple_alloc(&mut self, elements: Vec<JitValue>) -> Value {
-        let idx = self.tuple_pool.len();
-        self.tuple_pool.push(elements);
-
-        let ptr_val = self.builder.ins().iconst(types::I64, idx as i64);
-        ptr_val
+    fn build_heap_tuple(&mut self, elements: Vec<JitValue>) -> Result<(Value, Rc<TupleShape>), JitCompileError> {
+        let len_val = self.builder.ins().iconst(types::I64, elements.len() as i64);
+        let call = self.builder.ins().call(self.alloc_func, &[len_val]);
+        let ptr = self.builder.inst_results(call)[0];
+        let mut shape = Vec::with_capacity(elements.len());
+        for (i, element) in elements.into_iter().enumerate() {
+            let offset = 8 + 8 * i as i32;
+            match element {
+                JitValue::Object(val, kind) => {
+                    self.builder.ins().store(MemFlags::new(), val, ptr, offset);
+                    shape.push(ElementShape::Object(kind));
+                }
+                scalar => {
+                    let ty = scalar.to_jit_type().ok_or(JitCompileError::NotSupported)?;
+                    let val = scalar.into_value().unwrap();
+                    self.builder.ins().store(MemFlags::new(), val, ptr, offset);
+                    shape.push(ElementShape::Scalar(ty));
+                }
+            }
+        }
+        Ok((ptr, Rc::new(TupleShape(shape))))
+    }
+    fn unpack_heap_tuple(&mut self, ptr: Value, shape: &TupleShape) -> Vec<JitValue> {
+        let mut values = Vec::with_capacity(shape.0.len());
+        for (i, element_shape) in shape.0.iter().enumerate() {
+            let offset = 8 + 8 * i as i32;
+            let value = match element_shape {
+                ElementShape::Scalar(ty) => {
+                    let val = self.builder.ins().load(ty.to_cranelift(), MemFlags::new(), ptr, offset);
+                    JitValue::from_type_and_value(ty.clone(), val)
+                }
+                ElementShape::Object(kind) => {
+                    let val = self.builder.ins().load(types::I64, MemFlags::new(), ptr, offset);
+                    JitValue::Object(val, kind.clone())
+                }
+            };
+            values.push(value);
+        }
+        values
     }
-
     fn return_value(&mut self, val: JitValue) -> Result<(), JitCompileError> {
-        let cr_val = match val {
-            JitValue::Tuple(elements) => self.tuple_alloc(elements),
-            _ => val.into_value().ok_or(JitCompileError::NotSupported)?,
+        let (cr_val, ret_ty) = match val {
+            JitValue::Int(v)       => (v, JitType::Int),
+            JitValue::Float(v)     => (v, JitType::Float),
+            JitValue::Bool(v)      => (v, JitType::Bool),
+            JitValue::Object(p, _) => (p, JitType::Object),
+            //JitValue::None        => (self.none_ptr(), JitType::Object),
+            _ => return Err(JitCompileError::NotSupported),
         };
 
         // single abi type
         if self.sig.ret.is_none() {
-            self.sig.ret = Some(JitType::Int);
+            self.sig.ret = Some(ret_ty.clone());
             self.builder.func.signature
                 .returns
-                .push(AbiParam::new(types::I64));
+                .push(AbiParam::new(ret_ty.to_cranelift()));
         }
 
         self.builder.ins().return_(&[cr_val]);
@@ -573,7 +610,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             }
             Instruction::BuildTuple { count } => {
                 let elements = self.pop_multiple(count.get(arg) as usize);
-                self.stack.push(JitValue::Tuple(elements));
+                let (ptr, shape) = self.build_heap_tuple(elements)?;
+                self.stack.push(JitValue::Object(ptr, Rc::new(ObjectKind::Tuple(shape))));
                 Ok(())
             }
             Instruction::Call { argc } => {
@@ -584,7 +622,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                     let arg = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
                     args.push(arg.into_value().unwrap());
                 }
-
+                
                 // Pop self_or_null (should be Null for JIT-compiled recursive calls)
                 let self_or_null = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
                 if !matches!(self_or_null, JitValue::Null) {
@@ -791,17 +829,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
                 let local = self.variables[load_idx]
                     .as_ref()
                     .ok_or(JitCompileError::BadBytecode)?;
-                match local {
-                    Local::Scalar { var, ty } => {
-                        self.stack.push(JitValue::from_type_and_value(
-                            ty.clone(),
-                            self.builder.use_var(*var),
-                        ));
-                    },
-                    Local::Tuple { elements: _} => {
-                        self.stack.push(Self::local_to_value(self.builder, local)?);
-                    }
-                }
+                let value = Self::local_to_value(self.builder, local)?;
+                self.stack.push(value);
                 Ok(())
             }
             Instruction::StoreFastStoreFast { var_nums } => {
@@ -848,23 +877,23 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             }
             Instruction::UnpackSequence { count } => {
                 let val = self.stack.pop().ok_or(JitCompileError::BadBytecode)?;
-
-                let elements = match val {
-                    JitValue::Tuple(elements) => elements,
+                let (ptr, shape) = match val {
+                    JitValue::Object(ptr, kind) => match &*kind {
+                        ObjectKind::Tuple(shape) => (ptr, shape.clone()),
+                        _ => return Err(JitCompileError::NotSupported),
+                    },
                     _ => return Err(JitCompileError::NotSupported),
                 };
-
-                if elements.len() != count.get(arg) as usize {
+                if shape.0.len() != count.get(arg) as usize {
                     return Err(JitCompileError::NotSupported);
                 }
-
+                let elements = self.unpack_heap_tuple(ptr, &shape);
                 self.stack.extend(elements.into_iter().rev());
                 Ok(())
             }
             _ => Err(JitCompileError::NotSupported),
         }
     }
-
     fn compile_sub(&mut self, a: Value, b: Value) -> Value {
         let (out, carry) = self.builder.ins().ssub_overflow(a, b);
         self.builder.ins().trapnz(carry, TrapCode::INTEGER_OVERFLOW);
@@ -1320,6 +1349,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         self.builder.ins().jump(merge_block, &[nan_f.into()]);
 
         self.builder.switch_to_block(continue_neg_inf);
+
         // b is an integer here; convert b_floor to an i64.
         let b_i64 = self.builder.ins().fcvt_to_sint(i64_ty, b_floor);
         let one_i = self.builder.ins().iconst(i64_ty, 1);
@@ -1339,7 +1369,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
             even_block,
             &[inf_f.into()],
         );
-
+        
         self.builder.switch_to_block(odd_block);
         let phi_neg_inf = self.builder.block_params(odd_block)[0];
         self.builder.ins().jump(merge_block, &[phi_neg_inf.into()]);
@@ -1358,7 +1388,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         self.builder
             .ins()
             .brif(cmp_lt, a_neg_block, &[], a_pos_block, &[]);
-
+        
         // ----- Case: a > 0: Compute a^b = exp(b * ln(a)) using double–double arithmetic.
         self.builder.switch_to_block(a_pos_block);
         let ln_a_dd = self.dd_ln(a);
@@ -1396,7 +1426,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         let remainder = self.builder.ins().band(b_i64, one_i);
         let zero_i = self.builder.ins().iconst(i64_ty, 0);
         let is_odd = self.builder.ins().icmp(IntCC::NotEqual, remainder, zero_i);
-
+        
         let odd_block = self.builder.create_block();
         let even_block = self.builder.create_block();
         // Append block parameters for both branches:
@@ -1512,7 +1542,7 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         let is_odd = self.builder.ins().icmp_imm(IntCC::Equal, is_odd, 1);
         let mul_result = self.builder.ins().imul(result_phi, base_phi);
         let new_result = self.builder.ins().select(is_odd, mul_result, result_phi);
-
+        
         // Square the base and divide exponent by 2
         let squared_base = self.builder.ins().imul(base_phi, base_phi);
         let new_exp = self.builder.ins().sshr_imm(exp_phi, 1);
@@ -1531,7 +1561,8 @@ impl<'a, 'b> FunctionCompiler<'a, 'b> {
         self.builder.seal_block(loop_block);
         self.builder.seal_block(continue_block);
         self.builder.seal_block(exit_block);
-
+        
         res
     }
-}
+    
+}