From c9dac158eedab784e18832c34482917d3dc2f0dd Mon Sep 17 00:00:00 2001
From: Michael Zhang <iptq@protonmail.com>
Date: Mon, 8 Jun 2020 11:33:47 -0500
Subject: [PATCH] sleep

---
 src/main.rs         |  1 +
 src/screens/play.rs | 43 ++++++++++++++++++++++++++++++-------------
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 768f192..c9cd31e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -70,5 +70,6 @@ fn main() {
         target.finish().unwrap();
 
         prev = now;
+        std::thread::sleep(std::time::Duration::from_millis(17));
     }
 }
diff --git a/src/screens/play.rs b/src/screens/play.rs
index 0f75b64..7968bd1 100644
--- a/src/screens/play.rs
+++ b/src/screens/play.rs
@@ -15,8 +15,9 @@ const LEVEL_1: &str = include_str!("../../levels/level1.json");
 
 pub struct PlayScreen {
     animations: AnimationState,
-    levels: Vec<Level>,
-    current_level: usize,
+    levels: Vec<&'static str>,
+    current_level: Level,
+    current_level_idx: usize,
 }
 
 impl Screen for PlayScreen {
@@ -61,6 +62,11 @@ impl Screen for PlayScreen {
             btn_handler!(VirtualKeyCode::J, Board::Right, PushDir::Left);
             btn_handler!(VirtualKeyCode::K, Board::Right, PushDir::Down);
             btn_handler!(VirtualKeyCode::L, Board::Right, PushDir::Right);
+
+            if keymap.is_pressed(VirtualKeyCode::R) {
+                // restart the level
+                self.restart_level();
+            }
         }
 
         ScreenAction::None
@@ -74,34 +80,45 @@ impl Screen for PlayScreen {
 
 impl PlayScreen {
     pub fn get_current_level(&self) -> &Level {
-        self.levels.get(self.current_level).unwrap()
+        &self.current_level
     }
 
     pub fn get_current_level_mut(&mut self) -> &mut Level {
-        self.levels.get_mut(self.current_level).unwrap()
+        &mut self.current_level
     }
 
     pub fn new() -> PlayScreen {
         let levels = vec![
-            Level::from_json(&LEVEL_TUTORIAL),
-            Level::from_json(&LEVEL_TUTORIAL2),
-            Level::from_json(&LEVEL_1),
+            LEVEL_TUTORIAL,
+            LEVEL_TUTORIAL2,
+            LEVEL_1,
         ];
 
         PlayScreen {
             levels,
-            current_level: 0,
+            current_level: Level::from_json(&LEVEL_TUTORIAL),
+            current_level_idx: 0,
             animations: AnimationState::new(),
         }
     }
 
-    fn go_to_next_level(&mut self) {
-        self.current_level += 1;
+    fn restart_level(&mut self) {
+        self.switch_to_level(self.current_level_idx);
+    }
 
+    fn switch_to_level(&mut self, idx: usize) {
+        self.current_level = Level::from_json(self.levels[idx]);
+        self.current_level_idx = idx;
+    }
+
+    fn go_to_next_level(&mut self) {
         // TODO: make an actual win screen
-        if self.current_level >= self.levels.len() {
-            self.current_level = 0;
-        }
+        let next_level = if self.current_level_idx + 1 >= self.levels.len() {
+            0
+        } else {
+            self.current_level_idx + 1
+        };
+        self.switch_to_level(next_level);
     }
 
     fn check_win_condition(&mut self) {