path: root/2021/rust/day14/src/main.rs

use {
    anyhow::{
        anyhow,
        Context,
        Result,
    },
    regex::Regex,
    std::{
        collections::HashMap,
        env::args,
        fs::File,
        io::{
            BufRead,
            BufReader,
        },
        path::Path,
    },
};

#[derive(Debug)]
struct InsertionRule {
    pair: String,
    element: String,
}

impl InsertionRule {
    fn new(pair: &str, element: &str) -> Result<Self> {
        Ok(Self{
            pair: pair.into(),
            element: element.into(),
        })
    }

    fn new_apply(&self, input: &str) -> Result<String> {
        let mut iter = input.chars().take(2);
        let (left, right) = (
            iter.next().ok_or_else(|| anyhow!("Failure on pair: {}", input))?,
            iter.next().ok_or_else(|| anyhow!("Failure on pair: {}", input))?,
        );
        let mut output = format!("{}", left);

        if input.starts_with(&self.pair) {
            output += &format!("{}", self.element.chars().next()
                .ok_or_else(|| anyhow!("Failure with element for: {}", input))?);
        }
        output += &format!("{}", right);
        Ok(output)
    }
}

#[derive(Debug)]
struct Polymer {
    inner: String,
}

type CharCount = HashMap<char, usize>;
type CacheMap = HashMap<String, CharCount>;

impl Polymer {
    fn new(template: impl ToString) -> Self {
        Self {
            inner: template.to_string(),
        }
    }

    fn len(&self) -> usize {
        self.inner.len()
    }
    
    fn first(&self) -> Result<char> {
        self.inner.chars().next().ok_or_else(|| anyhow!("Could not get first character"))
    }

    fn transform<'a, I>(&self, rules: I) -> Result<Self>
        where I: Copy + IntoIterator<Item = &'a InsertionRule>
    {
        let first = if let Some(c) = self.inner.chars().next() {
            c.to_string()
        } else {
            String::from("")
        };
        let inner = format!("{}{}", first,
                self.inner.chars().zip(self.inner.chars().skip(1)).map(|(left, right)| {
                let input = format!("{}{}", left, right);
                for rule in rules {
                    let output = rule.new_apply(&input)?;
                    if input != output {
                        return Ok(output[1..].to_string());
                    }
                }
                Ok(input[1..].to_string())
            }).collect::<Result<String>>()?);
        Ok(Self {
            inner
        })
    }

    fn elements(&self) -> HashMap<char, usize> {
        let mut elements = HashMap::new();

        for c in self.inner.chars() {
            let entry = elements.entry(c).or_insert(0);
            *entry += 1;
        }

        elements
    }

    fn str_to_char_count(s: &str) -> HashMap<char, usize> {
        let mut elements = HashMap::new();
        for c in s.chars() {
          let entry = elements.entry(c).or_insert(0);
          *entry += 1;
        }

        elements
    }

    fn evaluate<'a, I>(&self, rules: I, depth: usize, skip_first: bool, cache: &mut CacheMap)
        -> Result<CharCount>
        where I: Copy + IntoIterator<Item = &'a InsertionRule>
    {
        if depth == 0 {
            let interesting = if skip_first {
                &self.inner[1..]
            } else {
                &self.inner
            };
            return Ok(Self::str_to_char_count(interesting));
        }
        let mut char_count = CharCount::new();
        if !skip_first {
            char_count.insert(self.first()?, 1);
        }

        for index in 0..=(self.len() - 2) {
            let sub_key = self.inner[index..index+2].to_string();
            let cache_key = &format!("{}-{}-{:?}", sub_key, depth, skip_first);
            if let Some(cached) = cache.get(cache_key) {
                for (key, value) in cached.iter() {
                    let entry = char_count.entry(*key).or_insert(0);
                    *entry += value;
                }
            } else {
                let sub_polymer = Polymer::new(&sub_key);
                let transformed = sub_polymer.transform(rules)?;

                let polymer_chars = &transformed.evaluate(rules, depth - 1, true, cache)?;
                cache.insert(cache_key.to_string(), polymer_chars.clone());
                for (key, value) in polymer_chars.iter() {
                    let entry = char_count.entry(*key).or_insert(0);
                    *entry += value;
                }
            }
        }
        Ok(char_count)
    }
}

fn read_input<T: AsRef<Path>>(filename: T) -> Result<(Polymer, Vec<InsertionRule>)> {
    let reader = BufReader::new(File::open(filename)?);
    let mut polymer = None;
    let mut rules = vec![];

    let re = Regex::new(r#"(?x)
        ^(?P<template>[A-Z]+)$
        |
        ^(?P<pair>[A-Z]{2})\ ->\ (?P<element>[A-Z])$
        "#).map_err(|err| anyhow!("Failed to compile regex: {}", err))?;

    for row in reader.lines() {
        let string = row?;
        if string.is_empty() {
            continue;
        }

        let caps = re.captures(&string).ok_or_else(|| anyhow!("Could not parse: {}", string))?;
        match (caps.name("template"), caps.name("pair"), caps.name("element")) {
            (Some(template), None, None) => polymer = Some(Polymer::new(template.as_str())),
            (None, Some(pair), Some(element)) => 
                rules.push(InsertionRule::new(pair.as_str(), element.as_str())?),
            _ => return Err(anyhow!("Could not parse: {}", string)),
        }
    }

    if let Some(p) = polymer {
        Ok((p, rules))
    } else {
        Err(anyhow!("Missing polymer template in indata"))
    }
}

fn part1<'a, I>(polymer: &Polymer, rules: I) -> Result<usize>
    where I: Copy + IntoIterator<Item = &'a InsertionRule>
{
    let mut p = polymer;
    let mut build;

    for _ in 0..10 {
        build = p.transform(rules)?;
        p = &build;
    }
    let elements = p.elements();
    let min = elements.values().min().ok_or_else(|| anyhow!("Could not find min"))?;
    let max = elements.values().max().ok_or_else(|| anyhow!("Could not find max"))?;
    Ok(max - min)
}

fn part2<'a, I>(polymer: &Polymer, rules: I) -> Result<usize>
    where I: Copy + IntoIterator<Item = &'a InsertionRule>
{
    let mut cache = HashMap::new();
    let elements = polymer.evaluate(rules, 40, false, &mut cache)?;

    let min = elements.values().min().ok_or_else(|| anyhow!("Could not find min"))?;
    let max = elements.values().max().ok_or_else(|| anyhow!("Could not find max"))?;
    Ok(max - min)
}

fn main() -> Result<()> {
    let ( do_part_1, do_part_2 ) = aoc::do_parts();

    let filename = args().nth(1).ok_or(anyhow!("Missing input filename"))?;
    let (polymer, rules) = read_input(filename).context("Could not read input")?;
    if do_part_1 {
        let solution = part1(&polymer, &rules).context("No solution for part 1")?;
        println!("Part1, solution found to be: {}", solution);
    }
    if do_part_2 {
        let solution = part2(&polymer, &rules).context("No solution for part 2")?;
        println!("Part2, solution found to be: {}", solution);
    }
    Ok(())
}