9
\$\begingroup\$

I solved the Advent of Code problems last year in Python, and I wanted to learn a new language this year, so I decided to write up Part 1 of on of the hardest problems from last year - problem 15 - in Rust.

While the problem in all its particulars is quite difficult to understand fully, the basic idea is quite simple: each Unit (with two teams: G, the Goblins, and E, the Elves) is present on a CombatGrid. Every single turn, they check if there's an enemy unit adjacent to them (where diagonal doesn't count as adjacent). If so, they attack that unit, reducing their HP. If not, they move using a weird and kind of fussy algorithm that boils down to BFS, and then proceed to attack if they're now in range of an enemy unit after moving. If not, they end their turn. This goes on until every unit of a particular team is dead.

As far as I can tell, this program is 100% correct in all its particulars - it produced the correct answer for my input and some others I tested. It uses BFS to calculate the correct moves for each unit, which I think is an appropriate enough algorithm, and it's definitely fast enough for the problem's needs.

What I'm really looking for review on is making this code more idiomatic, since I basically just read the Rust book and winged my way through writing this code. I'm not sure about the fact that I have to derive Copy and/or Clone for many of these types, it feels a bit "cheaty", and the structure of the code is also a little more awkward than it feels like it needs to be, but fighting with the borrow checker means some difficulty in correcting that problem.

use std::env;
use std::fs;
use std::fmt;
use std::error::Error;
use std::collections::BinaryHeap;
use hashbrown::HashMap;
use std::cmp::{Ordering, Reverse};
use unit::*;

pub fn main() -> Result<(), Box<dyn Error>> {
    let args = env::args().collect::<Vec<String>>();

    let input_filename = match args.len() {
        2 => &args[1],
        _ => "input.txt"
    };

    let string_grid = fs::read_to_string(input_filename)?;

    let mut combat_grid = parse_input(&string_grid)?;
    let mut full_rounds: usize = 0;

    println!("Start");
    print!("{}", combat_grid);
    println!("\n");

    while combat_grid.tick() {
        full_rounds += 1;
        println!("\n");
        println!("Round {}", full_rounds);
        print!("{}", combat_grid);
        println!("\n");
    }

    println!("Final");
    print!("{}", combat_grid);
    println!("\n");

    println!("Outcome: {}",
        full_rounds  * combat_grid.units.values().map(|u| u.hp).sum::<usize>());

    Ok(())
}

pub fn parse_input(string_grid: &str) -> Result<CombatGrid, String> {
    let mut grid = HashMap::new();
    let mut units = HashMap::new();
    let mut dimensions = (0, 0);

    for (y, row) in string_grid.lines().enumerate() {
        dimensions.1 += 1;

        for (x, character) in row.chars().enumerate() {
            dimensions.0 += 1;

            let current_location = Location { x, y };

            grid.insert(current_location, match character {
                '#' => Environment::Wall,
                '.' => Environment::Open,
                'G' | 'E' => {
                    units.insert(current_location, Unit {
                        team: if character == 'G' { UnitTeam::Goblin } else { UnitTeam::Elf },
                        location: current_location,
                        hp: 200,
                        attack_power: 3
                    });

                    Environment::Open
                },
                _ => {
                    return Err(format!("Invalid input character: {}", character));
                }
            });
        }
    }

    dimensions.0 /= dimensions.1;

    Ok(CombatGrid { grid, units, dimensions })
}

#[derive(Eq, PartialEq, Copy, Clone, Hash)]
pub struct Location {
    x: usize,
    y: usize
}

impl Location {
    fn adjacent(&self) -> [Self; 4] {
        [
            Location { x: self.x, y: self.y - 1 },
            Location { x: self.x, y: self.y + 1 },
            Location { x: self.x - 1, y: self.y },
            Location { x: self.x + 1, y: self.y }
        ]
    }
}

impl fmt::Debug for Location {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "({}, {})", self.x, self.y)
    }
}

impl Ord for Location {
    fn cmp(&self, other: &Self) -> Ordering {
        self.y.cmp(&other.y).then(self.x.cmp(&other.x))
    }
}

impl PartialOrd for Location {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

pub struct CombatGrid {
    pub grid: HashMap<Location, Environment>,
    pub units: HashMap<Location, Unit>,
    pub dimensions: (usize, usize),
}

impl fmt::Display for CombatGrid {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for y in 0..self.dimensions.1 {
            let mut row_units = Vec::new();

