Rust solution for /r/dailyprogrammer challenge #317 [intermediate]

challenge_317_intermediate.rs

#![feature(iter_map_while)]

#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
#[repr(align(2))]
pub struct Element([u8; 2]);

impl Element {
    fn as_str(&self) -> &str {
        let s = if self.0[1] == 0 { &self.0[..1] } else { &self.0 };
        unsafe { std::str::from_utf8_unchecked(s) }
    }
}

pub type Compound = flat_map::FlatMap<Element, i32>;

mod parsers {
    use super::{Compound, Element};
    use nom::{bytes::complete::tag, character::complete::digit1,
              combinator::{map, map_opt, opt, verify},
              branch::alt, multi::many1, sequence::{delimited, pair}};

    mod ast {
        pub type Formula   = Vec<Component>;
        pub type Allotrope = (super::Element, i32);
        pub type Compound  = (Formula, i32);
        pub enum Component { Allo(Allotrope), Comp(Compound) }
    }

    type Input<'a> = &'a [u8];
    type Result<'a, T> = nom::IResult<Input<'a>, T>;

    pub fn parse_formula(input: &str) -> Option<Compound> {
        fn visit(mut comp: Compound, component: ast::Component) -> Compound {
            let mut inc_count = |elem, c| *comp.entry(elem).or_insert(0) += c;
            match component {
                ast::Component::Allo((elem, count))            => inc_count(elem, count),
                ast::Component::Comp((components, multiplier)) => components.into_iter()
                                                                      .fold(Compound::new(), visit)
                                                                      .into_iter()
                                                                      .for_each(|(elem, count)| inc_count(elem, count * multiplier))
            }
            comp
        }

        if let Ok(([], raw)) = formula(input.as_bytes()) {
            Some(visit(Compound::new(), ast::Component::Comp((raw, 1))))
        } else {
            None
        }
    }

    fn formula(i: Input) -> Result<ast::Formula> {
        many1(
            alt((
                map(compound,  ast::Component::Comp),
                map(allotrope, ast::Component::Allo)
            ))
        )(i)
    }

    fn allotrope(i: Input) -> Result<ast::Allotrope> {
        let (i, upper) =     take_single_if(u8::is_ascii_uppercase)(i)?;
        let (i, lower) = opt(take_single_if(u8::is_ascii_lowercase))(i)?;
        let (i, count) = opt(extract_u16)(i)?;
        Ok((i, (Element([upper, lower.unwrap_or_default()]), count.unwrap_or(1).into())))
    }

    fn compound(i: Input) -> Result<ast::Compound> {
        pair(
            delimited(tag(b"("), formula, tag(b")")),
            map(extract_u16, i32::from)
        )(i)
    }

    fn extract_u16(i: Input) -> Result<u16> {
        use nom::ParseTo;
        map_opt(digit1, |n: Input| n.parse_to())(i)
    }

    fn take_single_if<'a, F>(pred: F) -> impl Fn(Input<'a>) -> Result<'a, u8>
    where F: Fn(&u8) -> bool {
        verify(
            |i: Input| if let Some((&head, tail)) = i.split_first() {
                           Ok((tail, head))
                       } else {
                           Err(nom::Err::Error((i, nom::error::ErrorKind::Eof)))
                       },
            pred
        )
    }
}
pub use parsers::parse_formula;

fn pnp(input: &str) {
    print!("{}", input);
    if let Some(comp) = parse_formula(input) {
        println!();
        comp.into_iter().for_each(|(elem, count)| println!("{:>6}: {}", elem.as_str(), count));
    } else {
        println!(" - failed to parse");
    }
    println!();
}

fn main() {
    if cfg!(debug_assertions) {
        [
            "C6H12O6",
            "CCl2F2",
            "NaHCO3",
            "C4H8(OH)2",
            "PbCl(NH3)2(COOH)2",
            "PbCl(NH3(H2O)4)2",
            "Cl((NaH)2CO3)2",
            "PbCl(NH3)2((CO)2OH)2",
            "(C3(H2O)3)2",
            "FPb((NO4)2(COOH)3)4"
        ].iter().for_each(|&input| pnp(input));
    } else {
        use std::io::BufRead;
        std::io::stdin().lock()
            .lines()
            .map_while(|line_res| line_res.ok().filter(|line| !line.is_empty()))
            .for_each(|input| pnp(&input));
    }
}

Updated from nom v3 to v5