rustpython_vm/vm/
context.rs

1use crate::{
2    builtins::{
3        bytes,
4        code::{self, PyCode},
5        descriptor::{
6            MemberGetter, MemberKind, MemberSetter, MemberSetterFunc, PyDescriptorOwned,
7            PyMemberDef, PyMemberDescriptor,
8        },
9        getset::PyGetSet,
10        object, pystr,
11        type_::PyAttributes,
12        PyBaseException, PyBytes, PyComplex, PyDict, PyDictRef, PyEllipsis, PyFloat, PyFrozenSet,
13        PyInt, PyIntRef, PyList, PyListRef, PyNone, PyNotImplemented, PyStr, PyStrInterned,
14        PyTuple, PyTupleRef, PyType, PyTypeRef,
15    },
16    class::{PyClassImpl, StaticType},
17    common::rc::PyRc,
18    exceptions,
19    function::{
20        HeapMethodDef, IntoPyGetterFunc, IntoPyNativeFn, IntoPySetterFunc, PyMethodDef,
21        PyMethodFlags,
22    },
23    intern::{InternableString, MaybeInternedString, StringPool},
24    object::{Py, PyObjectPayload, PyObjectRef, PyPayload, PyRef},
25    types::{PyTypeFlags, PyTypeSlots, TypeZoo},
26    PyResult, VirtualMachine,
27};
28use malachite_bigint::BigInt;
29use num_complex::Complex64;
30use num_traits::ToPrimitive;
31use rustpython_common::lock::PyRwLock;
32
33#[derive(Debug)]
34pub struct Context {
35    pub true_value: PyIntRef,
36    pub false_value: PyIntRef,
37    pub none: PyRef<PyNone>,
38    pub empty_tuple: PyTupleRef,
39    pub empty_frozenset: PyRef<PyFrozenSet>,
40    pub empty_str: &'static PyStrInterned,
41    pub empty_bytes: PyRef<PyBytes>,
42    pub ellipsis: PyRef<PyEllipsis>,
43    pub not_implemented: PyRef<PyNotImplemented>,
44
45    pub types: TypeZoo,
46    pub exceptions: exceptions::ExceptionZoo,
47    pub int_cache_pool: Vec<PyIntRef>,
48    // there should only be exact objects of str in here, no non-str objects and no subclasses
49    pub(crate) string_pool: StringPool,
50    pub(crate) slot_new_wrapper: PyMethodDef,
51    pub names: ConstName,
52}
53
54macro_rules! declare_const_name {
55    ($($name:ident,)*) => {
56        #[derive(Debug, Clone, Copy)]
57        #[allow(non_snake_case)]
58        pub struct ConstName {
59            $(pub $name: &'static PyStrInterned,)*
60        }
61
62        impl ConstName {
63            unsafe fn new(pool: &StringPool, typ: &PyTypeRef) -> Self {
64                Self {
65                    $($name: pool.intern(stringify!($name), typ.clone()),)*
66                }
67            }
68        }
69    }
70}
71
72declare_const_name! {
73    True,
74    False,
75    None,
76    NotImplemented,
77    Ellipsis,
78
79    // magic methods
80    __abs__,
81    __abstractmethods__,
82    __add__,
83    __aenter__,
84    __aexit__,
85    __aiter__,
86    __alloc__,
87    __all__,
88    __and__,
89    __anext__,
90    __annotations__,
91    __args__,
92    __await__,
93    __bases__,
94    __bool__,
95    __build_class__,
96    __builtins__,
97    __bytes__,
98    __call__,
99    __ceil__,
100    __cformat__,
101    __class__,
102    __classcell__,
103    __class_getitem__,
104    __complex__,
105    __contains__,
106    __copy__,
107    __deepcopy__,
108    __del__,
109    __delattr__,
110    __delete__,
111    __delitem__,
112    __dict__,
113    __dir__,
114    __div__,
115    __divmod__,
116    __doc__,
117    __enter__,
118    __eq__,
119    __exit__,
120    __file__,
121    __float__,
122    __floor__,
123    __floordiv__,
124    __format__,
125    __fspath__,
126    __ge__,
127    __get__,
128    __getattr__,
129    __getattribute__,
130    __getformat__,
131    __getitem__,
132    __getnewargs__,
133    __getstate__,
134    __gt__,
135    __hash__,
136    __iadd__,
137    __iand__,
138    __idiv__,
139    __ifloordiv__,
140    __ilshift__,
141    __imatmul__,
142    __imod__,
143    __import__,
144    __imul__,
145    __index__,
146    __init__,
147    __init_subclass__,
148    __instancecheck__,
149    __int__,
150    __invert__,
151    __ior__,
152    __ipow__,
153    __irshift__,
154    __isub__,
155    __iter__,
156    __itruediv__,
157    __ixor__,
158    __jit__,  // RustPython dialect
159    __le__,
160    __len__,
161    __length_hint__,
162    __lshift__,
163    __lt__,
164    __main__,
165    __match_args__,
166    __matmul__,
167    __missing__,
168    __mod__,
169    __module__,
170    __mro_entries__,
171    __mul__,
172    __name__,
173    __ne__,
174    __neg__,
175    __new__,
176    __next__,
177    __objclass__,
178    __or__,
179    __orig_bases__,
180    __orig_class__,
181    __origin__,
182    __parameters__,
183    __pos__,
184    __pow__,
185    __prepare__,
186    __qualname__,
187    __radd__,
188    __rand__,
189    __rdiv__,
190    __rdivmod__,
191    __reduce__,
192    __reduce_ex__,
193    __repr__,
194    __reversed__,
195    __rfloordiv__,
196    __rlshift__,
197    __rmatmul__,
198    __rmod__,
199    __rmul__,
200    __ror__,
201    __round__,
202    __rpow__,
203    __rrshift__,
204    __rshift__,
205    __rsub__,
206    __rtruediv__,
207    __rxor__,
208    __set__,
209    __setattr__,
210    __setitem__,
211    __setstate__,
212    __set_name__,
213    __slots__,
214    __slotnames__,
215    __str__,
216    __sub__,
217    __subclasscheck__,
218    __subclasshook__,
219    __subclasses__,
220    __sizeof__,
221    __truediv__,
222    __trunc__,
223    __xor__,
224
225    // common names
226    _attributes,
227    _fields,
228    _showwarnmsg,
229    decode,
230    encode,
231    keys,
232    items,
233    values,
234    version,
235    update,
236    copy,
237    flush,
238    close,
239    WarningMessage,
240}
241
242// Basic objects:
243impl Context {
244    pub const INT_CACHE_POOL_RANGE: std::ops::RangeInclusive<i32> = (-5)..=256;
245    const INT_CACHE_POOL_MIN: i32 = *Self::INT_CACHE_POOL_RANGE.start();
246
247    pub fn genesis() -> &'static PyRc<Self> {
248        rustpython_common::static_cell! {
249            static CONTEXT: PyRc<Context>;
250        }
251        CONTEXT.get_or_init(|| PyRc::new(Self::init_genesis()))
252    }
253
254    fn init_genesis() -> Self {
255        flame_guard!("init Context");
256        let types = TypeZoo::init();
257        let exceptions = exceptions::ExceptionZoo::init();
258
259        #[inline]
260        fn create_object<T: PyObjectPayload + PyPayload>(
261            payload: T,
262            cls: &'static Py<PyType>,
263        ) -> PyRef<T> {
264            PyRef::new_ref(payload, cls.to_owned(), None)
265        }
266
267        let none = create_object(PyNone, PyNone::static_type());
268        let ellipsis = create_object(PyEllipsis, PyEllipsis::static_type());
269        let not_implemented = create_object(PyNotImplemented, PyNotImplemented::static_type());
270
271        let int_cache_pool = Self::INT_CACHE_POOL_RANGE
272            .map(|v| {
273                PyRef::new_ref(
274                    PyInt::from(BigInt::from(v)),
275                    types.int_type.to_owned(),
276                    None,
277                )
278            })
279            .collect();
280
281        let true_value = create_object(PyInt::from(1), types.bool_type);
282        let false_value = create_object(PyInt::from(0), types.bool_type);
283
284        let empty_tuple = create_object(
285            PyTuple::new_unchecked(Vec::new().into_boxed_slice()),
286            types.tuple_type,
287        );
288        let empty_frozenset = PyRef::new_ref(
289            PyFrozenSet::default(),
290            types.frozenset_type.to_owned(),
291            None,
292        );
293
294        let string_pool = StringPool::default();
295        let names = unsafe { ConstName::new(&string_pool, &types.str_type.to_owned()) };
296
297        let slot_new_wrapper = PyMethodDef::new_const(
298            names.__new__.as_str(),
299            PyType::__new__,
300            PyMethodFlags::METHOD,
301            None,
302        );
303
304        let empty_str = unsafe { string_pool.intern("", types.str_type.to_owned()) };
305        let empty_bytes = create_object(PyBytes::from(Vec::new()), types.bytes_type);
306        Context {
307            true_value,
308            false_value,
309            none,
310            empty_tuple,
311            empty_frozenset,
312            empty_str,
313            empty_bytes,
314
315            ellipsis,
316            not_implemented,
317
318            types,
319            exceptions,
320            int_cache_pool,
321            string_pool,
322            slot_new_wrapper,
323            names,
324        }
325    }
326
327    pub fn intern_str<S: InternableString>(&self, s: S) -> &'static PyStrInterned {
328        unsafe { self.string_pool.intern(s, self.types.str_type.to_owned()) }
329    }
330
331    pub fn interned_str<S: MaybeInternedString + ?Sized>(
332        &self,
333        s: &S,
334    ) -> Option<&'static PyStrInterned> {
335        self.string_pool.interned(s)
336    }
337
338    #[inline(always)]
339    pub fn none(&self) -> PyObjectRef {
340        self.none.clone().into()
341    }
342
343    #[inline(always)]
344    pub fn ellipsis(&self) -> PyObjectRef {
345        self.ellipsis.clone().into()
346    }
347
348    #[inline(always)]
349    pub fn not_implemented(&self) -> PyObjectRef {
350        self.not_implemented.clone().into()
351    }
352
353    // universal pyref constructor
354    pub fn new_pyref<T, P>(&self, value: T) -> PyRef<P>
355    where
356        T: Into<P>,
357        P: PyPayload,
358    {
359        value.into().into_ref(self)
360    }
361
362    // shortcuts for common type
363
364    #[inline]
365    pub fn new_int<T: Into<BigInt> + ToPrimitive>(&self, i: T) -> PyIntRef {
366        if let Some(i) = i.to_i32() {
367            if Self::INT_CACHE_POOL_RANGE.contains(&i) {
368                let inner_idx = (i - Self::INT_CACHE_POOL_MIN) as usize;
369                return self.int_cache_pool[inner_idx].clone();
370            }
371        }
372        PyInt::from(i).into_ref(self)
373    }
374
375    #[inline]
376    pub fn new_bigint(&self, i: &BigInt) -> PyIntRef {
377        if let Some(i) = i.to_i32() {
378            if Self::INT_CACHE_POOL_RANGE.contains(&i) {
379                let inner_idx = (i - Self::INT_CACHE_POOL_MIN) as usize;
380                return self.int_cache_pool[inner_idx].clone();
381            }
382        }
383        PyInt::from(i.clone()).into_ref(self)
384    }
385
386    #[inline]
387    pub fn new_float(&self, value: f64) -> PyRef<PyFloat> {
388        PyFloat::from(value).into_ref(self)
389    }
390
391    #[inline]
392    pub fn new_complex(&self, value: Complex64) -> PyRef<PyComplex> {
393        PyComplex::from(value).into_ref(self)
394    }
395
396    #[inline]
397    pub fn new_str(&self, s: impl Into<pystr::PyStr>) -> PyRef<PyStr> {
398        pystr::PyStr::new_ref(s, self)
399    }
400
401    pub fn interned_or_new_str<S, M>(&self, s: S) -> PyRef<PyStr>
402    where
403        S: Into<PyStr> + AsRef<M>,
404        M: MaybeInternedString,
405    {
406        match self.interned_str(s.as_ref()) {
407            Some(s) => s.to_owned(),
408            None => self.new_str(s),
409        }
410    }
411
412    #[inline]
413    pub fn new_bytes(&self, data: Vec<u8>) -> PyRef<bytes::PyBytes> {
414        bytes::PyBytes::new_ref(data, self)
415    }
416
417    #[inline(always)]
418    pub fn new_bool(&self, b: bool) -> PyIntRef {
419        let value = if b {
420            &self.true_value
421        } else {
422            &self.false_value
423        };
424        value.clone()
425    }
426
427    #[inline(always)]
428    pub fn new_tuple(&self, elements: Vec<PyObjectRef>) -> PyTupleRef {
429        PyTuple::new_ref(elements, self)
430    }
431
432    #[inline(always)]
433    pub fn new_list(&self, elements: Vec<PyObjectRef>) -> PyListRef {
434        PyList::new_ref(elements, self)
435    }
436
437    #[inline(always)]
438    pub fn new_dict(&self) -> PyDictRef {
439        PyDict::new_ref(self)
440    }
441
442    pub fn new_class(
443        &self,
444        module: Option<&str>,
445        name: &str,
446        base: PyTypeRef,
447        slots: PyTypeSlots,
448    ) -> PyTypeRef {
449        let mut attrs = PyAttributes::default();
450        if let Some(module) = module {
451            attrs.insert(identifier!(self, __module__), self.new_str(module).into());
452        };
453        PyType::new_heap(
454            name,
455            vec![base],
456            attrs,
457            slots,
458            self.types.type_type.to_owned(),
459            self,
460        )
461        .unwrap()
462    }
463
464    pub fn new_exception_type(
465        &self,
466        module: &str,
467        name: &str,
468        bases: Option<Vec<PyTypeRef>>,
469    ) -> PyTypeRef {
470        let bases = if let Some(bases) = bases {
471            bases
472        } else {
473            vec![self.exceptions.exception_type.to_owned()]
474        };
475        let mut attrs = PyAttributes::default();
476        attrs.insert(identifier!(self, __module__), self.new_str(module).into());
477
478        let interned_name = self.intern_str(name);
479        PyType::new_heap(
480            name,
481            bases,
482            attrs,
483            PyTypeSlots {
484                name: interned_name.as_str(),
485                ..PyBaseException::make_slots()
486            },
487            self.types.type_type.to_owned(),
488            self,
489        )
490        .unwrap()
491    }
492
493    pub fn new_method_def<F, FKind>(
494        &self,
495        name: &'static str,
496        f: F,
497        flags: PyMethodFlags,
498        doc: Option<&'static str>,
499    ) -> PyRef<HeapMethodDef>
500    where
501        F: IntoPyNativeFn<FKind>,
502    {
503        let def = PyMethodDef {
504            name,
505            func: Box::leak(Box::new(f.into_func())),
506            flags,
507            doc,
508        };
509        let payload = HeapMethodDef::new(def);
510        PyRef::new_ref(payload, self.types.method_def.to_owned(), None)
511    }
512
513    #[inline]
514    pub fn new_member(
515        &self,
516        name: &str,
517        member_kind: MemberKind,
518        getter: fn(&VirtualMachine, PyObjectRef) -> PyResult,
519        setter: MemberSetterFunc,
520        class: &'static Py<PyType>,
521    ) -> PyRef<PyMemberDescriptor> {
522        let member_def = PyMemberDef {
523            name: name.to_owned(),
524            kind: member_kind,
525            getter: MemberGetter::Getter(getter),
526            setter: MemberSetter::Setter(setter),
527            doc: None,
528        };
529        let member_descriptor = PyMemberDescriptor {
530            common: PyDescriptorOwned {
531                typ: class.to_owned(),
532                name: self.intern_str(name),
533                qualname: PyRwLock::new(None),
534            },
535            member: member_def,
536        };
537        member_descriptor.into_ref(self)
538    }
539
540    pub fn new_readonly_getset<F, T>(
541        &self,
542        name: impl Into<String>,
543        class: &'static Py<PyType>,
544        f: F,
545    ) -> PyRef<PyGetSet>
546    where
547        F: IntoPyGetterFunc<T>,
548    {
549        let name = name.into();
550        let getset = PyGetSet::new(name, class).with_get(f);
551        PyRef::new_ref(getset, self.types.getset_type.to_owned(), None)
552    }
553
554    pub fn new_getset<G, S, T, U>(
555        &self,
556        name: impl Into<String>,
557        class: &'static Py<PyType>,
558        g: G,
559        s: S,
560    ) -> PyRef<PyGetSet>
561    where
562        G: IntoPyGetterFunc<T>,
563        S: IntoPySetterFunc<U>,
564    {
565        let name = name.into();
566        let getset = PyGetSet::new(name, class).with_get(g).with_set(s);
567        PyRef::new_ref(getset, self.types.getset_type.to_owned(), None)
568    }
569
570    pub fn new_base_object(&self, class: PyTypeRef, dict: Option<PyDictRef>) -> PyObjectRef {
571        debug_assert_eq!(
572            class.slots.flags.contains(PyTypeFlags::HAS_DICT),
573            dict.is_some()
574        );
575        PyRef::new_ref(object::PyBaseObject, class, dict).into()
576    }
577
578    pub fn new_code(&self, code: impl code::IntoCodeObject) -> PyRef<PyCode> {
579        let code = code.into_code_object(self);
580        PyRef::new_ref(PyCode { code }, self.types.code_type.to_owned(), None)
581    }
582}
583
584impl AsRef<Context> for Context {
585    fn as_ref(&self) -> &Self {
586        self
587    }
588}