1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use crate::bytecode::*;
use crate::marshal::{self, Read, ReadBorrowed, Write};

/// A frozen module. Holds a frozen code object and whether it is part of a package
#[derive(Copy, Clone)]
pub struct FrozenModule<B = &'static [u8]> {
    pub code: FrozenCodeObject<B>,
    pub package: bool,
}

#[derive(Copy, Clone)]
pub struct FrozenCodeObject<B> {
    pub bytes: B,
}

impl<B: AsRef<[u8]>> FrozenCodeObject<B> {
    /// Decode a frozen code object
    #[inline]
    pub fn decode<Bag: AsBag>(&self, bag: Bag) -> CodeObject<<Bag::Bag as ConstantBag>::Constant> {
        Self::_decode(self.bytes.as_ref(), bag.as_bag())
    }
    fn _decode<Bag: ConstantBag>(data: &[u8], bag: Bag) -> CodeObject<Bag::Constant> {
        let decompressed = lz4_flex::decompress_size_prepended(data)
            .expect("deserialize frozen CodeObject failed");
        marshal::deserialize_code(&mut &decompressed[..], bag)
            .expect("deserializing frozen CodeObject failed")
    }
}

impl FrozenCodeObject<Vec<u8>> {
    pub fn encode<C: Constant>(code: &CodeObject<C>) -> Self {
        let mut data = Vec::new();
        marshal::serialize_code(&mut data, code);
        let bytes = lz4_flex::compress_prepend_size(&data);
        FrozenCodeObject { bytes }
    }
}

#[repr(transparent)]
pub struct FrozenLib<B: ?Sized = [u8]> {
    pub bytes: B,
}

impl<B: AsRef<[u8]> + ?Sized> FrozenLib<B> {
    pub const fn from_ref(b: &B) -> &FrozenLib<B> {
        unsafe { &*(b as *const B as *const FrozenLib<B>) }
    }

    /// Decode a library to a iterable of frozen modules
    pub fn decode(&self) -> FrozenModulesIter<'_> {
        let mut data = self.bytes.as_ref();
        let remaining = data.read_u32().unwrap();
        FrozenModulesIter { remaining, data }
    }
}

impl<'a, B: AsRef<[u8]> + ?Sized> IntoIterator for &'a FrozenLib<B> {
    type Item = (&'a str, FrozenModule<&'a [u8]>);
    type IntoIter = FrozenModulesIter<'a>;
    fn into_iter(self) -> Self::IntoIter {
        self.decode()
    }
}

pub struct FrozenModulesIter<'a> {
    remaining: u32,
    data: &'a [u8],
}

impl<'a> Iterator for FrozenModulesIter<'a> {
    type Item = (&'a str, FrozenModule<&'a [u8]>);

    fn next(&mut self) -> Option<Self::Item> {
        if self.remaining > 0 {
            let entry = read_entry(&mut self.data).unwrap();
            self.remaining -= 1;
            Some(entry)
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.remaining as usize, Some(self.remaining as usize))
    }
}
impl ExactSizeIterator for FrozenModulesIter<'_> {}

fn read_entry<'a>(
    rdr: &mut &'a [u8],
) -> Result<(&'a str, FrozenModule<&'a [u8]>), marshal::MarshalError> {
    let len = rdr.read_u32()?;
    let name = rdr.read_str_borrow(len)?;
    let len = rdr.read_u32()?;
    let code_slice = rdr.read_slice_borrow(len)?;
    let code = FrozenCodeObject { bytes: code_slice };
    let package = rdr.read_u8()? != 0;
    Ok((name, FrozenModule { code, package }))
}

impl FrozenLib<Vec<u8>> {
    /// Encode the given iterator of frozen modules into a compressed vector of bytes
    pub fn encode<'a, I, B: AsRef<[u8]>>(lib: I) -> FrozenLib<Vec<u8>>
    where
        I: IntoIterator<Item = (&'a str, FrozenModule<B>)>,
        I::IntoIter: ExactSizeIterator + Clone,
    {
        let iter = lib.into_iter();
        let mut bytes = Vec::new();
        write_lib(&mut bytes, iter);
        Self { bytes }
    }
}

fn write_lib<'a, B: AsRef<[u8]>>(
    buf: &mut Vec<u8>,
    lib: impl ExactSizeIterator<Item = (&'a str, FrozenModule<B>)>,
) {
    marshal::write_len(buf, lib.len());
    for (name, module) in lib {
        write_entry(buf, name, module);
    }
}

fn write_entry(buf: &mut Vec<u8>, name: &str, module: FrozenModule<impl AsRef<[u8]>>) {
    marshal::write_vec(buf, name.as_bytes());
    marshal::write_vec(buf, module.code.bytes.as_ref());
    buf.write_u8(module.package as u8);
}