include_bytes_plus/
lib.rs

1//! Improved version of Rust's `include_bytes` macro that allows to reinterpret input as differently array.
2//!
3//! Due to inability to capture current file path in the stable Rust, this macro only accepts paths relative to crate's root.
4
5#![warn(missing_docs)]
6#![cfg_attr(feature = "cargo-clippy", allow(clippy::style))]
7
8extern crate proc_macro;
9
10use proc_macro::TokenStream;
11
12use core::{fmt, mem};
13
14#[cold]
15#[inline(never)]
16fn compile_error<T: core::fmt::Display>(text: T) -> TokenStream {
17    format!("core::compile_error!(\"{}\")", text).parse().unwrap()
18}
19
20enum Primitive {
21    U8,
22    U16,
23    U32,
24    U64,
25    U128,
26    U8LE,
27    U16LE,
28    U32LE,
29    U64LE,
30    U128LE,
31    U8BE,
32    U16BE,
33    U32BE,
34    U64BE,
35    U128BE,
36}
37
38impl Primitive {
39    const fn size(&self) -> usize {
40        match self {
41            Primitive::U8 => mem::size_of::<u8>(),
42            Primitive::U16 => mem::size_of::<u16>(),
43            Primitive::U32 => mem::size_of::<u32>(),
44            Primitive::U64 => mem::size_of::<u64>(),
45            Primitive::U128 => mem::size_of::<u128>(),
46            Primitive::U8LE => mem::size_of::<u8>(),
47            Primitive::U16LE => mem::size_of::<u16>(),
48            Primitive::U32LE => mem::size_of::<u32>(),
49            Primitive::U64LE => mem::size_of::<u64>(),
50            Primitive::U128LE => mem::size_of::<u128>(),
51            Primitive::U8BE => mem::size_of::<u8>(),
52            Primitive::U16BE => mem::size_of::<u16>(),
53            Primitive::U32BE => mem::size_of::<u32>(),
54            Primitive::U64BE => mem::size_of::<u64>(),
55            Primitive::U128BE => mem::size_of::<u128>(),
56        }
57    }
58}
59
60enum Type {
61    Primitive(Primitive),
62    Array(Primitive, usize),
63}
64
65impl Type {
66    ///returns number of bytes written.
67    fn parse(input: &str) -> Result<Self, TokenStream> {
68        match input {
69            "" => Err(compile_error("'as' is missing type")),
70            "u8" => Ok(Type::Primitive(Primitive::U8)),
71            "u16" => Ok(Type::Primitive(Primitive::U16)),
72            "u32" => Ok(Type::Primitive(Primitive::U32)),
73            "u64" => Ok(Type::Primitive(Primitive::U64)),
74            "u128" => Ok(Type::Primitive(Primitive::U128)),
75            "u8le" => Ok(Type::Primitive(Primitive::U8LE)),
76            "u16le" => Ok(Type::Primitive(Primitive::U16LE)),
77            "u32le" => Ok(Type::Primitive(Primitive::U32LE)),
78            "u64le" => Ok(Type::Primitive(Primitive::U64LE)),
79            "u128le" => Ok(Type::Primitive(Primitive::U128LE)),
80            "u8be" => Ok(Type::Primitive(Primitive::U8BE)),
81            "u16be" => Ok(Type::Primitive(Primitive::U16BE)),
82            "u32be" => Ok(Type::Primitive(Primitive::U32BE)),
83            "u64be" => Ok(Type::Primitive(Primitive::U64BE)),
84            "u128be" => Ok(Type::Primitive(Primitive::U128BE)),
85            other => {
86                if let Some(arg) = input.strip_prefix('[') {
87                    if let Some(arg) = arg.strip_suffix(']') {
88                        let mut arg_split = arg.split(';');
89                        let arr_type = arg_split.next().unwrap();
90                        let arr_size = match arg_split.next() {
91                            Some(size) => size,
92                            None => return Err(compile_error(format_args!("'as' array expression '{}' is missing size", other))),
93                        };
94
95                        if let Some(_) = arg_split.next() {
96                            return Err(compile_error(format_args!("'as' array expression '{}' has superfluous ';'", other)));
97                        }
98                        let arr_type = match arr_type {
99                            "u8" => Primitive::U8,
100                            "u16" => Primitive::U16,
101                            "u32" => Primitive::U32,
102                            "u64" => Primitive::U64,
103                            "u128" => Primitive::U128,
104                            "u8le" => Primitive::U8LE,
105                            "u16le" => Primitive::U16LE,
106                            "u32le" => Primitive::U32LE,
107                            "u64le" => Primitive::U64LE,
108                            "u128le" => Primitive::U128LE,
109                            "u8be" => Primitive::U8BE,
110                            "u16be" => Primitive::U16BE,
111                            "u32be" => Primitive::U32BE,
112                            "u64be" => Primitive::U64BE,
113                            "u128be" => Primitive::U128BE,
114                            invalid => return Err(compile_error(format_args!("'as' array expression '{}' has invalid type '{}'", other, invalid))),
115                        };
116                        match arr_size.parse() {
117                            Ok(0) => Err(compile_error(format_args!("'as' array expression '{}' has zero size, which makes no sense", other))),
118                            Ok(arr_size) => Ok(Type::Array(arr_type, arr_size)),
119                            Err(err) => Err(compile_error(format_args!("'as' array expression '{}' has invalid size: {}", other, err))),
120                        }
121                    } else {
122                        Err(compile_error(format_args!("'as' array expression '{}' is missing closing brackets", other)))
123                    }
124                } else {
125                    Err(compile_error(format_args!("'as' specifies unsupported type '{}'", other)))
126                }
127            },
128        }
129    }
130
131    ///returns number of bytes written.
132    fn write_bytes<O: fmt::Write>(&self, out: &mut O, bytes: &[u8]) -> usize {
133        match self {
134            Type::Primitive(primitive) => primitive.write_bytes(out, bytes),
135            Type::Array(primitive, size) => {
136                let mut written = 0;
137                let required_size = primitive.size() * size;
138
139                if required_size > bytes.len() {
140                    return written;
141                }
142
143                for chunk in bytes.chunks_exact(required_size) {
144                    out.write_str("[").expect("to write string");
145                    written += primitive.write_bytes(out, chunk);
146                    out.write_str("],").expect("to write string");
147                }
148
149                written
150            },
151        }
152    }
153}
154
155impl Primitive {
156    ///returns number of bytes written.
157    fn write_bytes<O: fmt::Write>(&self, out: &mut O, bytes: &[u8]) -> usize {
158        match self {
159            Primitive::U8 | Primitive::U8LE | Primitive::U8BE => {
160                for byte in bytes {
161                    core::fmt::write(out, format_args!("0x{:x}u8, ", byte)).expect("To write string");
162                }
163                bytes.len()
164            },
165            Primitive::U16 => {
166                let mut written = 0;
167                for chunk in bytes.chunks_exact(2) {
168                    written += chunk.len();
169                    let byte = u16::from_ne_bytes([chunk[0], chunk[1]]);
170                    core::fmt::write(out, format_args!("0x{:x}u16, ", byte)).expect("To write string");
171                }
172                written
173            },
174            Primitive::U32 => {
175                let mut written = 0;
176                for chunk in bytes.chunks_exact(4) {
177                    written += chunk.len();
178                    let byte = u32::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
179                    core::fmt::write(out, format_args!("0x{:x}u32, ", byte)).expect("To write string");
180                }
181                written
182            }
183            Primitive::U64 => {
184                let mut written = 0;
185                for chunk in bytes.chunks_exact(8) {
186                    written += chunk.len();
187                    let byte = u64::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7]]);
188                    core::fmt::write(out, format_args!("0x{:x}u64, ", byte)).expect("To write string");
189                }
190                written
191            },
192            Primitive::U128 => {
193                let mut written = 0;
194                for chunk in bytes.chunks_exact(16) {
195                    written += chunk.len();
196                    let byte = u128::from_ne_bytes([chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7], chunk[8], chunk[9], chunk[10], chunk[11], chunk[12], chunk[13], chunk[14], chunk[15]]);
197                    core::fmt::write(out, format_args!("0x{:x}u128, ", byte)).expect("To write string");
198                }
199                written
200            },
201            Primitive::U16LE => {
202                let mut written = 0;
203                for chunk in bytes.chunks_exact(2) {
204                    written += chunk.len();
205                    let byte = u16::from_le_bytes([chunk[0], chunk[1]]);
206                    core::fmt::write(out, format_args!("0x{:x}u16, ", byte)).expect("To write string");
207                }
208                written
209            },
210            Primitive::U32LE => {
211                let mut written = 0;
212                for chunk in bytes.chunks_exact(4) {
213                    written += chunk.len();
214                    let byte = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
215                    core::fmt::write(out, format_args!("0x{:x}u32, ", byte)).expect("To write string");
216                }
217                written
218            }
219            Primitive::U64LE => {
220                let mut written = 0;
221                for chunk in bytes.chunks_exact(8) {
222                    written += chunk.len();
223                    let byte = u64::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7]]);
224                    core::fmt::write(out, format_args!("0x{:x}u64, ", byte)).expect("To write string");
225                }
226                written
227            },
228            Primitive::U128LE => {
229                let mut written = 0;
230                for chunk in bytes.chunks_exact(16) {
231                    written += chunk.len();
232                    let byte = u128::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7], chunk[8], chunk[9], chunk[10], chunk[11], chunk[12], chunk[13], chunk[14], chunk[15]]);
233                    core::fmt::write(out, format_args!("0x{:x}u128, ", byte)).expect("To write string");
234                }
235                written
236            },
237            Primitive::U16BE => {
238                let mut written = 0;
239                for chunk in bytes.chunks_exact(2) {
240                    written += chunk.len();
241                    let byte = u16::from_be_bytes([chunk[0], chunk[1]]);
242                    core::fmt::write(out, format_args!("0x{:x}u16, ", byte)).expect("To write string");
243                }
244                written
245            },
246            Primitive::U32BE => {
247                let mut written = 0;
248                for chunk in bytes.chunks_exact(4) {
249                    written += chunk.len();
250                    let byte = u32::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
251                    core::fmt::write(out, format_args!("0x{:x}u32, ", byte)).expect("To write string");
252                }
253                written
254            }
255            Primitive::U64BE => {
256                let mut written = 0;
257                for chunk in bytes.chunks_exact(8) {
258                    written += chunk.len();
259                    let byte = u64::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7]]);
260                    core::fmt::write(out, format_args!("0x{:x}u64, ", byte)).expect("To write string");
261                }
262                written
263            },
264            Primitive::U128BE => {
265                let mut written = 0;
266                for chunk in bytes.chunks_exact(16) {
267                    written += chunk.len();
268                    let byte = u128::from_be_bytes([chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7], chunk[8], chunk[9], chunk[10], chunk[11], chunk[12], chunk[13], chunk[14], chunk[15]]);
269                    core::fmt::write(out, format_args!("0x{:x}u128, ", byte)).expect("To write string");
270                }
271                written
272            },
273        }
274    }
275}
276
277impl fmt::Display for Primitive {
278    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
279        match self {
280            Primitive::U8 => fmt.write_str("u8"),
281            Primitive::U16 => fmt.write_str("u16"),
282            Primitive::U32 => fmt.write_str("u32"),
283            Primitive::U64 => fmt.write_str("u64"),
284            Primitive::U128 => fmt.write_str("u128"),
285            Primitive::U8LE => fmt.write_str("u8le"),
286            Primitive::U16LE => fmt.write_str("u16le"),
287            Primitive::U32LE => fmt.write_str("u32le"),
288            Primitive::U64LE => fmt.write_str("u64le"),
289            Primitive::U128LE => fmt.write_str("u128le"),
290            Primitive::U8BE => fmt.write_str("u8be"),
291            Primitive::U16BE => fmt.write_str("u16be"),
292            Primitive::U32BE => fmt.write_str("u32be"),
293            Primitive::U64BE => fmt.write_str("u64be"),
294            Primitive::U128BE => fmt.write_str("u128be"),
295        }
296    }
297}
298
299impl fmt::Display for Type {
300    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
301        match self {
302            Type::Primitive(primitive) => fmt::Display::fmt(primitive, fmt),
303            Type::Array(primitive, size) => fmt.write_fmt(format_args!("[{}; {}]", primitive, size)),
304        }
305    }
306}
307
308struct Input<'a> {
309    file: &'a str,
310    typ: Type,
311}
312
313impl<'a> Input<'a> {
314    fn parse(input: &'a str) -> Result<Self, TokenStream> {
315        let (file, input) = if let Some(input) = input.strip_prefix('"') {
316            if let Some(end_file_idx) = input.find('"') {
317                (&input[..end_file_idx], &input[end_file_idx+1..])
318            } else {
319                return Err(compile_error("Missing '\"' at the end of file path"));
320            }
321        } else {
322            let mut split = input.split_whitespace();
323            let file = split.next().unwrap();
324            (file, &input[file.len()..])
325        };
326
327        let input = input.trim();
328        let mut split = input.split_whitespace();
329
330        let typ = match split.next() {
331            Some("as") => {
332                let arg = split.fold(String::new(), |mut acc, part| {
333                    acc.push_str(part);
334                    acc
335                });
336                let arg = arg.trim();
337
338                Type::parse(arg)?
339            },
340            Some(other) => return Err(compile_error(format_args!("Unsupported syntax after file name '{}'", other))),
341            None => Type::Primitive(Primitive::U8),
342        };
343
344        Ok(Self {
345            file,
346            typ,
347        })
348    }
349}
350
351#[proc_macro]
352///Includes a file as a reference to a byte array.
353///
354///This macro will yield an expression of type [u8; N] by default with content of file.
355///
356///To reinterpret it as different type add `as <type>` where type can be: `u8`, `u16`, `u32`, `u64` or `u128`.
357///
358///You can use an endianness suffix to change the expected endianness within a file.
359///When endianness is specified the macro shall treat integers within a file correspondingly to the suffix and
360///convert it to the target's native endianness. Allowed suffixes: le and be, eg. `u16le` or `u32be`.
361///
362///Without suffix, bytes are extracted as it is.
363///
364///# NOTE:
365///
366///Due to `Span::source_file` being unstable, the file is searched relative to crate root.
367///
368///# Supported types:
369///
370///- Primitive fixed sized unsigned integers;
371///
372///- Arrays with unsigned integers;
373///
374///# Usage:
375///
376///```
377///use include_bytes_plus::include_bytes;
378///
379///let bytes = include_bytes!("tests/include.in");
380///let bytes_u16 = include_bytes!("tests/include.in" as u16);
381///let bytes_u16_2 = include_bytes!("tests/include with whitespaces.in" as u16);
382///let bytes_u16_3 = include_bytes!("tests/include with whitespaces.in" as [u8; 48]);
383///let bytes_u16_4 = include_bytes!("tests/include with whitespaces.in" as [u16; 12]);
384///let bytes_u16be = include_bytes!("tests/include.in" as u16be);
385///let bytes_u16be_4 = include_bytes!("tests/include with whitespaces.in" as [u16be; 12]);
386///
387///assert_eq!(bytes.len(), bytes_u16.len() * 2);
388///assert_eq!(bytes.len(), bytes_u16_2.len() * 2);
389///assert_eq!(bytes_u16_3.len(), 1);
390///assert_eq!(bytes_u16_3[0].len(), 48);
391///assert_eq!(bytes_u16_4.len(), 2);
392///assert_eq!(bytes_u16_4[0].len(), 12);
393///assert_eq!(bytes_u16_4[1].len(), 12);
394///assert_eq!(bytes_u16be.len(), bytes_u16.len());
395///assert_eq!(bytes_u16be_4.len(), 2);
396///assert_eq!(bytes_u16be_4[0].len(), 12);
397///assert_eq!(bytes_u16be_4[1].len(), 12);
398///#[cfg(target_endian = "little")]
399///assert_ne!(bytes_u16_4, bytes_u16be_4);
400///```
401///
402///# Debugging timings:
403///
404///Set env variable `RUST_INCLUDE_BYTES_LOG=1` to enable logging of each parsed file statistics
405pub fn include_bytes(input: TokenStream) -> TokenStream {
406    let is_log = match std::env::var("RUST_INCLUDE_BYTES_LOG") {
407        Ok(res) => match res.as_str() {
408            "1" | "true" => true,
409            _ => false,
410        },
411        _ => false,
412    };
413
414    let now = std::time::Instant::now();
415
416    let input = input.to_string();
417    let input = input.trim();
418
419    let args = match Input::parse(input) {
420        Ok(args) => args,
421        Err(error) => return error,
422    };
423
424    if args.file.is_empty() {
425        return compile_error("Empty file name");
426    }
427
428    let mut file = match std::fs::File::open(args.file) {
429        Ok(file) => file,
430        Err(error) => return compile_error(format_args!("{}: Cannot open file: {}", args.file, error)),
431    };
432
433    let mut cursor = 0;
434    let mut file_len = 0;
435    let mut buf = [0u8; 4096];
436    let mut result = "[".to_owned();
437
438    loop {
439        match std::io::Read::read(&mut file, &mut buf[cursor..]) {
440            Ok(0) => {
441                result.push(']');
442                if cursor != 0 {
443                    return compile_error(format_args!("File input with size {}b cannot be reinterpret as {}", file_len, args.typ));
444                }
445                break;
446            },
447            Ok(size) => {
448                file_len += size;
449                let buf_len = cursor + size;
450                let written = args.typ.write_bytes(&mut result, &buf[..buf_len]);
451
452                if written > 0 {
453                    unsafe {
454                        core::ptr::copy(buf.as_ptr().add(written), buf.as_mut_ptr(), buf_len - written);
455                    }
456                    cursor = buf_len - written;
457                } else {
458                    //not enough data to write expression
459                    //wait another read.
460                    cursor = buf_len;
461                }
462            },
463            Err(error) => {
464                return compile_error(format_args!("{}: Error reading file: {}", args.file, error))
465            },
466        }
467    }
468
469    if is_log {
470        let elapsed = now.elapsed();
471        let secs = elapsed.as_secs();
472        let ms = elapsed.subsec_millis();
473
474        if secs > 0 {
475            println!("{}: parsed {}b in {}.{} seconds", args.file, file_len, secs, ms);
476        } else {
477            println!("{}: parsed {}b in {} ms", args.file, file_len, ms);
478        }
479    }
480
481    result.parse().expect("To parse")
482}