Skip to main content

rustpython_vm/vm/
context.rs

1use crate::{
2    PyResult, VirtualMachine,
3    builtins::{
4        PyByteArray, PyBytes, PyComplex, PyDict, PyDictRef, PyEllipsis, PyFloat, PyFrozenSet,
5        PyInt, PyIntRef, PyList, PyListRef, PyNone, PyNotImplemented, PyStr, PyStrInterned,
6        PyTuple, PyTupleRef, PyType, PyTypeRef, PyUtf8Str,
7        bool_::PyBool,
8        code::{self, PyCode},
9        descriptor::{
10            MemberGetter, MemberKind, MemberSetter, MemberSetterFunc, PyDescriptorOwned,
11            PyMemberDef, PyMemberDescriptor,
12        },
13        getset::PyGetSet,
14        object, pystr,
15        type_::PyAttributes,
16    },
17    bytecode::{self, CodeFlags, CodeUnit, Instruction},
18    class::StaticType,
19    common::rc::PyRc,
20    exceptions,
21    function::{
22        HeapMethodDef, IntoPyGetterFunc, IntoPyNativeFn, IntoPySetterFunc, PyMethodDef,
23        PyMethodFlags,
24    },
25    intern::{InternableString, MaybeInternedString, StringPool},
26    object::{Py, PyObjectPayload, PyObjectRef, PyPayload, PyRef},
27    types::{PyTypeFlags, PyTypeSlots, TypeZoo},
28};
29use malachite_bigint::BigInt;
30use num_complex::Complex64;
31use num_traits::ToPrimitive;
32use rustpython_common::lock::PyRwLock;
33use rustpython_compiler_core::{OneIndexed, SourceLocation};
34
35#[derive(Debug)]
36pub struct Context {
37    pub true_value: PyRef<PyBool>,
38    pub false_value: PyRef<PyBool>,
39    pub none: PyRef<PyNone>,
40    pub empty_tuple: PyTupleRef,
41    pub empty_frozenset: PyRef<PyFrozenSet>,
42    pub empty_str: &'static PyStrInterned,
43    pub empty_bytes: PyRef<PyBytes>,
44    pub ellipsis: PyRef<PyEllipsis>,
45    pub not_implemented: PyRef<PyNotImplemented>,
46
47    pub typing_no_default: PyRef<crate::stdlib::_typing::NoDefault>,
48
49    pub types: TypeZoo,
50    pub exceptions: exceptions::ExceptionZoo,
51    pub int_cache_pool: Vec<PyIntRef>,
52    pub(crate) latin1_char_cache: Vec<PyRef<PyStr>>,
53    pub(crate) ascii_char_cache: Vec<PyRef<PyStr>>,
54    pub(crate) init_cleanup_code: PyRef<PyCode>,
55    // there should only be exact objects of str in here, no non-str objects and no subclasses
56    pub(crate) string_pool: StringPool,
57    pub(crate) slot_new_wrapper: PyMethodDef,
58    pub names: ConstName,
59
60    // GC module state (callbacks and garbage lists)
61    pub gc_callbacks: PyListRef,
62    pub gc_garbage: PyListRef,
63}
64
65macro_rules! declare_const_name {
66    ($($name:ident$(: $s:literal)?,)*) => {
67        #[derive(Debug, Clone, Copy)]
68        #[allow(non_snake_case)]
69        pub struct ConstName {
70            $(pub $name: &'static PyStrInterned,)*
71        }
72
73        impl ConstName {
74            unsafe fn new(pool: &StringPool, typ: &Py<PyType>) -> Self {
75                Self {
76                    $($name: unsafe { pool.intern(declare_const_name!(@string $name $($s)?), typ.to_owned()) },)*
77                }
78            }
79        }
80    };
81    (@string $name:ident) => { stringify!($name) };
82    (@string $name:ident $string:literal) => { $string };
83}
84
85declare_const_name! {
86    True,
87    False,
88    None,
89    NotImplemented,
90    Ellipsis,
91
92    // magic methods
93    __abs__,
94    __abstractmethods__,
95    __add__,
96    __aenter__,
97    __aexit__,
98    __aiter__,
99    __alloc__,
100    __all__,
101    __and__,
102    __anext__,
103    __annotate__,
104    __annotate_func__,
105    __annotations__,
106    __annotations_cache__,
107    __args__,
108    __await__,
109    __bases__,
110    __bool__,
111    __build_class__,
112    __builtins__,
113    __bytes__,
114    __cached__,
115    __call__,
116    __ceil__,
117    __cformat__,
118    __class__,
119    __class_getitem__,
120    __classcell__,
121    __classdictcell__,
122    __complex__,
123    __contains__,
124    __copy__,
125    __deepcopy__,
126    __del__,
127    __delattr__,
128    __delete__,
129    __delitem__,
130    __dict__,
131    __dir__,
132    __div__,
133    __divmod__,
134    __doc__,
135    __enter__,
136    __eq__,
137    __exit__,
138    __file__,
139    __firstlineno__,
140    __float__,
141    __floor__,
142    __floordiv__,
143    __format__,
144    __fspath__,
145    __ge__,
146    __get__,
147    __getattr__,
148    __getattribute__,
149    __getformat__,
150    __getitem__,
151    __getnewargs__,
152    __getnewargs_ex__,
153    __getstate__,
154    __gt__,
155    __hash__,
156    __iadd__,
157    __iand__,
158    __idiv__,
159    __ifloordiv__,
160    __ilshift__,
161    __imatmul__,
162    __imod__,
163    __import__,
164    __imul__,
165    __index__,
166    __init__,
167    __init_subclass__,
168    __instancecheck__,
169    __int__,
170    __invert__,
171    __ior__,
172    __ipow__,
173    __irshift__,
174    __isub__,
175    __iter__,
176    __itruediv__,
177    __ixor__,
178    __jit__,  // RustPython dialect
179    __le__,
180    __len__,
181    __length_hint__,
182    __lshift__,
183    __lt__,
184    __main__,
185    __match_args__,
186    __matmul__,
187    __missing__,
188    __mod__,
189    __module__,
190    __mro_entries__,
191    __mul__,
192    __name__,
193    __ne__,
194    __neg__,
195    __new__,
196    __next__,
197    __objclass__,
198    __or__,
199    __orig_bases__,
200    __orig_class__,
201    __origin__,
202    __parameters__,
203    __pos__,
204    __pow__,
205    __prepare__,
206    __qualname__,
207    __radd__,
208    __rand__,
209    __rdiv__,
210    __rdivmod__,
211    __reduce__,
212    __reduce_ex__,
213    __repr__,
214    __reversed__,
215    __rfloordiv__,
216    __rlshift__,
217    __rmatmul__,
218    __rmod__,
219    __rmul__,
220    __ror__,
221    __round__,
222    __rpow__,
223    __rrshift__,
224    __rshift__,
225    __rsub__,
226    __rtruediv__,
227    __rxor__,
228    __set__,
229    __setattr__,
230    __setitem__,
231    __setstate__,
232    __set_name__,
233    __slots__,
234    __slotnames__,
235    __str__,
236    __sub__,
237    __subclasscheck__,
238    __subclasshook__,
239    __subclasses__,
240    __sizeof__,
241    __truediv__,
242    __trunc__,
243    __type_params__,
244    __typing_subst__,
245    __typing_is_unpacked_typevartuple__,
246    __typing_prepare_subst__,
247    __typing_unpacked_tuple_args__,
248    __weakref__,
249    __xor__,
250
251    // common names
252    _attributes,
253    _fields,
254    _defaultaction,
255    _onceregistry,
256    _showwarnmsg,
257    defaultaction,
258    onceregistry,
259    filters,
260    backslashreplace,
261    close,
262    copy,
263    decode,
264    encode,
265    flush,
266    ignore,
267    items,
268    keys,
269    modules,
270    n_fields,
271    n_sequence_fields,
272    n_unnamed_fields,
273    namereplace,
274    replace,
275    strict,
276    surrogateescape,
277    surrogatepass,
278    update,
279    utf_8: "utf-8",
280    values,
281    version,
282    WarningMessage,
283    xmlcharrefreplace,
284}
285
286// Basic objects:
287impl Context {
288    pub const INT_CACHE_POOL_RANGE: core::ops::RangeInclusive<i32> = (-5)..=256;
289    const INT_CACHE_POOL_MIN: i32 = *Self::INT_CACHE_POOL_RANGE.start();
290
291    pub fn genesis() -> &'static PyRc<Self> {
292        rustpython_common::static_cell! {
293            static CONTEXT: PyRc<Context>;
294        }
295        CONTEXT.get_or_init(|| {
296            let ctx = PyRc::new(Self::init_genesis());
297            // SAFETY: ctx is heap-allocated via PyRc and will be stored in
298            // the CONTEXT static cell, so the Context lives for 'static.
299            let ctx_ref: &'static Context = unsafe { &*PyRc::as_ptr(&ctx) };
300            crate::types::TypeZoo::extend(ctx_ref);
301            crate::exceptions::ExceptionZoo::extend(ctx_ref);
302            ctx
303        })
304    }
305
306    fn init_genesis() -> Self {
307        flame_guard!("init Context");
308        let types = TypeZoo::init();
309        let exceptions = exceptions::ExceptionZoo::init();
310
311        #[inline]
312        fn create_object<T: PyObjectPayload>(payload: T, cls: &'static Py<PyType>) -> PyRef<T> {
313            PyRef::new_ref(payload, cls.to_owned(), None)
314        }
315
316        let none = create_object(PyNone, PyNone::static_type());
317        let ellipsis = create_object(PyEllipsis, PyEllipsis::static_type());
318        let not_implemented = create_object(PyNotImplemented, PyNotImplemented::static_type());
319
320        let typing_no_default = create_object(
321            crate::stdlib::_typing::NoDefault,
322            crate::stdlib::_typing::NoDefault::static_type(),
323        );
324
325        let int_cache_pool = Self::INT_CACHE_POOL_RANGE
326            .map(|v| {
327                PyRef::new_ref(
328                    PyInt::from(BigInt::from(v)),
329                    types.int_type.to_owned(),
330                    None,
331                )
332            })
333            .collect();
334        let latin1_char_cache: Vec<PyRef<PyStr>> = (0u8..=255)
335            .map(|b| create_object(PyStr::from(char::from(b)), types.str_type))
336            .collect();
337        let ascii_char_cache = latin1_char_cache[..128].to_vec();
338
339        let true_value = create_object(PyBool(PyInt::from(1)), types.bool_type);
340        let false_value = create_object(PyBool(PyInt::from(0)), types.bool_type);
341
342        let empty_tuple = create_object(
343            PyTuple::new_unchecked(Vec::new().into_boxed_slice()),
344            types.tuple_type,
345        );
346        let empty_frozenset = PyRef::new_ref(
347            PyFrozenSet::default(),
348            types.frozenset_type.to_owned(),
349            None,
350        );
351
352        let string_pool = StringPool::default();
353        let names = unsafe { ConstName::new(&string_pool, types.str_type) };
354
355        let slot_new_wrapper = PyMethodDef::new_const(
356            names.__new__.as_str(),
357            PyType::__new__,
358            PyMethodFlags::METHOD,
359            None,
360        );
361        let init_cleanup_code = Self::new_init_cleanup_code(&types, &names);
362
363        let empty_str = unsafe { string_pool.intern("", types.str_type.to_owned()) };
364        let empty_bytes = create_object(PyBytes::from(Vec::new()), types.bytes_type);
365
366        // GC callbacks and garbage lists
367        let gc_callbacks = PyRef::new_ref(PyList::default(), types.list_type.to_owned(), None);
368        let gc_garbage = PyRef::new_ref(PyList::default(), types.list_type.to_owned(), None);
369
370        Self {
371            true_value,
372            false_value,
373            none,
374            empty_tuple,
375            empty_frozenset,
376            empty_str,
377            empty_bytes,
378            ellipsis,
379
380            not_implemented,
381            typing_no_default,
382
383            types,
384            exceptions,
385            int_cache_pool,
386            latin1_char_cache,
387            ascii_char_cache,
388            init_cleanup_code,
389            string_pool,
390            slot_new_wrapper,
391            names,
392
393            gc_callbacks,
394            gc_garbage,
395        }
396    }
397
398    fn new_init_cleanup_code(types: &TypeZoo, names: &ConstName) -> PyRef<PyCode> {
399        let loc = SourceLocation {
400            line: OneIndexed::MIN,
401            character_offset: OneIndexed::from_zero_indexed(0),
402        };
403        let instructions = [
404            CodeUnit {
405                op: Instruction::ExitInitCheck,
406                arg: 0.into(),
407            },
408            CodeUnit {
409                op: Instruction::ReturnValue,
410                arg: 0.into(),
411            },
412            CodeUnit {
413                op: Instruction::Resume {
414                    context: bytecode::Arg::marker(),
415                },
416                arg: 0.into(),
417            },
418        ];
419        let code = bytecode::CodeObject {
420            instructions: instructions.into(),
421            locations: vec![(loc, loc); instructions.len()].into_boxed_slice(),
422            flags: CodeFlags::OPTIMIZED,
423            posonlyarg_count: 0,
424            arg_count: 0,
425            kwonlyarg_count: 0,
426            source_path: names.__init__,
427            first_line_number: None,
428            max_stackdepth: 2,
429            obj_name: names.__init__,
430            qualname: names.__init__,
431            constants: core::iter::empty().collect(),
432            names: Vec::new().into_boxed_slice(),
433            varnames: Vec::new().into_boxed_slice(),
434            cellvars: Vec::new().into_boxed_slice(),
435            freevars: Vec::new().into_boxed_slice(),
436            localspluskinds: Vec::new().into_boxed_slice(),
437            linetable: Vec::new().into_boxed_slice(),
438            exceptiontable: Vec::new().into_boxed_slice(),
439        };
440        PyRef::new_ref(PyCode::new(code), types.code_type.to_owned(), None)
441    }
442
443    pub fn intern_str<S: InternableString>(&self, s: S) -> &'static PyStrInterned {
444        unsafe { self.string_pool.intern(s, self.types.str_type.to_owned()) }
445    }
446
447    pub fn interned_str<S: MaybeInternedString + ?Sized>(
448        &self,
449        s: &S,
450    ) -> Option<&'static PyStrInterned> {
451        self.string_pool.interned(s)
452    }
453
454    #[inline(always)]
455    pub fn none(&self) -> PyObjectRef {
456        self.none.clone().into()
457    }
458
459    #[inline(always)]
460    pub fn not_implemented(&self) -> PyObjectRef {
461        self.not_implemented.clone().into()
462    }
463
464    #[inline]
465    pub fn empty_tuple_typed<T>(&self) -> &Py<PyTuple<T>> {
466        let py: &Py<PyTuple> = &self.empty_tuple;
467        unsafe { core::mem::transmute(py) }
468    }
469
470    // universal pyref constructor
471    pub fn new_pyref<T, P>(&self, value: T) -> PyRef<P>
472    where
473        T: Into<P>,
474        P: PyPayload + core::fmt::Debug,
475    {
476        value.into().into_ref(self)
477    }
478
479    // shortcuts for common type
480
481    #[inline]
482    pub fn new_int<T: Into<BigInt> + ToPrimitive>(&self, i: T) -> PyIntRef {
483        if let Some(i) = i.to_i32()
484            && Self::INT_CACHE_POOL_RANGE.contains(&i)
485        {
486            let inner_idx = (i - Self::INT_CACHE_POOL_MIN) as usize;
487            return self.int_cache_pool[inner_idx].clone();
488        }
489        PyInt::from(i).into_ref(self)
490    }
491
492    #[inline]
493    pub fn new_bigint(&self, i: &BigInt) -> PyIntRef {
494        if let Some(i) = i.to_i32()
495            && Self::INT_CACHE_POOL_RANGE.contains(&i)
496        {
497            let inner_idx = (i - Self::INT_CACHE_POOL_MIN) as usize;
498            return self.int_cache_pool[inner_idx].clone();
499        }
500        PyInt::from(i.clone()).into_ref(self)
501    }
502
503    #[inline]
504    pub fn new_float(&self, value: f64) -> PyRef<PyFloat> {
505        PyFloat::from(value).into_ref(self)
506    }
507
508    #[inline]
509    pub fn new_complex(&self, value: Complex64) -> PyRef<PyComplex> {
510        PyComplex::from(value).into_ref(self)
511    }
512
513    #[inline]
514    pub fn latin1_char(&self, ch: u8) -> PyRef<PyStr> {
515        self.latin1_char_cache[ch as usize].clone()
516    }
517
518    #[inline]
519    fn latin1_singleton_index(s: &PyStr) -> Option<u8> {
520        let mut cps = s.as_wtf8().code_points();
521        let cp = cps.next()?;
522        if cps.next().is_some() {
523            return None;
524        }
525        u8::try_from(cp.to_u32()).ok()
526    }
527
528    #[inline]
529    pub fn new_str(&self, s: impl Into<pystr::PyStr>) -> PyRef<PyStr> {
530        let s = s.into();
531        if let Some(ch) = Self::latin1_singleton_index(&s) {
532            return self.latin1_char(ch);
533        }
534        s.into_ref(self)
535    }
536
537    #[inline]
538    pub fn new_utf8_str(&self, s: impl Into<PyUtf8Str>) -> PyRef<PyUtf8Str> {
539        s.into().into_ref(self)
540    }
541
542    pub fn interned_or_new_str<S, M>(&self, s: S) -> PyRef<PyStr>
543    where
544        S: Into<PyStr> + AsRef<M>,
545        M: MaybeInternedString,
546    {
547        match self.interned_str(s.as_ref()) {
548            Some(s) => s.to_owned(),
549            None => self.new_str(s),
550        }
551    }
552
553    #[inline]
554    pub fn new_bytes(&self, data: Vec<u8>) -> PyRef<PyBytes> {
555        if data.is_empty() {
556            self.empty_bytes.clone()
557        } else {
558            PyBytes::from(data).into_ref(self)
559        }
560    }
561
562    #[inline]
563    pub fn new_bytearray(&self, data: Vec<u8>) -> PyRef<PyByteArray> {
564        PyByteArray::from(data).into_ref(self)
565    }
566
567    #[inline(always)]
568    pub fn new_bool(&self, b: bool) -> PyRef<PyBool> {
569        let value = if b {
570            &self.true_value
571        } else {
572            &self.false_value
573        };
574        value.to_owned()
575    }
576
577    #[inline(always)]
578    pub fn new_tuple(&self, elements: Vec<PyObjectRef>) -> PyTupleRef {
579        PyTuple::new_ref(elements, self)
580    }
581
582    #[inline(always)]
583    pub fn new_list(&self, elements: Vec<PyObjectRef>) -> PyListRef {
584        PyList::from(elements).into_ref(self)
585    }
586
587    #[inline(always)]
588    pub fn new_dict(&self) -> PyDictRef {
589        PyDict::default().into_ref(self)
590    }
591
592    pub fn new_class(
593        &self,
594        module: Option<&str>,
595        name: &str,
596        base: PyTypeRef,
597        slots: PyTypeSlots,
598    ) -> PyTypeRef {
599        let mut attrs = PyAttributes::default();
600        if let Some(module) = module {
601            attrs.insert(identifier!(self, __module__), self.new_str(module).into());
602        };
603        PyType::new_heap(
604            name,
605            vec![base],
606            attrs,
607            slots,
608            self.types.type_type.to_owned(),
609            self,
610        )
611        .unwrap()
612    }
613
614    pub fn new_exception_type(
615        &self,
616        module: &str,
617        name: &str,
618        bases: Option<Vec<PyTypeRef>>,
619    ) -> PyTypeRef {
620        let bases = if let Some(bases) = bases {
621            bases
622        } else {
623            vec![self.exceptions.exception_type.to_owned()]
624        };
625        let mut attrs = PyAttributes::default();
626        attrs.insert(identifier!(self, __module__), self.new_str(module).into());
627
628        let interned_name = self.intern_str(name);
629        let slots = PyTypeSlots {
630            name: interned_name.as_str(),
631            basicsize: 0,
632            flags: PyTypeFlags::heap_type_flags() | PyTypeFlags::HAS_DICT,
633            ..PyTypeSlots::default()
634        };
635        PyType::new_heap(
636            name,
637            bases,
638            attrs,
639            slots,
640            self.types.type_type.to_owned(),
641            self,
642        )
643        .unwrap()
644    }
645
646    pub fn new_method_def<F, FKind>(
647        &self,
648        name: &'static str,
649        f: F,
650        flags: PyMethodFlags,
651        doc: Option<&'static str>,
652    ) -> PyRef<HeapMethodDef>
653    where
654        F: IntoPyNativeFn<FKind>,
655    {
656        let def = PyMethodDef {
657            name,
658            func: Box::leak(Box::new(f.into_func())),
659            flags,
660            doc,
661        };
662        let payload = HeapMethodDef::new(def);
663        PyRef::new_ref(payload, self.types.method_def.to_owned(), None)
664    }
665
666    #[inline]
667    pub fn new_member(
668        &self,
669        name: &str,
670        member_kind: MemberKind,
671        getter: fn(&VirtualMachine, PyObjectRef) -> PyResult,
672        setter: MemberSetterFunc,
673        class: &'static Py<PyType>,
674    ) -> PyRef<PyMemberDescriptor> {
675        let member_def = PyMemberDef {
676            name: name.to_owned(),
677            kind: member_kind,
678            getter: MemberGetter::Getter(getter),
679            setter: MemberSetter::Setter(setter),
680            doc: None,
681        };
682        let member_descriptor = PyMemberDescriptor {
683            common: PyDescriptorOwned {
684                typ: class.to_owned(),
685                name: self.intern_str(name),
686                qualname: PyRwLock::new(None),
687            },
688            member: member_def,
689        };
690        member_descriptor.into_ref(self)
691    }
692
693    pub fn new_readonly_getset<F, T>(
694        &self,
695        name: impl Into<String>,
696        class: &'static Py<PyType>,
697        f: F,
698    ) -> PyRef<PyGetSet>
699    where
700        F: IntoPyGetterFunc<T>,
701    {
702        let name = name.into();
703        let getset = PyGetSet::new(name, class).with_get(f);
704        PyRef::new_ref(getset, self.types.getset_type.to_owned(), None)
705    }
706
707    pub fn new_static_getset<G, S, T, U>(
708        &self,
709        name: impl Into<String>,
710        class: &'static Py<PyType>,
711        g: G,
712        s: S,
713    ) -> PyRef<PyGetSet>
714    where
715        G: IntoPyGetterFunc<T>,
716        S: IntoPySetterFunc<U>,
717    {
718        let name = name.into();
719        let getset = PyGetSet::new(name, class).with_get(g).with_set(s);
720        PyRef::new_ref(getset, self.types.getset_type.to_owned(), None)
721    }
722
723    /// Creates a new `PyGetSet` with a heap type.
724    ///
725    /// # Safety
726    /// In practice, this constructor is safe because a getset is always owned by its `class` type.
727    /// However, it can be broken if used unconventionally.
728    pub unsafe fn new_getset<G, S, T, U>(
729        &self,
730        name: impl Into<String>,
731        class: &Py<PyType>,
732        g: G,
733        s: S,
734    ) -> PyRef<PyGetSet>
735    where
736        G: IntoPyGetterFunc<T>,
737        S: IntoPySetterFunc<U>,
738    {
739        let class = unsafe { &*(class as *const _) };
740        self.new_static_getset(name, class, g, s)
741    }
742
743    pub fn new_base_object(&self, class: PyTypeRef, dict: Option<PyDictRef>) -> PyObjectRef {
744        debug_assert_eq!(
745            class.slots.flags.contains(PyTypeFlags::HAS_DICT),
746            dict.is_some()
747        );
748        PyRef::new_ref(object::PyBaseObject, class, dict).into()
749    }
750
751    pub fn new_code(&self, code: impl code::IntoCodeObject) -> PyRef<PyCode> {
752        let code = code.into_code_object(self);
753        PyRef::new_ref(PyCode::new(code), self.types.code_type.to_owned(), None)
754    }
755}
756
757impl AsRef<Self> for Context {
758    fn as_ref(&self) -> &Self {
759        self
760    }
761}