            for x in 0..self.dimensions.0 {
                let location = Location { x, y };

                if let Some(unit) = self.units.get(&location) {
                    write!(f, "{:?}", unit.team)?;
                    row_units.push(unit);
                } else if let Some(env) = self.grid.get(&location) {
                    write!(f, "{:?}", env)?;
                }
            }

            write!(f, "\t")?;

            for unit in row_units {
                write!(f, " [{:?}] ", unit)?;
            }

            writeln!(f)?;
        }

        Ok(())
    }
}

impl CombatGrid {
    pub fn tick(&mut self) -> bool {
        let mut unit_locations = self.units.keys().cloned().collect::<Vec<_>>();
        unit_locations.sort_unstable();

        for unit_location in unit_locations.iter() {
            // This unit may have since died by the hands of another
            // by the time we have gotten to it, so check if it's still there.
            let unit = match self.units.get(unit_location) {
                Some(unit) => unit.clone(),
                None => continue
            };

            let enemy_units = self.units
                .iter()
                .filter(|(_, u)| u.is_enemy(&unit))
                .map(|(l, u)| (*l, u.clone()))
                .collect::<HashMap<_, _>>();

            if enemy_units.is_empty() {
                return false; // Combat has ended, one team has won.
            }

            if let Some(attacked_unit_location) = unit.maybe_attack(&enemy_units) {
                self.attack_unit(unit_location, &attacked_unit_location);
                continue;
            }

            if let Some(move_location) = unit.maybe_move(&enemy_units, |l| self.is_open_fn(l)) {
                // Get the new Unit with the updated location. The old reference is stale
                // otherwise, leading to attack behaviour based on the old location, which never
                // actually works out, because the only reason any unit moves is because its
                // old location is not adjacent to any enemy unit.
                let unit = self.move_unit(unit_location, &move_location);

                if let Some(attacked_unit_location) = unit.maybe_attack(&enemy_units) {
                    self.attack_unit(&move_location, &attacked_unit_location);
                }
            }
        }

        true
    }

    fn attack_unit(&mut self, current_unit_location: &Location, attacked_unit_location: &Location) {
        let current_unit = &self.units[current_unit_location].clone();
        let mut attacked_unit = self.units.get_mut(attacked_unit_location).unwrap();

        // This protects against overflows in the usize
        attacked_unit.hp =
            attacked_unit.hp.saturating_sub(current_unit.attack_power);

        if attacked_unit.is_dead() {
            self.units.remove(attacked_unit_location);
        }
    }

    fn move_unit(&mut self, current_unit_location: &Location, new_location: &Location) -> Unit {
        let new_location = *new_location;
        let mut current_unit = self.units.remove(current_unit_location).unwrap();

        current_unit.location = new_location;
        self.units.insert(new_location, current_unit.clone());

        current_unit
    }

    fn is_open_fn(&self, location: &Location) -> bool {
        if self.units.contains_key(location) {
            false
        } else if let Some(env) = self.grid.get(location) {
            env == &Environment::Open
        } else { false }
    }
}

#[derive(Eq, PartialEq)]
pub enum Environment {
    Wall,
    Open
}

impl fmt::Debug for Environment {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", if self == &Self::Wall { '#' } else { '.' })
    }
}

mod unit {
    use super::*;

    #[derive(Eq, PartialEq, Copy, Clone)]
    pub enum UnitTeam { Goblin, Elf }

