day20.rs

use petgraph::prelude::*;
use regex_syntax::hir::{self, Hir, HirKind};
use regex_syntax::ParserBuilder;
use std::collections::{HashMap, HashSet};
use std::ops::Add;

type MyGraph = DiGraph<Pos, f32>;
type GraphMap = HashMap<Pos, NodeIndex>;

#[derive(Debug, Default, Clone, Copy, Hash, PartialEq, Eq)]
struct Pos(i32, i32);

impl Add for Pos {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Pos(self.0 + rhs.0, self.1 + rhs.1)
    }
}

const NORTH: Pos = Pos(0, 1);
const SOUTH: Pos = Pos(0, -1);
const WEST: Pos = Pos(-1, 0);
const EAST: Pos = Pos(1, 0);

#[derive(Debug, Default)]
struct Vis {
    gm: GraphMap,
    graph: MyGraph,
    positions: Vec<Vec<Pos>>,
    nestiness: Vec<(usize, usize)>,
}

impl hir::Visitor for Vis {
    type Output = Vec<f32>;
    type Err = ();

    fn start(&mut self) {
        const START: Pos = Pos(0, 0);
        self.positions.push(vec![START]);
        self.gm.insert(START, self.graph.add_node(START));
    }

    fn visit_pre(&mut self, hir: &Hir) -> Result<(), Self::Err> {
        match hir.kind() {
            HirKind::Literal(hir::Literal::Unicode(c)) => {
                for pos in self.positions.last_mut().unwrap() {
                    let npos = *pos
                        + match c {
                            'N' => NORTH,
                            'S' => SOUTH,
                            'W' => WEST,
                            'E' => EAST,
                            _ => unimplemented!("No more characters allowed!"),
                        };
                    let gm = &mut self.gm;
                    let graph = &mut self.graph;
                    let nix = *gm.entry(npos).or_insert_with(|| graph.add_node(npos));
                    self.graph.update_edge(self.gm[&pos], nix, 1.0);
                    *pos = npos;
                }
            }
            HirKind::Alternation(_) => {
                self.nestiness.push((self.positions.len() - 1, 1));
                let last = &self.positions[self.nestiness.last().unwrap().0];
                let mut nv = Vec::with_capacity(last.len());
                for pos in last {
                    nv.push(pos.clone())
                }
                self.positions.push(nv);
            }
            _ => {}
        }
        Ok(())
    }

    fn visit_post(&mut self, hir: &Hir) -> Result<(), Self::Err> {
        if let HirKind::Alternation(_) = hir.kind() {
                let (_, deepness) = self.nestiness.pop().unwrap();
                let mut nv = HashSet::new();
                for _ in 0..deepness {
                    nv.extend(self.positions.pop().unwrap().into_iter());
                }
                self.positions.push(nv.into_iter().collect());
        }
        Ok(())
    }

    fn visit_alternation_in(&mut self) -> Result<(), Self::Err> {
        self.nestiness.last_mut().unwrap().1 += 1;
        let last = &self.positions[self.nestiness.last().unwrap().0];
        let mut nv = Vec::with_capacity(last.len());
        for pos in last {
            nv.push(pos.clone())
        }
        self.positions.push(nv);
        Ok(())
    }

    fn finish(self) -> Result<Self::Output, Self::Err> {
        Ok(petgraph::algo::bellman_ford(&self.graph, self.gm[&Pos(0, 0)])
            .unwrap()
            .0)
    }
}

fn solve_a(s: &str) -> u32 {
    let hir = ParserBuilder::new().nest_limit(1000).build().parse(s).unwrap();
    hir::visit(&hir, Vis::default())
        .unwrap()
        .into_iter()
        .map(|f| f as u32)
        .max()
        .unwrap()
}
fn solve_b(s: &str) -> usize {
    let hir = ParserBuilder::new().nest_limit(1000).build().parse(s).unwrap();
    hir::visit(&hir, Vis::default())
        .unwrap()
        .into_iter()
        .map(|f| f as u32)
        .filter(|&n| n >= 1000)
        .count()
}

fn main() {
    let input = include_str!("../challenge.txt");;
    println!("A: {}", solve_a(&input[1..input.len() - 1]));
    println!("B: {}", solve_b(&input[1..input.len() - 1]));
}

#[test]
fn test_short() {
    assert_eq!(3, solve("WNE"));
}

#[test]
fn test_simple_branch() {
    assert_eq!(5, solve("W(N|WWWW)"));
}

#[test]
fn test_triple_branch() {
    assert_eq!(7, solve("W(NNE|WWWW|SWS)NN"));
}

#[test]
fn test_simple_group() {
    assert_eq!(4, solve("W(WW)W"));
}

#[test]
fn test_nested_branch() {
    assert_eq!(10, solve("ENWWW(NEEE|SSE(EE|N))"));
}

#[test]
fn test_empty_branch() {
    assert_eq!(18, solve("ENNWSWW(NEWS|)SSSEEN(WNSE|)EE(SWEN|)NNN"));
}

#[test]
fn more_examples() {
    assert_eq!(23, solve("ESSWWN(E|NNENN(EESS(WNSE|)SSS|WWWSSSSE(SW|NNNE)))"));
    assert_eq!(
        31,
        solve("WSSEESWWWNW(S|NENNEEEENN(ESSSSW(NWSW|SSEN)|WSWWN(E|WWS(E|SS))))")
    );
}