rustpython_compiler_core/
frozen.rs

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

rustpython_compiler_core/
frozen.rs
