advent-of-code/2022/src/days/day7.rs

use anyhow::{anyhow, Result};
use std::collections::BTreeMap;
use std::fs;

#[derive(Debug, Default)]
struct Dir {
    name: String,
    files: Vec<File>,
    directories: Vec<Dir>,
    total_size: usize,
}

#[derive(Debug)]
struct File {
    name: String,
    size: usize,
}

fn dig_directory<'a, 'b>(dirname: &'b String, fs: &'a mut Dir) -> Option<&'a mut Dir> {
    if dirname == "/" && fs.name == *dirname {
        return Some(fs);
    }
    for entry in fs.directories.iter_mut() {
        if entry.name == *dirname {
            return Some(entry);
        } else if dirname.starts_with(&entry.name) {
            return dig_directory(dirname, entry);
        }
    }

    None
}

fn just_sizes<'a>(dir: &'a mut Dir) -> usize {
    let mut size: usize = 0;

    for file in dir.files.iter() {
        size += file.size;
    }

    for subdir in dir.directories.iter_mut() {
        size += just_sizes(subdir);
    }

    dir.total_size = size;
    size
}

fn part_one_sizes<'a>(dir: &'a Dir) -> usize {
    let mut size: usize = 0;

    if dir.total_size > 100000 {
        size = 0
    } else {
        size += dir.total_size;
    }

    for subdir in dir.directories.iter() {
        size += part_one_sizes(subdir);
    }

    size
}

fn part_two_sizes<'a>(dir: &'a Dir, needed: u32) -> Option<(u32, &'a str)> {
    let mut map: BTreeMap<u32, &'a str> = BTreeMap::new();
    part_two_recurse(&mut map, dir);

    for (key, val) in map.iter() {
        if *key >= needed {
            return Some((*key, val));
        }
    }

    None
}

fn part_two_recurse<'a>(map: &mut BTreeMap<u32, &'a str>, dir: &'a Dir) {
    map.insert(
        (dir.total_size as u32).try_into().unwrap(),
        dir.name.as_str(),
    );
    for subdir in dir.directories.iter() {
        part_two_recurse(map, subdir);
    }
}

pub fn run() -> Result<()> {
    #[cfg(feature = "test_input")]
    let file_contents = fs::read_to_string("tests/day7.txt")?;

    #[cfg(not(feature = "test_input"))]
    let file_contents = fs::read_to_string("day7.txt")?;

    let mut root = Dir {
        name: String::from("/"),
        ..Default::default()
    };
    let mut current_path: Vec<&str> = Vec::with_capacity(25);
    let mut current = &mut root;

    for line in file_contents.lines() {
        if &line[0..1] == "$" {
            let cmd = &line[2..4];
            match cmd {
                "ls" => {
                    continue;
                }
                "cd" => {
                    let dirname = &line[5..];
                    current = &mut root;
                    if dirname == ".." {
                        current_path.pop().expect("going past root");
                        let newdir = format!("/{}", current_path.join("/"));
                        current = dig_directory(&newdir, current).expect("missing .. dir");
                    } else if dirname == "/" {
                        current_path.clear();
                    } else {
                        current_path.push(dirname);
                        let newdir = format!("/{}", current_path.join("/"));
                        current = dig_directory(&newdir, current).expect("missing subdir");
                    }
                    continue;
                }
                _ => {
                    return Err(anyhow!("unknown command {}", line));
                }
            }
        } else if line.chars().next().unwrap().is_numeric() {
            let mut fields = line.split_whitespace();
            let size: usize = fields
                .next()
                .expect("missing file size")
                .parse()
                .expect("invalid file size");
            let name = fields.next().expect("missing file name");

            current_path.push(name);
            let new_filename = format!("/{}", current_path.join("/"));
            current_path.pop();

            if !current.files.iter().any(|x| x.name == new_filename) {
                current.files.push(File {
                    name: new_filename,
                    size,
                });
            }
        } else if line.starts_with("dir") {
            let current_dirname = &line[4..];

            current_path.push(current_dirname);
            let new_dirname = format!("/{}", current_path.join("/"));
            current_path.pop();

            if !current.directories.iter().any(|x| x.name == new_dirname) {
                let dir = Dir {
                    name: new_dirname,
                    ..Default::default()
                };
                current.directories.push(dir);
            }
        } else {
            return Err(anyhow!("malformed line {}", line));
        }
    }

    just_sizes(&mut root);
    println!("part one {}", part_one_sizes(&root));

    let total_disk = 70000000;
    let min_free_space = 30000000;
    let free_space = total_disk - root.total_size as u32;
    let needed_space = min_free_space - free_space;
    println!(
        "total size used {} free {} needed {}",
        root.total_size, free_space, needed_space
    );

    if let Some((rm_space, rm_name)) = part_two_sizes(&root, needed_space) {
        println!(
            "part two: should delete {} to reclaim {} total free space {}",
            rm_name,
            rm_space,
            free_space + rm_space
        );
    } else {
        println!("could not find part 2");
    }

    Ok(())
}