dumb_shit/2022-aoc/src/d7.rs

use aoc_runner_derive::{aoc as aoc_run, aoc_generator};
use lazy_static::lazy_static;
use regex::Regex;
use std::collections::{HashMap, HashSet};

type Files = HashMap<String, u32>;
type Dirs = HashMap<String, Node>;

#[derive(Default, Debug)]
struct Node {
    files: HashSet<String>,
    dirs: HashSet<String>,
}

enum Parsing {
    Command,
    Listing,
}

fn command(c: &str, cwd: &mut Vec<String>) -> Parsing {
    if c.starts_with("$ cd") {
        let dest = String::from_utf8(c.as_bytes()[5..].to_vec()).unwrap();
        if dest == ".." {
            let _ = cwd.pop();
        } else {
            cwd.push(dest);
        }
        return Parsing::Command;
    }
    Parsing::Listing
}

fn listing(c: &str, cwd: &[String], files: &mut Files, node: &mut Node) {
    lazy_static! {
        static ref FILE: Regex = Regex::new(r"^(\d+) ([a-z.]+$)").unwrap();
        static ref DIR: Regex = Regex::new(r"^dir ([a-z.]+$)").unwrap();
    }
    let mut path = cwd.to_vec();
    if let Some(cap) = FILE.captures(c) {
        let sz = cap[1].parse().unwrap();
        let name = cap[2].to_owned();
        path.push(name);
        let name = path.join("/");
        files.insert(name.clone(), sz);
        node.files.insert(name);
    } else if let Some(cap) = DIR.captures(c) {
        let name = cap[1].to_owned();
        path.push(name);
        let name = path.join("/");
        node.dirs.insert(name);
    } else {
        dbg!("what the fuck", c);
    }
}

#[aoc_generator(day7)]
fn parse_input_day1(input: &str) -> (Dirs, Files) {
    let mut dirs = Dirs::new();
    let mut files = Files::new();
    let mut cwd = Vec::new();
    let mut pstatus = Parsing::Command;
    let mut node = Node::default();
    for line in input.lines() {
        match pstatus {
            Parsing::Command => {
                if line.starts_with('$') {
                    pstatus = command(line, &mut cwd);
                } else {
                    pstatus = Parsing::Listing;
                    listing(line, &cwd, &mut files, &mut node);
                }
            }
            Parsing::Listing => {
                if line.starts_with('$') {
                    // we are just finished with a listing, so add the node to the map
                    dirs.insert(cwd.join("/"), node);
                    node = Node::default();
                    pstatus = command(line, &mut cwd);
                } else {
                    listing(line, &cwd, &mut files, &mut node);
                }
            }
        }
    }
    dirs.insert(cwd.join("/"), node);
    (dirs, files)
}

#[aoc_run(day7, part1)]
fn part1((dirs, files): &(Dirs, Files)) -> u32 {
    let mut sizes = Vec::new();
    for (name, _node) in dirs.iter() {
        sizes.push((name, size(name, dirs, files)));
    }
    sizes
        .into_iter()
        .filter_map(|e| if e.1 <= 100000 { Some(e.1) } else { None })
        .sum()
}

#[aoc_run(day7, part2)]
fn part2((dirs, files): &(Dirs, Files)) -> u32 {
    let mut sizes = Vec::new();
    for (name, _node) in dirs.iter() {
        sizes.push((name, size(name, dirs, files)));
    }
    let total = size("/", dirs, files);
    let avail = 70_000_000 - total;
    let needed = 30_000_000 - avail;

    sizes.sort_by(|a, b| (a.1).cmp(&b.1));
    sizes.iter().find(|x| x.1 >= needed).unwrap().1
}

fn size(name: &str, dirs: &Dirs, files: &Files) -> u32 {
    let mut out = 0;
    if let Some(node) = dirs.get(name) {
        for file in node.files.iter() {
            out += files.get(file).unwrap();
        }
        for dir in node.dirs.iter() {
            out += size(dir, dirs, files);
        }
    }

    out
}