externref/
signature.rs

1//! Function signatures recorded into a custom section of WASM modules.
2
3use std::str;
4
5use crate::error::{ReadError, ReadErrorKind};
6
7/// Builder for [`BitSlice`]s that can be used in const contexts.
8#[doc(hidden)] // used by macro; not public (yet?)
9#[derive(Debug)]
10pub struct BitSliceBuilder<const BYTES: usize> {
11    bytes: [u8; BYTES],
12    bit_len: usize,
13}
14
15#[doc(hidden)] // not public yet
16impl<const BYTES: usize> BitSliceBuilder<BYTES> {
17    #[must_use]
18    pub const fn with_set_bit(mut self, bit_idx: usize) -> Self {
19        assert!(bit_idx < self.bit_len);
20        self.bytes[bit_idx / 8] |= 1 << (bit_idx % 8);
21        self
22    }
23
24    pub const fn build(&self) -> BitSlice<'_> {
25        BitSlice {
26            bytes: &self.bytes,
27            bit_len: self.bit_len,
28        }
29    }
30}
31
32/// Slice of bits. This type is used to mark [`Resource`](crate::Resource) args
33/// in imported / exported functions.
34// Why invent a new type? Turns out that existing implementations (e.g., `bv` and `bitvec`)
35// cannot be used in const contexts.
36#[derive(Debug, Clone, Copy)]
37#[cfg_attr(test, derive(PartialEq, Eq))]
38pub struct BitSlice<'a> {
39    bytes: &'a [u8],
40    bit_len: usize,
41}
42
43impl BitSlice<'static> {
44    #[doc(hidden)]
45    pub const fn builder<const BYTES: usize>(bit_len: usize) -> BitSliceBuilder<BYTES> {
46        assert!(BYTES > 0);
47        assert!(bit_len > (BYTES - 1) * 8 && bit_len <= BYTES * 8);
48        BitSliceBuilder {
49            bytes: [0_u8; BYTES],
50            bit_len,
51        }
52    }
53}
54
55impl<'a> BitSlice<'a> {
56    /// Returns the number of bits in this slice.
57    pub fn bit_len(&self) -> usize {
58        self.bit_len
59    }
60
61    /// Checks if a bit with the specified 0-based index is set.
62    pub fn is_set(&self, idx: usize) -> bool {
63        if idx > self.bit_len {
64            return false;
65        }
66        let mask = 1 << (idx % 8);
67        self.bytes[idx / 8] & mask > 0
68    }
69
70    /// Iterates over the indexes of set bits in this slice.
71    pub fn set_indices(&self) -> impl Iterator<Item = usize> + '_ {
72        (0..self.bit_len).filter(|&idx| self.is_set(idx))
73    }
74
75    /// Returns the number of set bits in this slice.
76    pub fn count_ones(&self) -> usize {
77        let ones: u32 = self.bytes.iter().copied().map(u8::count_ones).sum();
78        ones as usize
79    }
80
81    fn read_from_section(buffer: &mut &'a [u8], context: &str) -> Result<Self, ReadError> {
82        let bit_len = read_u32(buffer, || format!("length for {context}"))? as usize;
83        let byte_len = (bit_len + 7) / 8;
84        if buffer.len() < byte_len {
85            Err(ReadErrorKind::UnexpectedEof.with_context(context))
86        } else {
87            let bytes = &buffer[..byte_len];
88            *buffer = &buffer[byte_len..];
89            Ok(Self { bytes, bit_len })
90        }
91    }
92}
93
94macro_rules! write_u32 {
95    ($buffer:ident, $value:expr, $pos:expr) => {{
96        let value: u32 = $value;
97        let pos: usize = $pos;
98        $buffer[pos] = (value & 0xff) as u8;
99        $buffer[pos + 1] = ((value >> 8) & 0xff) as u8;
100        $buffer[pos + 2] = ((value >> 16) & 0xff) as u8;
101        $buffer[pos + 3] = ((value >> 24) & 0xff) as u8;
102    }};
103}
104
105fn read_u32(buffer: &mut &[u8], context: impl FnOnce() -> String) -> Result<u32, ReadError> {
106    if buffer.len() < 4 {
107        Err(ReadErrorKind::UnexpectedEof.with_context(context()))
108    } else {
109        let value = u32::from_le_bytes(buffer[..4].try_into().unwrap());
110        *buffer = &buffer[4..];
111        Ok(value)
112    }
113}
114
115fn read_str<'a>(buffer: &mut &'a [u8], context: &str) -> Result<&'a str, ReadError> {
116    let len = read_u32(buffer, || format!("length for {context}"))? as usize;
117    if buffer.len() < len {
118        Err(ReadErrorKind::UnexpectedEof.with_context(context))
119    } else {
120        let string = str::from_utf8(&buffer[..len])
121            .map_err(|err| ReadErrorKind::Utf8(err).with_context(context))?;
122        *buffer = &buffer[len..];
123        Ok(string)
124    }
125}
126
127/// Kind of a function with [`Resource`](crate::Resource) args or return type.
128#[derive(Debug)]
129#[cfg_attr(test, derive(PartialEq, Eq))]
130pub enum FunctionKind<'a> {
131    /// Function exported from a WASM module.
132    Export,
133    /// Function imported to a WASM module from the module with the enclosed name.
134    Import(&'a str),
135}
136
137impl<'a> FunctionKind<'a> {
138    const fn len_in_custom_section(&self) -> usize {
139        match self {
140            Self::Export => 4,
141            Self::Import(module_name) => 4 + module_name.len(),
142        }
143    }
144
145    #[allow(clippy::cast_possible_truncation)] // `TryFrom` cannot be used in const fns
146    const fn write_to_custom_section<const N: usize>(
147        &self,
148        mut buffer: [u8; N],
149    ) -> ([u8; N], usize) {
150        match self {
151            Self::Export => {
152                write_u32!(buffer, u32::MAX, 0);
153                (buffer, 4)
154            }
155
156            Self::Import(module_name) => {
157                write_u32!(buffer, module_name.len() as u32, 0);
158                let mut pos = 4;
159                while pos - 4 < module_name.len() {
160                    buffer[pos] = module_name.as_bytes()[pos - 4];
161                    pos += 1;
162                }
163                (buffer, pos)
164            }
165        }
166    }
167
168    fn read_from_section(buffer: &mut &'a [u8]) -> Result<Self, ReadError> {
169        if buffer.len() >= 4 && buffer[..4] == [0xff; 4] {
170            *buffer = &buffer[4..];
171            Ok(Self::Export)
172        } else {
173            let module_name = read_str(buffer, "module name")?;
174            Ok(Self::Import(module_name))
175        }
176    }
177}
178
179/// Information about a function with [`Resource`](crate::Resource) args or return type.
180///
181/// This information is written to a custom section of a WASM module and is then used
182/// during module [post-processing].
183///
184/// [post-processing]: crate::processor
185#[derive(Debug)]
186#[cfg_attr(test, derive(PartialEq, Eq))]
187pub struct Function<'a> {
188    /// Kind of this function.
189    pub kind: FunctionKind<'a>,
190    /// Name of this function.
191    pub name: &'a str,
192    /// Bit slice marking [`Resource`](crate::Resource) args / return type.
193    pub externrefs: BitSlice<'a>,
194}
195
196impl<'a> Function<'a> {
197    /// Name of a custom section in WASM modules where `Function` declarations are stored.
198    /// `Function`s can be read from this section using [`Self::read_from_section()`].
199    // **NB.** Keep synced with the `declare_function!()` macro below.
200    pub const CUSTOM_SECTION_NAME: &'static str = "__externrefs";
201
202    /// Computes length of a custom section for this function signature.
203    #[doc(hidden)]
204    pub const fn custom_section_len(&self) -> usize {
205        self.kind.len_in_custom_section() + 4 + self.name.len() + 4 + self.externrefs.bytes.len()
206    }
207
208    #[doc(hidden)]
209    #[allow(clippy::cast_possible_truncation)] // `TryFrom` cannot be used in const fns
210    pub const fn custom_section<const N: usize>(&self) -> [u8; N] {
211        debug_assert!(N == self.custom_section_len());
212        let (mut buffer, mut pos) = self.kind.write_to_custom_section([0_u8; N]);
213
214        write_u32!(buffer, self.name.len() as u32, pos);
215        pos += 4;
216        let mut i = 0;
217        while i < self.name.len() {
218            buffer[pos] = self.name.as_bytes()[i];
219            pos += 1;
220            i += 1;
221        }
222
223        write_u32!(buffer, self.externrefs.bit_len as u32, pos);
224        pos += 4;
225        let mut i = 0;
226        while i < self.externrefs.bytes.len() {
227            buffer[pos] = self.externrefs.bytes[i];
228            i += 1;
229            pos += 1;
230        }
231
232        buffer
233    }
234
235    /// Reads function information from a WASM custom section. After reading, the `buffer`
236    /// is advanced to trim the bytes consumed by the parser.
237    ///
238    /// This crate does not provide tools to read custom sections from a WASM module;
239    /// use a library like [`walrus`] or [`wasmparser`] for this purpose.
240    ///
241    /// # Errors
242    ///
243    /// Returns an error if the custom section is malformed.
244    ///
245    /// [`walrus`]: https://docs.rs/walrus/
246    /// [`wasmparser`]: https://docs.rs/wasmparser/
247    pub fn read_from_section(buffer: &mut &'a [u8]) -> Result<Self, ReadError> {
248        let kind = FunctionKind::read_from_section(buffer)?;
249        Ok(Self {
250            kind,
251            name: read_str(buffer, "function name")?,
252            externrefs: BitSlice::read_from_section(buffer, "externref bit slice")?,
253        })
254    }
255}
256
257#[macro_export]
258#[doc(hidden)]
259macro_rules! declare_function {
260    ($signature:expr) => {
261        const _: () = {
262            const FUNCTION: $crate::Function = $signature;
263
264            #[cfg_attr(target_arch = "wasm32", link_section = "__externrefs")]
265            static DATA_SECTION: [u8; FUNCTION.custom_section_len()] = FUNCTION.custom_section();
266        };
267    };
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    #[test]
275    fn function_serialization() {
276        const FUNCTION: Function = Function {
277            kind: FunctionKind::Import("module"),
278            name: "test",
279            externrefs: BitSlice::builder::<1>(3).with_set_bit(1).build(),
280        };
281
282        const SECTION: [u8; FUNCTION.custom_section_len()] = FUNCTION.custom_section();
283
284        assert_eq!(SECTION[..4], [6, 0, 0, 0]); // little-endian module name length
285        assert_eq!(SECTION[4..10], *b"module");
286        assert_eq!(SECTION[10..14], [4, 0, 0, 0]); // little-endian fn name length
287        assert_eq!(SECTION[14..18], *b"test");
288        assert_eq!(SECTION[18..22], [3, 0, 0, 0]); // little-endian bit slice length
289        assert_eq!(SECTION[22], 2); // bit slice
290
291        let mut section_reader = &SECTION as &[u8];
292        let restored_function = Function::read_from_section(&mut section_reader).unwrap();
293        assert_eq!(restored_function, FUNCTION);
294    }
295
296    #[test]
297    fn export_fn_serialization() {
298        const FUNCTION: Function = Function {
299            kind: FunctionKind::Export,
300            name: "test",
301            externrefs: BitSlice::builder::<1>(3).with_set_bit(1).build(),
302        };
303
304        const SECTION: [u8; FUNCTION.custom_section_len()] = FUNCTION.custom_section();
305
306        assert_eq!(SECTION[..4], [0xff, 0xff, 0xff, 0xff]);
307        assert_eq!(SECTION[4..8], [4, 0, 0, 0]); // little-endian fn name length
308        assert_eq!(SECTION[8..12], *b"test");
309
310        let mut section_reader = &SECTION as &[u8];
311        let restored_function = Function::read_from_section(&mut section_reader).unwrap();
312        assert_eq!(restored_function, FUNCTION);
313    }
314}
externref/signature.rs

externref/
signature.rs
