Add byte_map! parser macro

Author: ictrobot

crates/utils/src/parser/macros.rs

@@ -1,3 +1,77 @@
+/// Macro to define a parser which consumes a single byte and maps it using a lookup table.
+///
+/// This macro is a wrapper around [`parser::byte_lut`](crate::parser::byte_lut) to allow defining
+/// the lookup table using a match-like syntax. Each expression must be const and evaluate to a
+/// value of the same copy type.
+///
+/// # Examples
+/// ```
+/// # use utils::parser::{Parser, self};
+/// let parser = parser::byte_map!(
+///     b'#' => true,
+///     b'.' | b'S' => false,
+/// );
+/// assert_eq!(parser.parse(b"#.S##"), Ok((true, &b".S##"[..])));
+/// assert_eq!(parser.parse(b".S##"), Ok((false, &b"S##"[..])));
+/// assert_eq!(parser.parse(b"S##"), Ok((false, &b"##"[..])));
+///
+/// let (err, remaining) = parser.parse(b"abc").unwrap_err();
+/// assert_eq!(err.to_string(), "expected one of '#', '.', 'S'");
+/// assert_eq!(remaining, &b"abc"[..]);
+/// ```
+#[macro_export]
+macro_rules! parser_byte_map {
+    (
+        $($($l:literal)|+ => $e:expr),+$(,)?
+    ) => {{
+        // `let _: u8 = $l` ensures $l is used in the repetition and also ensures all the literals
+        // are byte literals
+        const COUNT: usize = 0usize $($(+ {let _: u8 = $l; 1usize})+)+;
+        const LEN: usize = 14 + 5 * COUNT;
+        const {
+            assert!(COUNT >= 2, "at least two literals must be provided");
+        }
+
+        // Once concat_bytes! is stabilized this error message can be created in the macro similar
+        // to parser_literal_map!
+        const ERROR: [u8; LEN] = {
+            let mut result = [0u8; LEN];
+            let (prefix, vals) = result.split_at_mut(16);
+            prefix.copy_from_slice(b"expected one of ");
+
+            let mut i = 0;
+            let literals = [$($($l),+),+];
+            while i < COUNT {
+                vals[i * 5] = b'\'';
+                vals[i * 5 + 1] = literals[i];
+                vals[i * 5 + 2] = b'\'';
+                if i + 1 < COUNT {
+                    vals[i * 5 + 3] = b',';
+                    vals[i * 5 + 4] = b' ';
+                }
+                i += 1;
+            }
+
+            result
+        };
+
+        $crate::parser::byte_lut(&const {
+            // Don't use a const item for the lut to avoid naming the value type
+            let mut lut = [None; 256];
+            $($(
+                assert!(lut[$l as usize].is_none(), "duplicate literal");
+                lut[$l as usize] = Some($e);
+            )+)+
+            lut
+        }, const {
+            match str::from_utf8(&ERROR) {
+                Ok(v) => v,
+                Err(_) => panic!("one or more of the provided literals is invalid unicode"),
+            }
+        })
+    }};
+}
+
 /// Macro to define a parser for one or more string literals, mapping the results.
 ///
 /// This is a replacement for

crates/utils/src/parser/mod.rs

@@ -15,8 +15,9 @@ pub use error::ParseError;
 pub use iterator::{ParserIterator, ParserMatchesIterator};
 pub use number::{i8, i16, i32, i64, i128, number_range, u8, u16, u32, u64, u128};
 pub use one_of::one_of;
-pub use simple::{byte, byte_range, constant, eof, eol, noop, take_while, take_while1};
+pub use simple::{byte, byte_lut, byte_range, constant, eof, eol, noop, take_while, take_while1};
 
+pub use crate::parser_byte_map as byte_map;
 pub use crate::parser_literal_map as literal_map;
 pub use crate::parser_parsable_enum as parsable_enum;
 pub use crate::parser_parse_tree as parse_tree;

crates/utils/src/parser/simple.rs

@@ -39,6 +39,56 @@ pub fn byte() -> Byte {
     Byte()
 }
 
