2022/src/bin/day07.rs

use std::{cell::RefCell, cmp::min, collections::BTreeMap, rc::Rc, str::FromStr};

use aoc_2022::*;

use color_eyre::{eyre::Result, Report};
use itertools::Itertools;
use tracing::*;
use tracing_subscriber::EnvFilter;

type FileHandle = Rc<RefCell<AocFile>>;

#[derive(Debug)]
enum AocFile {
    File(usize),
    Directory(BTreeMap<String, FileHandle>),
}

impl Default for AocFile {
    fn default() -> Self {
        AocFile::Directory(BTreeMap::new())
    }
}

impl AocFile {
    fn dir() -> FileHandle {
        Rc::new(RefCell::new(AocFile::Directory(BTreeMap::new())))
    }

    fn file(size: usize) -> FileHandle {
        Rc::new(RefCell::new(AocFile::File(size)))
    }

    fn add_file(&mut self, name: &str, size: usize) {
        if let AocFile::Directory(files) = self {
            files.insert(name.to_string(), AocFile::file(size));
            return;
        }

        panic!("cannot cd in file")
    }

    fn cd(&mut self, dir: &str) -> FileHandle {
        if let AocFile::Directory(files) = self {
            return files.entry(dir.to_string()).or_default().clone();
        }

        panic!("cannot cd in file")
    }

    fn size_under_100000(&self) -> (bool, usize, usize) {
        match self {
            AocFile::File(s) => (false, 0, *s),
            AocFile::Directory(files) => {
                let (running_total, size) = files
                    .values()
                    .map(|f| f.borrow().size_under_100000())
                    .fold((0_usize, 0_usize), |(mut running, size), (dir, r, s)| {
                        if dir && s <= 100000 {
                            running += s;
                        }

                        (r + running, size + s)
                    });

                (true, running_total, size)
            }
        }
    }

    fn smallest_bigger(&self, required: usize) -> (bool, usize, usize) {
        match self {
            AocFile::File(s) => (false, usize::MAX, *s),
            AocFile::Directory(files) => {
                let (mut dir_min, size) = files
                    .values()
                    .map(|f| f.borrow().smallest_bigger(required))
                    .fold(
                        (usize::MAX, 0_usize),
                        |(mut current_min, size), (dir, r_min, s)| {
                            if dir && s >= required {
                                // we got back from a subdirectory
                                current_min = min(current_min, s);
                            }

                            // also check local solutions in subdirectories
                            (min(current_min, r_min), size + s)
                        },
                    );

                if size >= required {
                    dir_min = min(dir_min, size);
                }

                (true, dir_min, size)
            }
        }
    }
}

struct Filesystem {
    root: FileHandle,
    cwd: Vec<FileHandle>,
}

impl Filesystem {
    fn new() -> Filesystem {
        let root = AocFile::dir();

        Filesystem {
            root: root.clone(),
            cwd: vec![root],
        }
    }

    fn cd(&mut self, dir: &str) {
        match dir {
            ".." => {
                self.cwd.pop();
            }
            "/" => {
                self.cwd.clear();
                self.cwd.push(self.root.clone());
            }
            _ => {
                let idx = self.cwd.len() - 1;
                let subdir = self.cwd[idx].borrow_mut().cd(dir);
                self.cwd.push(subdir);
            }
        }
    }

    fn touch(&mut self, name: &str, size: usize) {
        let idx = self.cwd.len() - 1;
        self.cwd[idx].borrow_mut().add_file(name, size);
    }

    fn run_command(&mut self, command: &str) {
        if command.starts_with("cd ") {
            let parts = command.split_ascii_whitespace().collect_vec();
            self.cd(parts[1]);
            return;
        }

        for file in command.lines().skip(1) {
            let parts = file.split_ascii_whitespace().collect_vec();
            if parts[0] == "dir" {
                /* no-op */
            } else {
                let name = parts[1];
                let size: usize = parts[0].parse().unwrap();
                self.touch(name, size);
            }
        }
    }
}

impl FromStr for Filesystem {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut fs = Filesystem::new();

        for command in s.trim_start_matches("$ ").split("\n$ ") {
            fs.run_command(command);
        }

        Ok(fs)
    }
}

type Input = Filesystem;
type Output = usize;

fn part_1(input: &Input) -> Output {
    input.root.borrow().size_under_100000().1
}

fn part_2(input: &Input) -> Output {
    let (total, needed) = (70000000, 30000000);
    let used = input.root.borrow().size_under_100000().2;

    // to_be_freed >= needed - (total - used)
    let to_be_freed = needed - (total - used);

    input.root.borrow().smallest_bigger(to_be_freed).1
}

fn parse_input(pathname: &str) -> Input {
    file_to_string(pathname).parse().unwrap()
}

fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(EnvFilter::from_default_env())
        .with_target(false)
        .with_file(true)
        .with_line_number(true)
        .without_time()
        .compact()
        .init();
    color_eyre::install()?;

    let input = parse_input("inputs/day07.txt");

    info!("Part 1: {}", part_1(&input));
    info!("Part 2: {}", part_2(&input));

    Ok(())
}

#[cfg(test)]
mod day_07 {
    use super::*;

    #[test]
    fn test_part_1() {
        let sample = parse_input("samples/day07.txt");
        assert_eq!(part_1(&sample), 95437);
    }

    #[test]
    fn test_part_2() {
        let sample = parse_input("samples/day07.txt");
        assert_eq!(part_2(&sample), 24933642);
    }
}