Prim's Algorithm in Rust

primmaze.rs

use std::env;
use std::fmt::{Display, Formatter};
use std::ops::{Index, IndexMut};
use rand::RngCore; // cargo add rand
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.")
    };
    println!("{}", Maze::generate(dimension));
}

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) -> [Coord; 4] {
        [self.n(), self.s(), self.e(), 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),
            }
        }
    }

    /// 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 Maze {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for y in 0..(self.width) {
            for x in 0..self.width {
                let cell = self[(x, y).into()];
                let ch = match cell % 16 {
                    0 => '*',
                    1 => '\u{2575}', // n
                    2 => '\u{2577}', // s
                    3 => '\u{2502}', // ns
                    4 => '\u{2576}', // e
                    5 => '\u{2514}', // ne
                    6 => '\u{250C}', // se
                    7 => '\u{251C}', // nse
                    8 => '\u{2574}', // w
                    9 => '\u{2518}', // nw
                    10 => '\u{2510}', // sw
                    11 => '\u{2524}', // nsw
                    12 => '\u{2500}', // ew
                    13 => '\u{2534}', // new
                    14 => '\u{252C}', // sew
                    15 => '\u{253C}', // nsew
                    _ => unreachable!()
                };
                write!(f, "{}", ch)?;
            }
            writeln!(f)?;
        }
        Ok(())
    }
}