AdvenOfCode2022/src/prob7.rs

use camino::Utf8PathBuf;
use id_tree::{InsertBehavior, Node, Tree};
use nom::{
    branch::alt,
    bytes::complete::{tag, take_while1},
    combinator::{all_consuming, map},
    sequence::{preceded, separated_pair},
    Finish, IResult,
};

#[derive(Debug)]
struct Ls;

#[derive(Debug)]
struct Cd(Utf8PathBuf);

#[derive(Debug)]
enum Entry {
    Dir(Utf8PathBuf),
    File(u64, Utf8PathBuf),
}

#[derive(Debug)]
enum Command {
    Ls,
    Cd(Utf8PathBuf),
}

#[derive(Debug)]
enum Line {
    Command(Command),
    Entry(Entry),
}

impl From<Ls> for Command {
    fn from(_ls: Ls) -> Self {
        Command::Ls
    }
}

impl From<Cd> for Command {
    fn from(cd: Cd) -> Self {
        Command::Cd(cd.0)
    }
}

#[allow(dead_code)]
#[derive(Debug, Default)]
struct FsEntry {
    path: Utf8PathBuf,
    size: u64,
}

fn parse_entry(input: &str) -> IResult<&str, Entry> {
    let parse_file = map(
        separated_pair(nom::character::complete::u64, tag(" "), parse_path),
        |(size, path)| Entry::File(size, path),
    );
    let parse_dir = map(preceded(tag("dir "), parse_path), Entry::Dir);

    alt((parse_file, parse_dir))(input)
}

fn parse_path(input: &str) -> IResult<&str, Utf8PathBuf> {
    map(
        take_while1(|c: char| "abcdefghijklmnopqrstuvwxyz./".contains(c)),
        Into::into,
    )(input)
}

fn parse_ls(input: &str) -> IResult<&str, Ls> {
    map(tag("ls"), |_| Ls)(input)
}
fn parse_cd(input: &str) -> IResult<&str, Cd> {
    map(preceded(tag("cd "), parse_path), Cd)(input)
}

fn parse_command(line: &str) -> IResult<&str, Command> {
    let (input, _) = tag("$ ")(line)?;
    alt((map(parse_ls, Into::into), map(parse_cd, Into::into)))(input)
}

fn parse_line(input: &str) -> IResult<&str, Line> {
    alt((
        map(parse_command, Line::Command),
        map(parse_entry, Line::Entry),
    ))(input)
}

fn total_size(tree: &Tree<FsEntry>, node: &Node<FsEntry>) -> color_eyre::Result<u64> {
    let mut total = node.data().size;
    for child in node.children() {
        total += total_size(tree, tree.get(child)?)?;
    }
    Ok(total)
}

fn generate_tree(content: &str) -> color_eyre::Result<Tree<FsEntry>> {
    let lines = content.lines().map(|line| {
        all_consuming(parse_line)(line)
            .finish()
            .unwrap_or_else(|err| panic!("{err}"))
            .1
    });

    let mut tree = Tree::<FsEntry>::new();
    let root = tree.insert(
        Node::new(FsEntry {
            path: "/".into(),
            size: 0,
        }),
        InsertBehavior::AsRoot,
    )?;
    let mut curr = root;

    for line in lines {
        match line {
            Line::Command(cmd) => match cmd {
                Command::Ls => {}
                Command::Cd(path) => match path.as_str() {
                    "/" => {}
                    ".." => {
                        curr = tree.get(&curr)?.parent().unwrap().clone();
                    }
                    _ => {
                        curr = tree.insert(
                            Node::new(FsEntry {
                                path: path.clone(),
                                size: 0,
                            }),
                            InsertBehavior::UnderNode(&curr),
                        )?;
                    }
                },
            },
            Line::Entry(entry) => match entry {
                Entry::Dir(_) => {}
                Entry::File(size, name) => {
                    tree.insert(
                        Node::new(FsEntry { path: name, size }),
                        InsertBehavior::UnderNode(&curr),
                    )?;
                }
            },
        }
    }
    Ok(tree)
}

pub fn solve_part1(content: &str) -> color_eyre::Result<u64> {
    let tree = generate_tree(content)?;
    let sum = tree
        .traverse_pre_order(tree.root_node_id().unwrap())?
        .filter(|n| !n.children().is_empty())
        .map(|n| total_size(&tree, n).unwrap())
        .filter(|&s| s <= 100_000)
        .sum::<u64>();
    Ok(sum)
}

pub fn solve_part2(content: &str) -> color_eyre::Result<u64> {
    let total_space = 70000000_u64;
    let tree = generate_tree(content)?;
    let used_space = total_size(&tree, tree.get(tree.root_node_id().unwrap())?)?;
    let free_space = total_space.checked_sub(used_space).unwrap();
    let needed_free_space = 30000000_u64;
    let minimum_space_to_free = needed_free_space.checked_sub(free_space).unwrap();

    let size_to_remove = tree
        .traverse_pre_order(tree.root_node_id().unwrap())?
        .filter(|n| !n.children().is_empty())
        .map(|n| total_size(&tree, n).unwrap())
        .filter(|&s| s >= minimum_space_to_free)
        .min().unwrap();
    Ok(size_to_remove)
}