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

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

#[derive(Clone,Copy,Debug,PartialEq)]
enum NodeType {
    Small,
    Large,
}

#[derive(Clone,Debug,PartialEq)]
struct Node {
    name: String,
    size: NodeType,
}

impl Node {
    fn new<S: ToString>(name: S) -> Result<Self> {
        let n = name.to_string();
        let size = if n.chars().filter(|c| c.is_lowercase()).count() == n.len() {
            Ok(NodeType::Small)
        } else if n.chars().filter(|c| c.is_uppercase()).count() == n.len() {
            Ok(NodeType::Large) 
        } else {
            Err(anyhow!(r#"Could not determine whether {} is a large or small cave"#, n))
        }?;
        Ok(Self {
            name: n,
            size,
        })
    }

    fn size(&self) -> NodeType {
        self.size
    }
}

#[derive(Clone,Debug)]
struct NodeList {
    inner: Vec<Node>,
}

impl NodeList {
    fn new<I, S>(init: I) -> Result<Self>
        where I: Clone + IntoIterator<Item = (S, S)>, S: ToString
    {
        let mut nodelist = NodeList { inner: vec![], };
        for (this, other) in init.into_iter() {
            for name in [this.to_string(), other.to_string()].iter() {
                if !nodelist.contains_key(name) {
                    let node = Node::new(name)?;
                    nodelist.push(node);
                }
            }
        }
        Ok(nodelist)
    }

    fn names(&self) -> Vec<&str> {
        (0..self.inner.len()).into_iter().map(|index| {
            self.inner[index].name.as_ref()
        }).collect()
    }

    fn push(&mut self, node: Node) {
        self.inner.push(node);
    }

    fn contains_key(&self, name: &str) -> bool {
        for index in 0..self.inner.len() {
            if self.inner[index].name == name {
                return true;
            }
        }
        false
    }

    fn get(&self, name: &str) -> Option<&Node> {
        for index in 0..self.inner.len() {
            if self.inner[index].name == name {
                if let Some(node) = self.inner.get(index) {
                    return Some(node)
                }
            }
        }
        None
    }
}

#[derive(Clone,Debug)]
struct Edge<'a> {
    nodes: Vec<&'a Node>,
}

impl<'a> Edge<'a> {
    fn new(nodes: Vec<&'a Node>) -> Self {
        Self {
            nodes,
        }
    }

    fn contains(&self, needle: &Node) -> bool {
        for node in &self.nodes {
            if **node == *needle {
                return true;
            }
        }
        false
    }
}

#[derive(Clone,Debug)]
struct EdgeList<'a> {
    inner: Vec<Edge<'a>>,
}

impl<'a> EdgeList<'a> {
    fn new<I, S>(init: I, nodelist: &'a NodeList) -> Result<Self>
        where I: IntoIterator<Item = (S, S)>, S: ToString
    {
        let inner = init.into_iter().map(|(this, other)| {
            let this_node = nodelist.get(&this.to_string())
                .ok_or(anyhow!("Failed creating EdgeList"))?;
            let other_node = nodelist.get(&other.to_string())
                .ok_or(anyhow!("Failed creating EdgeList"))?;
            Ok(Edge::new(vec![this_node, other_node]))
        }).collect::<Result<Vec<_>>>()?;
        Ok(Self {
            inner,
        })
    }

    fn neighbours(&self, this: &Node) -> Vec<&Node> {
        let mut nodes = vec![];
        for edge in self.inner.iter() {
            if edge.contains(this) {
                for node in edge.nodes.iter().filter(|node| ***node != *this) {
                    nodes.push(*node);
                }
            }
        }

        nodes
    }
}

#[derive(Clone,Copy,Debug,Eq,PartialEq)]
enum State {
    Explorable,
    Unexplored,
    Visited,
    Exhausted,
}

#[derive(Clone,Debug)]
struct VisitList {
    inner: HashMap<String, Cell<State>>,
}

impl VisitList {
    fn new<'a, I>(init: I, state: State) -> Self
        where I: IntoIterator<Item = &'a str>
    {
        let inner = init.into_iter()
            .map(|key| (String::from(key), Cell::new(state))).collect();

        Self {
            inner,
        }
    }

    fn state(&self, key: &str) -> Option<State> {
        self.inner.get(key).map(|cell| cell.get())
    }

    fn set_state(&self, key: &str, state: State) -> Result<()> {
        self.inner.get(key)
            .ok_or_else(|| anyhow!("Could not set state for non-existent: {}", key))?.set(state);
        Ok(())
    }
}

#[derive(Clone,Debug)]
struct Cave<'a> {
    nodelist: &'a NodeList,
    edgelist: EdgeList<'a>,
}

