p16.rs

use crate::utils::{read_input, Result};

#[derive(Clone, Copy)]
pub enum OperatorType {
    Sum,
    Product,
    Min,
    Max,
    Gt,
    Lt,
    Eq,
}

impl OperatorType {
    pub fn from_num(n: u64) -> Self {
        match n {
            0 => OperatorType::Sum,
            1 => OperatorType::Product,
            2 => OperatorType::Min,
            3 => OperatorType::Max,
            5 => OperatorType::Gt,
            6 => OperatorType::Lt,
            7 => OperatorType::Eq,
            _ => unreachable!(),
        }
    }
}

#[derive(Clone)]
pub enum PacketContent {
    Literal(u64),
    Operator {
        operator: OperatorType,
        inner: Vec<Packet>,
    },
}

impl PacketContent {
    pub fn eval(&self) -> u64 {
        match self {
            &Self::Literal(v) => v,
            Self::Operator { operator, inner, .. } => match operator {
                OperatorType::Sum => Self::eval_inner(inner).sum(),
                OperatorType::Product => Self::eval_inner(inner).product(),
                OperatorType::Min => Self::eval_inner(inner).min().unwrap_or(0),
                OperatorType::Max => Self::eval_inner(inner).max().unwrap_or(0),
                OperatorType::Gt => Self::cmp_op(inner, u64::gt),
                OperatorType::Lt => Self::cmp_op(inner, u64::lt),
                OperatorType::Eq => Self::cmp_op(inner, u64::eq),
            }
        }
    }

    fn cmp_op(inner: &[Packet], op: fn(&u64, &u64) -> bool) -> u64 {
        inner
            .get(0)
            .and_then(|a|
                inner
                    .get(1)
                    .map(|b| if op(&a.eval(), &b.eval()) { 1 } else { 0 }))
            .unwrap_or(0)
    }

    fn eval_inner(inner: &[Packet]) -> impl Iterator<Item=u64> + '_ {
        inner.iter().map(|p| p.eval())
    }
}

#[derive(Clone)]
pub struct Packet {
    version: u64,
    content: PacketContent,
}

impl Packet {
    fn from_hex(input: &str) -> Result<Self> {
        let binary = utils::to_binary(input);

        // otherwise borrow checker hates me
        if let Ok((_, packet)) = parser::packet(&binary) {
            Ok(packet)
        } else {
            Err("invalid input".into())
        }
    }

    fn sum_version(&self) -> u64 {
        self.version + match &self.content {
            PacketContent::Operator { inner, .. } => inner
                .into_iter()
                .map(|p| p.sum_version())
                .sum(),
            _ => 0,
        }
    }

    fn eval(&self) -> u64 {
        self.content.eval()
    }
}

mod utils {
    use itertools::Itertools;

    pub fn to_binary(input: &str) -> String {
        input
            .trim()
            .chars()
            .map(|c| match c {
                '0' => "0000",
                '1' => "0001",
                '2' => "0010",
                '3' => "0011",
                '4' => "0100",
                '5' => "0101",
                '6' => "0110",
                '7' => "0111",
                '8' => "1000",
                '9' => "1001",
                'A' => "1010",
                'B' => "1011",
                'C' => "1100",
                'D' => "1101",
                'E' => "1110",
                'F' => "1111",
                _ => "",
            })
            .join("")
    }
}

mod parser {
    use itertools::Itertools;
    use nom::bytes::complete::take;
    use nom::combinator::map_res;
    use nom::IResult;
    use nom::multi::{count, many0};
    use crate::p_16::{OperatorType, Packet, PacketContent};

    fn from_binary(input: &str) -> Result<u64, std::num::ParseIntError> {
        u64::from_str_radix(input, 2)
    }

    fn from_owned_binary(input: String) -> Result<u64, std::num::ParseIntError> {
        u64::from_str_radix(input.as_str(), 2)
    }

    fn n_bin(input: &str, n: usize) -> IResult<&str, u64> {
        map_res(
            take(n),
            from_binary,
        )(input)
    }

    fn content_literal_str(input: &str) -> IResult<&str, String> {
        let mut input = input;
        let mut parts = vec![];
        loop {
            let (current_input, last_flag) = take(1_usize)(input)?;
            let (current_input, part) = take(4_usize)(current_input)?;
            parts.push(part);
            input = current_input;
            if last_flag == "0" {
                break;
            }
        }

        let binary = parts.iter().join("");

        Ok((input, binary))
    }

    fn content_literal(input: &str) -> IResult<&str, u64> {
        map_res(
            content_literal_str,
            from_owned_binary,
        )(input)
    }

    fn content_operator_type_length(input: &str, type_id: u64) -> IResult<&str, PacketContent> {
        let (input, len) = n_bin(input, 15)?;
        let (input, to_process) = take(len)(input)?;
        let (_, inner) = many0(packet)(to_process)?;

        Ok((input, PacketContent::Operator {
            inner,
            operator: OperatorType::from_num(type_id),
        }))
    }

    fn content_operator_type_amount(input: &str, type_id: u64) -> IResult<&str, PacketContent> {
        let (input, amount) = n_bin(input, 11)?;
        let (input, inner) = count(packet, amount as usize)(input)?;

        Ok((input, PacketContent::Operator {
            operator: OperatorType::from_num(type_id),
            inner,
        }))
    }

    pub fn packet(input: &str) -> IResult<&str, Packet> {
        let (input, version) = n_bin(input, 3)?;
        let (input, type_id) = n_bin(input, 3)?;

        if type_id == 4 {
            let (input, content) = content_literal(input)?;

            return Ok((input, Packet {
                version,
                content: PacketContent::Literal(content),
            }));
        }

        let (input, len_id) = n_bin(input, 1)?;
        let (input, content) = match len_id {
            0 => content_operator_type_length(input, type_id)?,
            1 => content_operator_type_amount(input, type_id)?,
            _ => unreachable!(),
        };

        Ok((input, Packet { version, content }))
    }
}

pub fn run_first(example: bool) -> Result<u64> {
    let input = read_input(16, example)?;
    let packet = Packet::from_hex(&input)?;

    Ok(packet.sum_version())
}

pub fn run_second(example: bool) -> Result<u64> {
    let input = read_input(16, example)?;
    let packet = Packet::from_hex(&input)?;

    Ok(packet.eval())
}