Prim's Algorithm + a solver in Rust

primmaze2.rs

use std::collections::VecDeque;
use std::env;
use std::fmt::{Display, Formatter};
use std::ops::{Index, IndexMut};
use rand::RngCore;
use crate::Dir::*;

fn main() {
    let args: Vec<_> = env::args().collect();
    let dimension = if args.len() < 2 {
        println!("No dimension, defaulting to 20");
        20
    } else {
        args[1].parse::<usize>().expect("That didn't look like a number.")
    };

    let method = if args.len() < 3 {
        println!("No method, defaulting to bfs (pass a dimension and \"dead_end\" to change)");
        Methods::Bfs
    } else if args[2] == "dead_end" { Methods::DeadEnd }
    else { Methods::Bfs };

    let maze = Maze::generate(dimension);
    let solution = match method {
        Methods::Bfs => bfs_solve(maze),
        Methods::DeadEnd => dead_end_solve(maze)
    };

    println!("{}", solution);
}

pub enum Methods { Bfs, DeadEnd }

pub enum Dir { North, South, East, West }

/// A coord in the maze
#[derive(Copy, Clone, Debug, PartialEq)]
struct Coord { x: i32, y: i32 }

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

impl From<(i32, i32)> for Coord {
    fn from(value: (i32, i32)) -> Self {
        Self { x: value.0, y: value.1 }
    }
}

impl From<(usize, usize)> for Coord {
    fn from(value: (usize, usize)) -> Self {
        Self { x: value.0 as i32, y: value.1 as i32 }
    }
}

impl Coord {
    fn n(&self) -> Coord { (self.x, self.y - 1).into() }
    fn s(&self) -> Coord { (self.x, self.y + 1).into() }
    fn e(&self) -> Coord { (self.x + 1, self.y).into() }
    fn w(&self) -> Coord { (self.x - 1, self.y).into() }
    fn neighbors(&self) -> [(u8, Coord); 4] {
        [(1, self.n()), (2, self.s()), (4, self.e()), (8, self.w())]
    }
}

/// The list of places we might conceivably expand the maze to
#[derive(Clone, Debug)]
struct Frontier(pub Vec<Coord>);

impl Frontier {
    fn new() -> Self {
        Self(Vec::new())
    }

    fn add(&mut self, coord: Coord) {
        self.0.push(coord)
    }

