1 files changed, 104 insertions, 22 deletions

diff --git a/src/main.rs b/src/main.rs
index b8705c7..d5d2ad1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,6 +15,7 @@ impl fmt::Display for Cell {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         if self.flagged {
             write!(f, "F")?;
+            return Ok(());
         }
 
         if !self.revealed {
@@ -69,8 +70,12 @@ impl Grid {
 
         let mut mines_left = mines;
         while mines_left > 0 {
-            let row = rng.gen_range(0..self.0.len());
-            let col = rng.gen_range(0..self.0[0].len());
+            let height = self.0.len();
+            let width = self.0[0].len();
+
+            let row = rng.gen_range(0..height);
+            let col = rng.gen_range(0..width);
+
             let cell: &mut Cell = &mut self.0[row][col];
             if cell.mine {
                 continue;
@@ -82,6 +87,25 @@ impl Grid {
         }
     }
 
+    fn on_neighbors<F>(&mut self, x: usize, y: usize, mut f: F)
+    where
+        F: FnMut(&mut Cell, (usize, usize)),
+    {
+        for row in y.saturating_sub(1)..=y.saturating_add(1) {
+            for col in x.saturating_sub(1)..=x.saturating_add(1) {
+                if x == col && y == row {
+                    continue;
+                }
+
+                let Some(cell) = self.0.get_mut(row).and_then(|row| row.get_mut(col)) else {
+                    continue;
+                };
+
+                f(cell, (row, col));
+            }
+        }
+    }
+
     pub fn new(rows: usize, cols: usize, mines: u8) -> Option<Self> {
         if rows < 1 || cols < 1 {
             return None;
@@ -112,7 +136,7 @@ enum GridCommand {
 
 struct GameState {
     grid: Grid,
-    win: bool,
+    play: bool,
     turn: u32,
 }
 
@@ -168,18 +192,33 @@ fn parse_command(cmd: &str) -> Result<GridCommand, &'static str> {
     let coord = parse_coord(coord_str).or(Err("Invalid coordinates"))?;
 
     match command_part.to_lowercase().as_str() {
-        "flag" => Ok(GridCommand::FLAG(coord)),
-        "reveal" => Ok(GridCommand::REVEAL(coord)),
+        "f" | "flag" => Ok(GridCommand::FLAG(coord)),
+        "r" | "reveal" => Ok(GridCommand::REVEAL(coord)),
         _ => Err("Unknown command"),
     }
 }
 
-fn reveal_at(game_state: &mut GameState, coord: (usize, usize)) {
+fn reveal_at(game_state: &mut GameState, coord: (usize, usize)) -> Result<bool, &'static str> {
+    let Some(cell) = game_state
+        .grid
+        .0
+        .get_mut(coord.0)
+        .and_then(|row| row.get_mut(coord.1))
+    else {
+        return Err("invalid coordinates");
+    };
+
     let mut seen = HashSet::new();
     let mut stack = Vec::<(usize, usize)>::new();
     stack.push(coord);
 
+    if cell.mine {
+        println!("that's a mine.");
+        return Ok(false);
+    }
+
     while let Some((y, x)) = stack.pop() {
+        let coord = (y, x);
         if seen.contains(&coord) {
             continue;
         }
@@ -188,45 +227,87 @@ fn reveal_at(game_state: &mut GameState, coord: (usize, usize)) {
         let Some(cell) = game_state.grid.0.get_mut(y).and_then(|row| row.get_mut(x)) else {
             continue;
         };
+        cell.revealed = true;
+
         if cell.neighbors > 0 || cell.mine {
             continue;
         }
 
-        cell.revealed = true;
+        game_state
+            .grid
+            .on_neighbors(x, y, |_cell, coord| stack.push(coord));
+    }
 
-        for row in y.saturating_sub(1)..=y.saturating_add(1) {
-            for col in x.saturating_sub(1)..=x.saturating_add(1) {
-                if x == col && y == row {
-                    continue;
-                }
+    return Ok(true);
+}
 
-                stack.push((row, col));
-            }
-        }
-    }
+fn flag_at(game_state: &mut GameState, coord: (usize, usize)) -> Result<bool, &'static str> {
+    let Some(cell) = game_state
+        .grid
+        .0
+        .get_mut(coord.0)
+        .and_then(|row| row.get_mut(coord.1))
+    else {
+        return Err("invalid coordinates");
+    };
+    cell.flagged = !cell.flagged;
+
+    let won = game_state.grid.0.iter().fold(true, |won, row| {
+        won && row
+            .iter()
+            .fold(won, |won, cell| won && cell.flagged == cell.mine)
+    });
+
+    return Ok(!won);
 }
 
 fn play_game(game_state: &mut GameState) {
-    if game_state.win {
+    if !game_state.play {
         return;
     }
 
     println!("{}\n", game_state);
     print!("> ");
     let _ = io::stdout().flush();
-
     let mut buffer = String::new();
     let stdin = io::stdin();
-    stdin.read_line(&mut buffer);
+    match stdin.read_line(&mut buffer) {
+        Ok(_) => (),
+        Err(e) => {
+            eprintln!("error: {}", e);
+            return play_game(game_state);
+        }
+    }
 
     let cmd = match parse_command(buffer.as_str()) {
         Ok(command) => command,
         Err(e) => {
-            println!("error: {}", e);
+            eprintln!("error: {}", e);
             return play_game(game_state);
         }
     };
 
+    match cmd {
+        GridCommand::REVEAL(coord) => {
+            game_state.play = match reveal_at(game_state, coord) {
+                Ok(continue_play) => continue_play,
+                Err(e) => {
+                    eprintln!("error: {}", e);
+                    return play_game(game_state);
+                }
+            }
+        }
+        GridCommand::FLAG(coord) => {
+            game_state.play = match flag_at(game_state, coord) {
+                Ok(continue_play) => continue_play,
+                Err(e) => {
+                    eprintln!("error: {}", e);
+                    return play_game(game_state);
+                }
+            }
+        }
+    }
+
     game_state.turn += 1;
     play_game(game_state);
 }
@@ -236,11 +317,12 @@ fn main() {
         eprintln!("failed to initialize grid");
         process::exit(1);
     };
-
     let mut game_state = GameState {
         grid,
-        win: false,
+        play: true,
         turn: 1,
     };
     play_game(&mut game_state);
+
+    main();
 }