impl<'a> Cave<'a> {
    fn new(nodelist: &'a NodeList, edgelist: EdgeList<'a>) -> Result<Self> {
        Ok(Self{
            nodelist,
            edgelist,
        })
    }

    fn explorable<'b, 'c: 'b>(&'b self, node: &'c Node) -> Explorable {
        Explorable::new(&self.edgelist, node)
    }

    fn recurse(&self, visited: VisitList, node: &Node) -> Result<usize> {
        let mut paths = 0;
        let node_state = visited.state(&node.name)
            .ok_or_else(|| anyhow!("No state for {}", node.name))?;
        match (node.name.as_str(), node.size(), node_state) {
            ("start", _, _) => {
                visited.set_state(&node.name, State::Exhausted)?;
            },
            ("end", _, _) => {
                return Ok(1);
            },
            (_, _,               State::Exhausted) => {
                return Ok(0);
            }
            (_, NodeType::Small, State::Explorable) => {
                visited.set_state(&node.name, State::Visited)?;
            }
            (_, NodeType::Small, State::Unexplored) => {
                visited.set_state(&node.name, State::Exhausted)?;
            }
            (_, NodeType::Small, State::Visited) => {
                for cave in self.nodelist.inner.iter() {
                    let cave_state = visited.state(&cave.name);
                    if cave.size() == NodeType::Small {
                        if cave_state == Some(State::Explorable) {
                            visited.set_state(&cave.name, State::Unexplored)?;
                        } else {
                            visited.set_state(&cave.name, State::Exhausted)?;
                        }
                    }
                }
                visited.set_state(&node.name, State::Exhausted)?;
            }
            (_, NodeType::Large, State::Explorable | State::Unexplored | State::Visited) => {
                visited.set_state(&node.name, State::Visited)?;
            }
        }

        let explorable = self.explorable(node)
            .collect::<Vec<_>>();

         for next_node in explorable {
             visited.set_state("end", State::Unexplored)?;
             let next_state = visited.state(&next_node.name)
                 .ok_or_else(|| anyhow!("Missing state for: {}", next_node.name))?;

             match (next_node.size(), next_state) {
                 (_,               State::Exhausted)  => {
                     continue;
                 },
                 (NodeType::Small, State::Explorable | State::Unexplored | State::Visited) => {
                     paths += self.recurse(visited.clone(), next_node)?;
                 },
                 (NodeType::Large, State::Explorable | State::Unexplored | State::Visited) => {
                     paths += self.recurse(visited.clone(), next_node)?;
                 },
             }
        }
        Ok(paths)
    }

    fn walk_all(&self, state: State) -> Result<usize> {
        let nodelist = &self.nodelist;
        let node_names = nodelist.names();
        let start = nodelist.get("start").ok_or(anyhow!("Could not find start node"))?;
        let visited = VisitList::new(node_names, state);

        self.recurse(visited, start)
    }
}

struct Explorable<'a, 'b> {
    positition: usize,
    edgelist: &'a EdgeList<'a>,
    node: &'b Node,
}

impl<'a, 'b> Explorable<'a, 'b> {
    fn new(edgelist: &'a EdgeList, node: &'b Node) -> Self {
        Self {
            positition: 0,
            edgelist,
            node,
        }
    }
}

impl<'a, 'b> Iterator for Explorable<'a, 'b> {
    type Item = &'a Node;

    fn next(&mut self) -> Option<Self::Item> {
        let neighbours = self.edgelist.neighbours(self.node);
        let next = neighbours.into_iter().nth(self.positition);
        self.positition += 1;
        next
    }
}

fn read_input<T: AsRef<Path>>(filename: T) -> Result<Vec<(String, String)>> {
    let reader = BufReader::new(File::open(filename)?);

    reader.lines().map(
        |v| {
            let s = v?;
            let (left, right) = s.split_once('-')
                .ok_or_else(|| anyhow!("Unable to parse: {}", s))?;
            Ok((String::from(left), String::from(right)))
        }
    ).collect()
}

fn part1(cave: Cave) -> Result<usize> {
    let paths = cave.walk_all(State::Unexplored)?;
    Ok(paths)
}

fn part2(cave: Cave) -> Result<usize> {
    let paths = cave.walk_all(State::Explorable)?;
    Ok(paths)
}

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 input = read_input(filename).context("Could not read input")?;
    let nodelist = NodeList::new(input.clone())?;
    let edgelist = EdgeList::new(input, &nodelist)?;
    let cave = Cave::new(&nodelist, edgelist)?;
    if do_part_1 {
        let solution = part1(cave.clone()).context("No solution for part 1")?;
        println!("Part1, solution found to be: {}", solution);
    }
    if do_part_2 {
        let solution = part2(cave).context("No solution for part 2")?;
        println!("Part2, solution found to be: {}", solution);
    }
    Ok(())
}