rustpython_compiler_core/
frozen.rs

1use crate::bytecode::*;
2use crate::marshal::{self, Read, ReadBorrowed, Write};
3
4/// A frozen module. Holds a frozen code object and whether it is part of a package
5#[derive(Copy, Clone)]
6pub struct FrozenModule<B = &'static [u8]> {
7    pub code: FrozenCodeObject<B>,
8    pub package: bool,
9}
10
11#[derive(Copy, Clone)]
12pub struct FrozenCodeObject<B> {
13    pub bytes: B,
14}
15
16impl<B: AsRef<[u8]>> FrozenCodeObject<B> {
17    /// Decode a frozen code object
18    #[inline]
19    pub fn decode<Bag: AsBag>(&self, bag: Bag) -> CodeObject<<Bag::Bag as ConstantBag>::Constant> {
20        Self::_decode(self.bytes.as_ref(), bag.as_bag())
21    }
22    fn _decode<Bag: ConstantBag>(data: &[u8], bag: Bag) -> CodeObject<Bag::Constant> {
23        let decompressed = lz4_flex::decompress_size_prepended(data)
24            .expect("deserialize frozen CodeObject failed");
25        marshal::deserialize_code(&mut &decompressed[..], bag)
26            .expect("deserializing frozen CodeObject failed")
27    }
28}
29
30impl FrozenCodeObject<Vec<u8>> {
31    pub fn encode<C: Constant>(code: &CodeObject<C>) -> Self {
32        let mut data = Vec::new();
33        marshal::serialize_code(&mut data, code);
34        let bytes = lz4_flex::compress_prepend_size(&data);
35        FrozenCodeObject { bytes }
36    }
37}
38
39#[repr(transparent)]
40pub struct FrozenLib<B: ?Sized = [u8]> {
41    pub bytes: B,
42}
43
44impl<B: AsRef<[u8]> + ?Sized> FrozenLib<B> {
45    pub const fn from_ref(b: &B) -> &FrozenLib<B> {
46        unsafe { &*(b as *const B as *const FrozenLib<B>) }
47    }
48
49    /// Decode a library to a iterable of frozen modules
50    pub fn decode(&self) -> FrozenModulesIter<'_> {
51        let mut data = self.bytes.as_ref();
52        let remaining = data.read_u32().unwrap();
53        FrozenModulesIter { remaining, data }
54    }
55}
56
57impl<'a, B: AsRef<[u8]> + ?Sized> IntoIterator for &'a FrozenLib<B> {
58    type Item = (&'a str, FrozenModule<&'a [u8]>);
59    type IntoIter = FrozenModulesIter<'a>;
60    fn into_iter(self) -> Self::IntoIter {
61        self.decode()
62    }
63}
64
65pub struct FrozenModulesIter<'a> {
66    remaining: u32,
67    data: &'a [u8],
68}
69
70impl<'a> Iterator for FrozenModulesIter<'a> {
71    type Item = (&'a str, FrozenModule<&'a [u8]>);
72
73    fn next(&mut self) -> Option<Self::Item> {
74        if self.remaining > 0 {
75            let entry = read_entry(&mut self.data).unwrap();
76            self.remaining -= 1;
77            Some(entry)
78        } else {
79            None
80        }
81    }
82
83    fn size_hint(&self) -> (usize, Option<usize>) {
84        (self.remaining as usize, Some(self.remaining as usize))
85    }
86}
87impl ExactSizeIterator for FrozenModulesIter<'_> {}
88
89fn read_entry<'a>(
90    rdr: &mut &'a [u8],
91) -> Result<(&'a str, FrozenModule<&'a [u8]>), marshal::MarshalError> {
92    let len = rdr.read_u32()?;
93    let name = rdr.read_str_borrow(len)?;
94    let len = rdr.read_u32()?;
95    let code_slice = rdr.read_slice_borrow(len)?;
96    let code = FrozenCodeObject { bytes: code_slice };
97    let package = rdr.read_u8()? != 0;
98    Ok((name, FrozenModule { code, package }))
99}
100
101impl FrozenLib<Vec<u8>> {
102    /// Encode the given iterator of frozen modules into a compressed vector of bytes
103    pub fn encode<'a, I, B: AsRef<[u8]>>(lib: I) -> FrozenLib<Vec<u8>>
104    where
105        I: IntoIterator<Item = (&'a str, FrozenModule<B>)>,
106        I::IntoIter: ExactSizeIterator + Clone,
107    {
108        let iter = lib.into_iter();
109        let mut bytes = Vec::new();
110        write_lib(&mut bytes, iter);
111        Self { bytes }
112    }
113}
114
115fn write_lib<'a, B: AsRef<[u8]>>(
116    buf: &mut Vec<u8>,
117    lib: impl ExactSizeIterator<Item = (&'a str, FrozenModule<B>)>,
118) {
119    marshal::write_len(buf, lib.len());
120    for (name, module) in lib {
121        write_entry(buf, name, module);
122    }
123}
124
125fn write_entry(buf: &mut Vec<u8>, name: &str, module: FrozenModule<impl AsRef<[u8]>>) {
126    marshal::write_vec(buf, name.as_bytes());
127    marshal::write_vec(buf, module.code.bytes.as_ref());
128    buf.write_u8(module.package as u8);
129}
rustpython_compiler_core/frozen.rs

rustpython_compiler_core/
frozen.rs