+#[derive(Copy, Clone)]
+pub struct ByteLut<'a, O> {
+    lut: &'a [Option<O>; 256],
+    error: &'static str,
+}
+impl<'i, O: Copy> Parser<'i> for ByteLut<'_, O> {
+    type Output = O;
+    type Then<T: Parser<'i>> = Then2<Self, T>;
+
+    #[inline]
+    fn parse(&self, input: &'i [u8]) -> ParseResult<'i, Self::Output> {
+        if let [byte, remaining @ ..] = input
+            && let Some(output) = self.lut[*byte as usize]
+        {
+            Ok((output, remaining))
+        } else {
+            Err((ParseError::Custom(self.error), input))
+        }
+    }
+}
+
+/// Parser that consumes a single byte and maps it using a lookup table.
+///
+/// Equivalent to `parser::byte().map_res(|b| LOOKUP[b as usize].ok_or("expected ..."))`, which is
+/// usually faster than an equivalent match statement in the closure.
+///
+/// See also [`parser::byte_map!`](crate::parser::byte_map!) which wraps this function, allowing a
+/// match-like syntax to be used to define the lookup table.
+///
+/// # Examples
+/// ```
+/// # use utils::parser::{self, Parser, ParseError};
+/// const LOOKUP: [Option<bool>; 256] = {
+///     let mut x = [None; 256];
+///     x['#' as usize] = Some(true);
+///     x['.' as usize] = Some(false);
+///     x
+/// };
+///
+/// let parser = parser::byte_lut(&LOOKUP, "expected '#' or '.'");
+/// assert_eq!(parser.parse(b"#..##"), Ok((true, &b"..##"[..])));
+/// assert_eq!(parser.parse(b"..##"), Ok((false, &b".##"[..])));
+/// assert_eq!(parser.parse(b"abc"), Err((ParseError::Custom("expected '#' or '.'"), &b"abc"[..])));
+/// ```
+#[inline]
+#[must_use]
+pub fn byte_lut<'a, T: Copy>(lut: &'a [Option<T>; 256], error: &'static str) -> ByteLut<'a, T> {
+    ByteLut { lut, error }
+}
+
 #[derive(Copy, Clone)]
 pub struct ByteRange {
     min: u8,

crates/year2024/src/day19.rs

@@ -24,23 +24,18 @@ struct TrieNode {
     is_terminal: bool,
 }
 
-const LETTER_LOOKUP: [Option<Stripe>; 256] = {
-    let mut x = [None; 256];
-    x[b'w' as usize] = Some(Stripe::W);
-    x[b'u' as usize] = Some(Stripe::U);
-    x[b'b' as usize] = Some(Stripe::B);
-    x[b'r' as usize] = Some(Stripe::R);
-    x[b'g' as usize] = Some(Stripe::G);
-    x
-};
-
 const MAX_PATTERN_LENGTH: usize = 8;
 const MAX_DESIGN_LENGTH: usize = 64;
 
 impl Day19 {
     pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
-        let letter = parser::byte()
-            .map_res(|b| LETTER_LOOKUP[b as usize].ok_or("expected 'w', 'u', 'b', 'r' or 'g'"));
+        let letter = parser::byte_map!(
+            b'w' => Stripe::W,
+            b'u' => Stripe::U,
+            b'b' => Stripe::B,
+            b'r' => Stripe::R,
+            b'g' => Stripe::G,
+        );
 
         let Some((patterns, designs)) = input.split_once("\n\n") else {
             return Err(InputError::new(input, 0, "expected patterns and designs"));

crates/year2024/src/day25.rs

@@ -7,20 +7,12 @@ pub struct Day25 {
     keys: Vec<u32>,
 }
 
-const LOOKUP: [Option<bool>; 256] = {
-    let mut x = [None; 256];
-    x['#' as usize] = Some(true);
-    x['.' as usize] = Some(false);
-    x
-};
-
 impl Day25 {
     pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
         let mut locks = Vec::with_capacity(250);
         let mut keys = Vec::with_capacity(250);
 
-        for item in parser::byte()
-            .map_res(|b| LOOKUP[b as usize].ok_or("expected '.' or '#'"))
+        for item in parser::byte_map!(b'#' => true, b'.' => false)
             .repeat_n::<5, _>(parser::noop())
             .repeat_n::<7, _>(parser::eol())
             .with_consumed()