    impl fmt::Debug for UnitTeam {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{}", if self == &Self::Elf { 'E' } else { 'G' })
        }
    }

    #[derive(Eq, PartialEq, Clone)]
    pub struct Unit {
        pub team: UnitTeam,
        pub location: Location,
        pub hp: usize,
        pub attack_power: usize
    }

    impl fmt::Debug for Unit {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            write!(f, "{:?}({}) @ {:?}", self.team, self.hp, self.location)
        }
    }

    impl Unit {
        pub fn is_enemy(&self, other: &Self) -> bool {
            self.team != other.team
        }

        pub fn is_dead(&self) -> bool {
            self.hp == 0
        }

        pub fn maybe_attack(&self, enemy_units: &HashMap<Location, Unit>) -> Option<Location> {
            let mut adjacent_enemy_units = enemy_units
                .values()
                .filter(|u| self.location.adjacent().contains(&u.location))
                .collect::<Vec<_>>();

            adjacent_enemy_units.sort_unstable_by(|a, b| {
                match a.hp.cmp(&b.hp) {
                    Ordering::Equal => a.location.cmp(&b.location),
                    hp_cmp => hp_cmp
                }
            });

            adjacent_enemy_units.reverse();
            adjacent_enemy_units.pop().map(|u| u.location)
        }

        pub fn maybe_move(
            &self,
            enemy_units: &HashMap<Location, Unit>,
            is_open_fn: impl Fn(&Location) -> bool
        ) -> Option<Location> {
            let mut frontier = self.location.adjacent()
                .iter()
                .cloned()
                .filter(|l| is_open_fn(l))
                .map(|l| Reverse(SearchNode {
                    distance: 1,
                    current_location: l,
                    starting_location: l,
                }))
                .collect::<BinaryHeap<_>>();

            let mut explored: Vec<Location> = Vec::new();

            while let Some(Reverse(next)) = frontier.pop() {
                for next_adjacent in next.current_location.adjacent().iter().cloned() {
                    if explored.contains(&next_adjacent) { continue; }

                    if !is_open_fn(&next_adjacent) {
                        if enemy_units.contains_key(&next_adjacent) {
                            return Some(next.starting_location);
                        }

                        continue;
                    }

                    frontier.push(Reverse(SearchNode {
                        distance: next.distance + 1,
                        current_location: next_adjacent,
                        starting_location: next.starting_location
                    }));

                    explored.push(next_adjacent);
                }
            }

            None
        }
    }

    // Private helper to make maybe_move easier to keep track of
    #[derive(Debug, Eq, PartialEq, Copy, Clone)]
    struct SearchNode {
        distance: usize,
        current_location: Location,
        starting_location: Location,
    }

    impl Ord for SearchNode {
        fn cmp(&self, other: &Self) -> Ordering {
            self.distance.cmp(&other.distance).then(
                self.current_location.cmp(&other.current_location).then(
                    self.starting_location.cmp(&other.starting_location)
                )
            )
        }
    }

    impl PartialOrd for SearchNode {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
            Some(self.cmp(&other))
        }
    }
}
\$\endgroup\$
2
  • 3
    \$\begingroup\$ I originally wanted to ask you to add the challenge description to the question, but damn ..., that's a long read ;-) \$\endgroup\$
    – AlexV
    Nov 13, 2019 at 17:12
  • 2
    \$\begingroup\$ @AlexV I added a simple problem description to help alleviate the issue, since I'm not asking for improvements to the algorithm or the logic, but simply the idiomaticity of it. \$\endgroup\$
    – naiveai
    Nov 14, 2019 at 7:13

1 Answer 1

4
\$\begingroup\$

The code look good. Here's some nitpicks:

rustfmt

First of all, I ran cargo fmt on your code. There were some differences (trailing comma, linebreaks, etc.).

anyhow

Consider changing the return type of all functions to use anyhow::Result, which provides contexts and backtracks.

return Err(format!("Invalid input character: {}", character));

can be simplified to

bail!("Invalid input character: {}", character);

Matching on enums with two variants

For enums with two variants, it is probably clearer to write

if self == &Self::Wall { '#' } else { '.' }

as

match self {
    Self::Wall => '#',
    Self::Open => '.',
}

sort_unstable_by_key

adjacent_enemy_units.sort_unstable_by(|a, b| match a.hp.cmp(&b.hp) {
    Ordering::Equal => a.location.cmp(&b.location),
    hp_cmp => hp_cmp,
});

can be written as

adjacent_enemy_units.sort_unstable_by_key(|unit| (unit.hp, unit.location));

FnMut

FnMut is more general than Fn, and can be used in the signature of maybe_move:

pub fn maybe_move(
    &self,
    enemy_units: &HashMap<Location, Unit>,
    is_open_fn: impl Fn(&Location) -> bool,
) -> Option<Location> {

derive

impl Ord for SearchNode {
    fn cmp(&self, other: &Self) -> Ordering {
        self.distance.cmp(&other.distance).then(
            self.current_location
                .cmp(&other.current_location)
                .then(self.starting_location.cmp(&other.starting_location)),
        )
    }
}

impl PartialOrd for SearchNode {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(&other))
    }
}

These implementations are equivalent to the automatic versions, so use derive.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.