    fn random(&mut self) -> Option<Coord> {
        if self.0.len() > 1 {
            let idx = rand::thread_rng().next_u32() as usize % self.0.len();

            let cell = self.0[idx];

            if idx == self.0.len() - 1 {
                self.0.pop();
            } else {
                self.0[idx] = self.0.pop().unwrap();
            }

            Some(cell)
        } else {
            self.0.pop()
        }
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

/// The maze itself, stored in a row-major Vec<u8>
#[derive(Clone)]
struct Maze {
    pub width: usize,
    cells: Vec<u8>
}

impl Index<Coord> for Maze {
    type Output = u8;

    fn index(&self, index: Coord) -> &Self::Output {
        if !self.in_bounds(index) { panic!("Coord {} out of bounds", index) }
        &self.cells[index.x as usize + index.y as usize * self.width]
    }
}

impl IndexMut<Coord> for Maze {
    fn index_mut(&mut self, index: Coord) -> &mut Self::Output {
        if !self.in_bounds(index) { panic!("Coord {} out of bounds", index) }
        &mut self.cells[index.x as usize + index.y as usize * self.width]
    }
}

impl Maze {
    fn new(width: usize) -> Self {
        Self {
            width,
            cells: vec![0; width * width]
        }
    }

    /// Generate a square maze of the given side length
    fn generate(width: usize) -> Self {
        let mut maze = Self::new(width);
        let mut frontier = Frontier::new();

        // Pick a random cell to start with
        let start = rand::thread_rng().next_u32() as usize % (width * width);
        let start: Coord = (start % width, start / width).into();

        // We need this cell to appear "connected," meaning, make it not 0.
        // But there's nothing yet to connect it to! So we flip a bit we don't care about:
        maze[start] = 16;

        // Start by adding all neighbors of that cell to the frontier
        maze.frontierize_neighbors(start, &mut frontier);

        // Now, while we have anything in the frontier:
        while !frontier.is_empty() {
            // Pull a random cell out to visit
            let current = frontier.random().unwrap();

            // Connect it to the maze
            maze.connect_random_neighbor(current);

            // Add its neighbors to the frontier
            maze.frontierize_neighbors(current, &mut frontier);
        }

        maze
    }

    /// Returns whether a given coordinate is a valid spot in the maze
    fn in_bounds(&self, coord: Coord) -> bool {
        coord.x >= 0 && coord.y >= 0 && (coord.x as usize) < self.width && (coord.y as usize) < self.width
    }

    /// Get a cell from the maze, on None if the coord is out of bounds
    fn get(&self, index: Coord) -> Option<u8> {
        if self.in_bounds(index) {
            Some(self.cells[index.x as usize + index.y as usize * self.width])
        } else {
            None
        }
    }

    /// If a cell is next to a connected cell, that is not the current cell,
    /// then it must already have been added to the frontier
    fn in_frontier(&self, coord: Coord, current: Coord) -> bool {
        for (_, c) in coord.neighbors() {
            if c != current && self.in_bounds(c) && self[c] != 0 {
                return true
            }
        }
        false
    }

    /// Add to the frontier all neighbors of this point that are in the maze,
    /// unvisited, and not already in the frontier
    fn frontierize_neighbors(&self, coord: Coord, frontier: &mut Frontier) {
        for (_, c) in coord.neighbors() {
            if self.in_bounds(c) && self[c] == 0 && !self.in_frontier(c, coord) {
                frontier.add(c)
            }
        }
    }

    /// Cells are a bit field: n: 1, s: 2, e: 4, w: 8
    fn connect(&mut self, coord: Coord, dir: Dir, reciprocate: bool) {
        match dir {
            North => self[coord] |= 1,
            South => self[coord] |= 2,
            East => self[coord] |= 4,
            West => self[coord] |= 8
        }

        if reciprocate {
            match dir {
                North => self.connect(coord.n(), South, false),
                South => self.connect(coord.s(), North, false),
                East => self.connect(coord.e(), West, false),
                West => self.connect(coord.w(), East, false),
            }
        }
    }

    /// Disconnects a cell from all its neighbors
    fn disconnect(&mut self, coord: Coord) {
        if self[coord] & 1 != 0 { self[coord.n()] &= !2 }
        if self[coord] & 2 != 0 { self[coord.s()] &= !1 }
        if self[coord] & 4 != 0 { self[coord.e()] &= !8 }
        if self[coord] & 8 != 0 { self[coord.w()] &= !4 }
        self[coord] = 0
    }

    /// Connect this cell to a random neighbor that's already connected to the maze
    fn connect_random_neighbor(&mut self, source: Coord) {
        let mut neighbors = Vec::with_capacity(4);
        for (_, c) in source.neighbors() {
            if let Some(v) = self.get(c) {
                if v > 0 { neighbors.push(c) }
            }
        }

        // This can't ever happen, we can only have gotten to this cell by it being
        // on the frontier
        if neighbors.is_empty() { panic!("Something went horribly wrong...") }

        let idx = rand::thread_rng().next_u32() as usize % neighbors.len();
        let target = neighbors[idx];

        // Now connect target to coord
        if target.y == source.y + 1 {
            self.connect(source, South, true);
        }
        if target.y == source.y - 1 {
            self.connect(source, North, true);
        }
        if target.x == source.x - 1 {
            self.connect(source, West, true);
        }
        if target.x == source.x + 1 {
            self.connect(source, East, true);
        }
    }
}

impl Display for Solution {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let (maze, path) = (&self.maze, &self.path);
        // Write the top row of outer walls
        write!(f, " ")?;
        for _ in 0..maze.width {
            write!(f, "\u{2581}")?;
        }

        for y in 0..maze.width {
            write!(f, "\n\u{2595}")?; // Next line and the left edge

            for x in 0..maze.width {
                let cell = maze[(x, y).into()];
                let ch = match (cell & 2 > 0, cell & 4 > 0) { // connections south, east?
                    (false, false) => '\u{1FB7F}',
                    (false, true) => '\u{2581}',
                    (true, false) => '\u{2595}',
                    (true, true) => ' ',
                };

                if path[x + y * maze.width] {
                    write!(f, "\x1b[42m{}\x1b[0m", ch)?;
                } else {
                    write!(f, "{}", ch)?;
                }
            }
        }

        writeln!(f)
    }
}

/// A "Solution" is a solved maze: a `Maze` along with a vec of which cells are on the
/// shortest path from the top-left to bottom-right. The `path` is a row-major grid the
/// same dimensions as the maze.
#[derive(Clone)]
struct Solution {
    maze: Maze,
    path: Vec<bool>
}

/// Solve the maze by dead end filling: repeatedly fill in any cell with only one neighbor (except
/// the start and goal cells) until all such are removed, and what's left must be the shortest path.
fn dead_end_solve(maze: Maze) -> Solution {
    let mut path = maze.clone();
    loop {
        let mut changed = false;
        // Loop over all cells except the first and last; the top-left and bottom right
        for n in 1..path.cells.len() - 1 {
            if matches!(path.cells[n] % 16, 1 | 2 | 4 | 8) {
                changed = true;
                let coord: Coord = (n % maze.width, n / maze.width).into();
                path.disconnect(coord)
            }
        }
        if !changed { break }
    }

    Solution {
        maze,
        path: path.cells.into_iter().map(|v| v % 16 > 0).collect()
    }
}

/// Solve the maze using breadth-first search
fn bfs_solve(maze: Maze) -> Solution {
    // A list of all the unvisited-but-reachable cells, in the order we found them. Start with the
    // starting cell (hardcoded to the upper left)
    let mut open: VecDeque<Coord> = vec![(0, 0).into()].into();
    // A grid of coords: for each cell, the coord of the neighbor we reached that cell from
    let mut visited: Vec<Option<Coord>> = vec![None; maze.width * maze.width];

    let index_for = |c: Coord| c.x as usize + c.y as usize * maze.width;

    loop {
        // Grab the oldest thing off the open list. This would only ever fail if there's no path,
        // and because we generated the maze, we know there is.
        let current = open.pop_front().unwrap();
        // If we're at the goal (hardcoded to the lower right corner) then we're done.
        if current.x as usize == maze.width - 1 && current.y as usize == maze.width - 1 { break }
        let cell = maze[current];
        // For each neighbor of the current cell: if we're connected, and that neighbor hasn't been
        // visited, then add it to visited and open with the current cell as its coord.
        for (bit, nbr) in current.neighbors() {
            if cell & bit != 0 && visited[index_for(nbr)].is_none() {
                visited[index_for(nbr)] = Some(current);
                open.push_back(nbr)
            }
        }
    }

    // At this point the visited grid has values for a path from start to goal, but probably a lot
    // of other values as well. We need to convert it to a grid of bools, showing the one true path.
    let mut path = vec![false; maze.width * maze.width];

    // Start at the goal and mark that as part of the path.
    let mut current: Coord = (maze.width - 1, maze.width - 1).into();
    path[maze.width * maze.width - 1] = true;

    // Walk backwards along the path, until we reach the start cell. Visited stores, for each cell,
    // which cell we visited it from.
    while current.x != 0 || current.y != 0 {
        current = visited[index_for(current)].unwrap();
        path[index_for(current)] = true
    }

    Solution {
        maze,
        path
    }
}