soroban_wasmi/module/
custom_section.rs

1use core::slice;
2use std::{str, vec::Vec};
3
4/// Wasm custom sections.
5#[derive(Default, Debug)]
6pub struct CustomSections {
7    inner: CustomSectionsInner,
8}
9
10impl CustomSections {
11    /// Returns an iterator over the [`CustomSection`]s stored in `self`.
12    #[inline]
13    pub fn iter(&self) -> CustomSectionsIter {
14        self.inner.iter()
15    }
16}
17
18/// A builder for [`CustomSections`].
19#[derive(Default, Debug)]
20pub struct CustomSectionsBuilder {
21    inner: CustomSectionsInner,
22}
23
24impl CustomSectionsBuilder {
25    /// Pushes a new custom section segment to the [`CustomSectionsBuilder`].
26    #[inline]
27    pub fn push(&mut self, name: &str, data: &[u8]) {
28        self.inner.push(name, data);
29    }
30
31    /// Finalize construction of the [`CustomSections`].
32    #[inline]
33    pub fn finish(self) -> CustomSections {
34        CustomSections { inner: self.inner }
35    }
36}
37
38/// Internal representation of [`CustomSections`].
39#[derive(Debug, Default)]
40pub struct CustomSectionsInner {
41    /// The name and data lengths of each Wasm custom section.
42    items: Vec<CustomSectionInner>,
43    /// The combined name and data of all Wasm custom sections.
44    names_and_data: Vec<u8>,
45}
46
47/// Internal representation of a Wasm [`CustomSection`].
48#[derive(Debug, Copy, Clone)]
49pub struct CustomSectionInner {
50    /// The length in bytes of the Wasm custom section name.
51    len_name: usize,
52    /// The length in bytes of the Wasm custom section data.
53    len_data: usize,
54}
55
56impl CustomSectionsInner {
57    /// Pushes a new custom section segment to the [`CustomSectionsBuilder`].
58    #[inline]
59    pub fn push(&mut self, name: &str, data: &[u8]) {
60        let name_bytes = name.as_bytes();
61        self.names_and_data.extend_from_slice(name_bytes);
62        self.names_and_data.extend_from_slice(data);
63        self.items.push(CustomSectionInner {
64            len_name: name_bytes.len(),
65            len_data: data.len(),
66        })
67    }
68
69    /// Returns an iterator over the [`CustomSection`]s stored in `self`.
70    #[inline]
71    pub fn iter(&self) -> CustomSectionsIter {
72        CustomSectionsIter {
73            items: self.items.iter(),
74            names_and_data: &self.names_and_data[..],
75        }
76    }
77}
78
79/// A Wasm custom section.
80#[derive(Debug)]
81pub struct CustomSection<'a> {
82    /// The name of the custom section.
83    name: &'a str,
84    /// The undecoded data of the custom section.
85    data: &'a [u8],
86}
87
88impl<'a> CustomSection<'a> {
89    /// Returns the name or identifier of the [`CustomSection`].
90    #[inline]
91    pub fn name(&self) -> &'a str {
92        self.name
93    }
94
95    /// Returns a shared reference to the data of the [`CustomSection`].
96    #[inline]
97    pub fn data(&self) -> &'a [u8] {
98        self.data
99    }
100}
101
102/// An iterator over the custom sections of a Wasm module.
103#[derive(Debug)]
104pub struct CustomSectionsIter<'a> {
105    items: slice::Iter<'a, CustomSectionInner>,
106    names_and_data: &'a [u8],
107}
108
109impl<'a> Iterator for CustomSectionsIter<'a> {
110    type Item = CustomSection<'a>;
111
112    #[inline]
113    fn size_hint(&self) -> (usize, Option<usize>) {
114        self.items.size_hint()
115    }
116
117    #[inline]
118    fn next(&mut self) -> Option<Self::Item> {
119        let item = self.items.next()?;
120        let names_and_data = self.names_and_data;
121        let (name, names_and_data) = names_and_data.split_at(item.len_name);
122        let (data, names_and_data) = names_and_data.split_at(item.len_data);
123        self.names_and_data = names_and_data;
124        // Safety: We encoded this part of the data buffer from the bytes of a string previously.
125        let name = unsafe { str::from_utf8_unchecked(name) };
126        Some(CustomSection { name, data })
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn it_works() {
136        let mut builder = CustomSectionsBuilder::default();
137        builder.push("A", b"first");
138        builder.push("B", b"second");
139        builder.push("C", b"third");
140        builder.push("", b"fourth"); // empty name
141        builder.push("E", &[]); // empty data
142        let custom_sections = builder.finish();
143        let mut iter = custom_sections.iter();
144        assert_eq!(
145            iter.next().map(|s| (s.name(), s.data())),
146            Some(("A", &b"first"[..]))
147        );
148        assert_eq!(
149            iter.next().map(|s| (s.name(), s.data())),
150            Some(("B", &b"second"[..]))
151        );
152        assert_eq!(
153            iter.next().map(|s| (s.name(), s.data())),
154            Some(("C", &b"third"[..]))
155        );
156        assert_eq!(
157            iter.next().map(|s| (s.name(), s.data())),
158            Some(("", &b"fourth"[..]))
159        );
160        assert_eq!(
161            iter.next().map(|s| (s.name(), s.data())),
162            Some(("E", &b""[..]))
163        );
164        assert_eq!(iter.next().map(|s| (s.name(), s.data())), None);
165    }
166}