Skip to main content

rustpython_vm/vm/
vm_ops.rs

1use super::VirtualMachine;
2use crate::stdlib::_warnings;
3use crate::{
4    Py, PyRef,
5    builtins::{PyInt, PyStr, PyStrInterned, PyStrRef, PyType, PyUtf8Str},
6    object::{AsObject, PyObject, PyObjectRef, PyResult},
7    protocol::{PyNumberBinaryOp, PyNumberTernaryOp},
8    types::PyComparisonOp,
9};
10use num_traits::ToPrimitive;
11
12/// [CPython `method_is_overloaded`](https://github.com/python/cpython/blob/v3.14.3/Objects/typeobject.c#L9849-L9879)
13fn method_is_overloaded(
14    class_a: &Py<PyType>,
15    class_b: &Py<PyType>,
16    rop_name: Option<&'static PyStrInterned>,
17    vm: &VirtualMachine,
18) -> PyResult<bool> {
19    let Some(rop_name) = rop_name else {
20        return Ok(false);
21    };
22    let Some(method_b) = class_b.get_attr(rop_name) else {
23        return Ok(false);
24    };
25    class_a.get_attr(rop_name).map_or(Ok(true), |method_a| {
26        vm.identical_or_equal(&method_a, &method_b).map(|eq| !eq)
27    })
28}
29
30macro_rules! binary_func {
31    ($fn:ident, $op_slot:ident, $op:expr) => {
32        pub fn $fn(&self, a: &PyObject, b: &PyObject) -> PyResult {
33            self.binary_op(a, b, PyNumberBinaryOp::$op_slot, $op)
34        }
35    };
36}
37
38macro_rules! ternary_func {
39    ($fn:ident, $op_slot:ident, $op:expr) => {
40        pub fn $fn(&self, a: &PyObject, b: &PyObject, c: &PyObject) -> PyResult {
41            self.ternary_op(a, b, c, PyNumberTernaryOp::$op_slot, $op)
42        }
43    };
44}
45
46macro_rules! inplace_binary_func {
47    ($fn:ident, $iop_slot:ident, $op_slot:ident, $op:expr) => {
48        pub fn $fn(&self, a: &PyObject, b: &PyObject) -> PyResult {
49            self.binary_iop(
50                a,
51                b,
52                PyNumberBinaryOp::$iop_slot,
53                PyNumberBinaryOp::$op_slot,
54                $op,
55            )
56        }
57    };
58}
59
60macro_rules! inplace_ternary_func {
61    ($fn:ident, $iop_slot:ident, $op_slot:ident, $op:expr) => {
62        pub fn $fn(&self, a: &PyObject, b: &PyObject, c: &PyObject) -> PyResult {
63            self.ternary_iop(
64                a,
65                b,
66                c,
67                PyNumberTernaryOp::$iop_slot,
68                PyNumberTernaryOp::$op_slot,
69                $op,
70            )
71        }
72    };
73}
74
75/// Collection of operators
76impl VirtualMachine {
77    #[inline]
78    pub fn bool_eq(&self, a: &PyObject, b: &PyObject) -> PyResult<bool> {
79        a.rich_compare_bool(b, PyComparisonOp::Eq, self)
80    }
81
82    pub fn identical_or_equal(&self, a: &PyObject, b: &PyObject) -> PyResult<bool> {
83        if a.is(b) {
84            Ok(true)
85        } else {
86            self.bool_eq(a, b)
87        }
88    }
89
90    pub fn bool_seq_lt(&self, a: &PyObject, b: &PyObject) -> PyResult<Option<bool>> {
91        let value = if a.rich_compare_bool(b, PyComparisonOp::Lt, self)? {
92            Some(true)
93        } else if !self.bool_eq(a, b)? {
94            Some(false)
95        } else {
96            None
97        };
98        Ok(value)
99    }
100
101    pub fn bool_seq_gt(&self, a: &PyObject, b: &PyObject) -> PyResult<Option<bool>> {
102        let value = if a.rich_compare_bool(b, PyComparisonOp::Gt, self)? {
103            Some(true)
104        } else if !self.bool_eq(a, b)? {
105            Some(false)
106        } else {
107            None
108        };
109        Ok(value)
110    }
111
112    pub fn length_hint_opt(&self, iter: PyObjectRef) -> PyResult<Option<usize>> {
113        match iter.length(self) {
114            Ok(len) => return Ok(Some(len)),
115            Err(e) => {
116                if !e.fast_isinstance(self.ctx.exceptions.type_error) {
117                    return Err(e);
118                }
119            }
120        }
121        let hint = match self.get_method(iter, identifier!(self, __length_hint__)) {
122            Some(hint) => hint?,
123            None => return Ok(None),
124        };
125        let result = match hint.call((), self) {
126            Ok(res) => {
127                if res.is(&self.ctx.not_implemented) {
128                    return Ok(None);
129                }
130                res
131            }
132            Err(e) => {
133                return if e.fast_isinstance(self.ctx.exceptions.type_error) {
134                    Ok(None)
135                } else {
136                    Err(e)
137                };
138            }
139        };
140        let hint = result
141            .downcast_ref::<PyInt>()
142            .ok_or_else(|| {
143                self.new_type_error(format!(
144                    "'{}' object cannot be interpreted as an integer",
145                    result.class().name()
146                ))
147            })?
148            .try_to_primitive::<isize>(self)?;
149        if hint.is_negative() {
150            Err(self.new_value_error("__length_hint__() should return >= 0"))
151        } else {
152            Ok(Some(hint as usize))
153        }
154    }
155
156    /// Checks that the multiplication is able to be performed. On Ok returns the
157    /// index as a usize for sequences to be able to use immediately.
158    pub fn check_repeat_or_overflow_error(&self, length: usize, n: isize) -> PyResult<usize> {
159        if n <= 0 {
160            Ok(0)
161        } else {
162            let n = n as usize;
163            if length > crate::stdlib::sys::MAXSIZE as usize / n {
164                Err(self.new_overflow_error("repeated value are too long"))
165            } else {
166                Ok(n)
167            }
168        }
169    }
170
171    /// Calling scheme used for binary operations:
172    ///
173    /// Order operations are tried until either a valid result or error:
174    ///   `b.rop(b,a)[*], a.op(a,b), b.rop(b,a)`
175    ///
176    /// `[*]` - only when Py_TYPE(a) != Py_TYPE(b) && Py_TYPE(b) is a subclass of Py_TYPE(a)
177    pub fn binary_op1(&self, a: &PyObject, b: &PyObject, op_slot: PyNumberBinaryOp) -> PyResult {
178        let class_a = a.class();
179        let class_b = b.class();
180
181        // Number slots are inherited, direct access is O(1)
182        let slot_a = class_a.slots.as_number.left_binary_op(op_slot);
183        let slot_a_addr = slot_a.map(|x| x as usize);
184        let mut slot_b = None;
185        let left_b_addr;
186
187        if !class_a.is(class_b) {
188            let slot_bb = class_b.slots.as_number.right_binary_op(op_slot);
189            if slot_bb.map(|x| x as usize) != slot_a_addr {
190                slot_b = slot_bb;
191            }
192            left_b_addr = class_b
193                .slots
194                .as_number
195                .left_binary_op(op_slot)
196                .map(|x| x as usize);
197        } else {
198            left_b_addr = slot_a_addr;
199        }
200
201        if let Some(slot_a) = slot_a {
202            if let Some(slot_bb) = slot_b
203                && class_b.fast_issubclass(class_a)
204                && (slot_a_addr != left_b_addr
205                    || method_is_overloaded(
206                        class_a,
207                        class_b,
208                        op_slot.right_method_name(self),
209                        self,
210                    )?)
211            {
212                let ret = slot_bb(a, b, self)?;
213                if !ret.is(&self.ctx.not_implemented) {
214                    return Ok(ret);
215                }
216                slot_b = None;
217            }
218            let ret = slot_a(a, b, self)?;
219            if !ret.is(&self.ctx.not_implemented) {
220                return Ok(ret);
221            }
222        }
223
224        if let Some(slot_b) = slot_b {
225            let ret = slot_b(a, b, self)?;
226            if !ret.is(&self.ctx.not_implemented) {
227                return Ok(ret);
228            }
229        }
230
231        Ok(self.ctx.not_implemented())
232    }
233
234    pub fn binary_op(
235        &self,
236        a: &PyObject,
237        b: &PyObject,
238        op_slot: PyNumberBinaryOp,
239        op: &str,
240    ) -> PyResult {
241        let result = self.binary_op1(a, b, op_slot)?;
242        if !result.is(&self.ctx.not_implemented) {
243            return Ok(result);
244        }
245        Err(self.new_unsupported_bin_op_error(a, b, op))
246    }
247
248    /// Binary in-place operators
249    ///
250    /// The in-place operators are defined to fall back to the 'normal',
251    /// non in-place operations, if the in-place methods are not in place.
252    ///
253    /// - If the left hand object has the appropriate struct members, and
254    ///   they are filled, call the appropriate function and return the
255    ///   result.  No coercion is done on the arguments; the left-hand object
256    ///   is the one the operation is performed on, and it's up to the
257    ///   function to deal with the right-hand object.
258    ///
259    /// - Otherwise, in-place modification is not supported. Handle it exactly as
260    ///   a non in-place operation of the same kind.
261    fn binary_iop1(
262        &self,
263        a: &PyObject,
264        b: &PyObject,
265        iop_slot: PyNumberBinaryOp,
266        op_slot: PyNumberBinaryOp,
267    ) -> PyResult {
268        if let Some(slot) = a.class().slots.as_number.left_binary_op(iop_slot) {
269            let x = slot(a, b, self)?;
270            if !x.is(&self.ctx.not_implemented) {
271                return Ok(x);
272            }
273        }
274        self.binary_op1(a, b, op_slot)
275    }
276
277    fn binary_iop(
278        &self,
279        a: &PyObject,
280        b: &PyObject,
281        iop_slot: PyNumberBinaryOp,
282        op_slot: PyNumberBinaryOp,
283        op: &str,
284    ) -> PyResult {
285        let result = self.binary_iop1(a, b, iop_slot, op_slot)?;
286        if !result.is(&self.ctx.not_implemented) {
287            return Ok(result);
288        }
289        Err(self.new_unsupported_bin_op_error(a, b, op))
290    }
291
292    fn ternary_op(
293        &self,
294        a: &PyObject,
295        b: &PyObject,
296        c: &PyObject,
297        op_slot: PyNumberTernaryOp,
298        op_str: &str,
299    ) -> PyResult {
300        let class_a = a.class();
301        let class_b = b.class();
302        let class_c = c.class();
303
304        // Number slots are inherited, direct access is O(1)
305        let slot_a = class_a.slots.as_number.left_ternary_op(op_slot);
306        let slot_a_addr = slot_a.map(|x| x as usize);
307        let mut slot_b = None;
308        let left_b_addr;
309
310        if !class_a.is(class_b) {
311            let slot_bb = class_b.slots.as_number.right_ternary_op(op_slot);
312            if slot_bb.map(|x| x as usize) != slot_a_addr {
313                slot_b = slot_bb;
314            }
315            left_b_addr = class_b
316                .slots
317                .as_number
318                .left_ternary_op(op_slot)
319                .map(|x| x as usize);
320        } else {
321            left_b_addr = slot_a_addr;
322        }
323
324        if let Some(slot_a) = slot_a {
325            if let Some(slot_bb) = slot_b
326                && class_b.fast_issubclass(class_a)
327                && (slot_a_addr != left_b_addr
328                    || method_is_overloaded(
329                        class_a,
330                        class_b,
331                        op_slot.right_method_name(self),
332                        self,
333                    )?)
334            {
335                let ret = slot_bb(a, b, c, self)?;
336                if !ret.is(&self.ctx.not_implemented) {
337                    return Ok(ret);
338                }
339                slot_b = None;
340            }
341            let ret = slot_a(a, b, c, self)?;
342            if !ret.is(&self.ctx.not_implemented) {
343                return Ok(ret);
344            }
345        }
346
347        if let Some(slot_b) = slot_b {
348            let ret = slot_b(a, b, c, self)?;
349            if !ret.is(&self.ctx.not_implemented) {
350                return Ok(ret);
351            }
352        }
353
354        if let Some(slot_c) = class_c.slots.as_number.left_ternary_op(op_slot)
355            && slot_a.is_some_and(|slot_a| !core::ptr::fn_addr_eq(slot_a, slot_c))
356            && slot_b.is_some_and(|slot_b| !core::ptr::fn_addr_eq(slot_b, slot_c))
357        {
358            let ret = slot_c(a, b, c, self)?;
359            if !ret.is(&self.ctx.not_implemented) {
360                return Ok(ret);
361            }
362        }
363
364        Err(if self.is_none(c) {
365            self.new_type_error(format!(
366                "unsupported operand type(s) for {}: \
367                '{}' and '{}'",
368                op_str,
369                a.class(),
370                b.class()
371            ))
372        } else {
373            self.new_type_error(format!(
374                "unsupported operand type(s) for {}: \
375                '{}' and '{}', '{}'",
376                op_str,
377                a.class(),
378                b.class(),
379                c.class()
380            ))
381        })
382    }
383
384    fn ternary_iop(
385        &self,
386        a: &PyObject,
387        b: &PyObject,
388        c: &PyObject,
389        iop_slot: PyNumberTernaryOp,
390        op_slot: PyNumberTernaryOp,
391        op_str: &str,
392    ) -> PyResult {
393        if let Some(slot) = a.class().slots.as_number.left_ternary_op(iop_slot) {
394            let x = slot(a, b, c, self)?;
395            if !x.is(&self.ctx.not_implemented) {
396                return Ok(x);
397            }
398        }
399        self.ternary_op(a, b, c, op_slot, op_str)
400    }
401
402    binary_func!(_sub, Subtract, "-");
403    binary_func!(_mod, Remainder, "%");
404    binary_func!(_divmod, Divmod, "divmod");
405    binary_func!(_lshift, Lshift, "<<");
406    binary_func!(_rshift, Rshift, ">>");
407    binary_func!(_and, And, "&");
408    binary_func!(_xor, Xor, "^");
409    binary_func!(_or, Or, "|");
410    binary_func!(_floordiv, FloorDivide, "//");
411    binary_func!(_truediv, TrueDivide, "/");
412    binary_func!(_matmul, MatrixMultiply, "@");
413
414    inplace_binary_func!(_isub, InplaceSubtract, Subtract, "-=");
415    inplace_binary_func!(_imod, InplaceRemainder, Remainder, "%=");
416    inplace_binary_func!(_ilshift, InplaceLshift, Lshift, "<<=");
417    inplace_binary_func!(_irshift, InplaceRshift, Rshift, ">>=");
418    inplace_binary_func!(_iand, InplaceAnd, And, "&=");
419    inplace_binary_func!(_ixor, InplaceXor, Xor, "^=");
420    inplace_binary_func!(_ior, InplaceOr, Or, "|=");
421    inplace_binary_func!(_ifloordiv, InplaceFloorDivide, FloorDivide, "//=");
422    inplace_binary_func!(_itruediv, InplaceTrueDivide, TrueDivide, "/=");
423    inplace_binary_func!(_imatmul, InplaceMatrixMultiply, MatrixMultiply, "@=");
424
425    ternary_func!(_pow, Power, "** or pow()");
426    inplace_ternary_func!(_ipow, InplacePower, Power, "**=");
427
428    pub fn _add(&self, a: &PyObject, b: &PyObject) -> PyResult {
429        let result = self.binary_op1(a, b, PyNumberBinaryOp::Add)?;
430        if !result.is(&self.ctx.not_implemented) {
431            return Ok(result);
432        }
433        // Check if concat slot is available directly, matching PyNumber_Add behavior
434        let seq = a.sequence_unchecked();
435        if let Some(f) = seq.slots().concat.load() {
436            let result = f(seq, b, self)?;
437            if !result.is(&self.ctx.not_implemented) {
438                return Ok(result);
439            }
440        }
441        Err(self.new_unsupported_bin_op_error(a, b, "+"))
442    }
443
444    pub fn _iadd(&self, a: &PyObject, b: &PyObject) -> PyResult {
445        let result = self.binary_iop1(a, b, PyNumberBinaryOp::InplaceAdd, PyNumberBinaryOp::Add)?;
446        if !result.is(&self.ctx.not_implemented) {
447            return Ok(result);
448        }
449        // Check inplace_concat or concat slot directly, matching PyNumber_InPlaceAdd behavior
450        let seq = a.sequence_unchecked();
451        let slots = seq.slots();
452        if let Some(f) = slots.inplace_concat.load().or_else(|| slots.concat.load()) {
453            let result = f(seq, b, self)?;
454            if !result.is(&self.ctx.not_implemented) {
455                return Ok(result);
456            }
457        }
458        Err(self.new_unsupported_bin_op_error(a, b, "+="))
459    }
460
461    pub fn _mul(&self, a: &PyObject, b: &PyObject) -> PyResult {
462        let result = self.binary_op1(a, b, PyNumberBinaryOp::Multiply)?;
463        if !result.is(&self.ctx.not_implemented) {
464            return Ok(result);
465        }
466        if let Ok(seq_a) = a.try_sequence(self) {
467            let n = b
468                .try_index(self)?
469                .as_bigint()
470                .to_isize()
471                .ok_or_else(|| self.new_overflow_error("repeated bytes are too long"))?;
472            return seq_a.repeat(n, self);
473        } else if let Ok(seq_b) = b.try_sequence(self) {
474            let n = a
475                .try_index(self)?
476                .as_bigint()
477                .to_isize()
478                .ok_or_else(|| self.new_overflow_error("repeated bytes are too long"))?;
479            return seq_b.repeat(n, self);
480        }
481        Err(self.new_unsupported_bin_op_error(a, b, "*"))
482    }
483
484    pub fn _imul(&self, a: &PyObject, b: &PyObject) -> PyResult {
485        let result = self.binary_iop1(
486            a,
487            b,
488            PyNumberBinaryOp::InplaceMultiply,
489            PyNumberBinaryOp::Multiply,
490        )?;
491        if !result.is(&self.ctx.not_implemented) {
492            return Ok(result);
493        }
494        if let Ok(seq_a) = a.try_sequence(self) {
495            let n = b
496                .try_index(self)?
497                .as_bigint()
498                .to_isize()
499                .ok_or_else(|| self.new_overflow_error("repeated bytes are too long"))?;
500            return seq_a.inplace_repeat(n, self);
501        } else if let Ok(seq_b) = b.try_sequence(self) {
502            let n = a
503                .try_index(self)?
504                .as_bigint()
505                .to_isize()
506                .ok_or_else(|| self.new_overflow_error("repeated bytes are too long"))?;
507            /* Note that the right hand operand should not be
508             * mutated in this case so inplace_repeat is not
509             * used. */
510            return seq_b.repeat(n, self);
511        }
512        Err(self.new_unsupported_bin_op_error(a, b, "*="))
513    }
514
515    pub fn _abs(&self, a: &PyObject) -> PyResult<PyObjectRef> {
516        self.get_special_method(a, identifier!(self, __abs__))?
517            .ok_or_else(|| self.new_unsupported_unary_error(a, "abs()"))?
518            .invoke((), self)
519    }
520
521    pub fn _pos(&self, a: &PyObject) -> PyResult {
522        self.get_special_method(a, identifier!(self, __pos__))?
523            .ok_or_else(|| self.new_unsupported_unary_error(a, "unary +"))?
524            .invoke((), self)
525    }
526
527    pub fn _neg(&self, a: &PyObject) -> PyResult {
528        self.get_special_method(a, identifier!(self, __neg__))?
529            .ok_or_else(|| self.new_unsupported_unary_error(a, "unary -"))?
530            .invoke((), self)
531    }
532
533    pub fn _invert(&self, a: &PyObject) -> PyResult {
534        const STR: &str = "Bitwise inversion '~' on bool is deprecated and will be removed in Python 3.16. \
535            This returns the bitwise inversion of the underlying int object and is usually not what you expect from negating a bool. \
536            Use the 'not' operator for boolean negation or ~int(x) if you really want the bitwise inversion of the underlying int.";
537        if a.fast_isinstance(self.ctx.types.bool_type) {
538            _warnings::warn(
539                self.ctx.exceptions.deprecation_warning,
540                STR.to_owned(),
541                1,
542                self,
543            )?;
544        }
545        self.get_special_method(a, identifier!(self, __invert__))?
546            .ok_or_else(|| self.new_unsupported_unary_error(a, "unary ~"))?
547            .invoke((), self)
548    }
549
550    // PyObject_Format
551    pub fn format(&self, obj: &PyObject, format_spec: PyStrRef) -> PyResult<PyStrRef> {
552        if format_spec.is_empty() {
553            let obj = match obj.to_owned().downcast_exact::<PyStr>(self) {
554                Ok(s) => return Ok(s.into_pyref()),
555                Err(obj) => obj,
556            };
557            if obj.class().is(self.ctx.types.int_type) {
558                return obj.str(self);
559            }
560        }
561        let bound_format = self
562            .get_special_method(obj, identifier!(self, __format__))?
563            .ok_or_else(|| {
564                self.new_type_error(format!(
565                    "Type {} doesn't define __format__",
566                    obj.class().name()
567                ))
568            })?;
569        let formatted = bound_format.invoke((format_spec,), self)?;
570        formatted.downcast().map_err(|result| {
571            self.new_type_error(format!(
572                "__format__ must return a str, not {}",
573                &result.class().name()
574            ))
575        })
576    }
577    pub fn format_utf8(&self, obj: &PyObject, format_spec: PyStrRef) -> PyResult<PyRef<PyUtf8Str>> {
578        self.format(obj, format_spec)?.try_into_utf8(self)
579    }
580
581    pub fn _contains(&self, haystack: &PyObject, needle: &PyObject) -> PyResult<bool> {
582        let seq = haystack.sequence_unchecked();
583        seq.contains(needle, self)
584    }
585}