move main loop to state

3 files changed, 25 insertions, 19 deletions

diff --git a/src/main.rs b/src/main.rs
index 1149f83..21b6465 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,10 +5,7 @@ use crate::state::State;
 use crate::state::Threads;
 
 use std::io::{stdin, stdout};
-use termion::event::Key;
-use termion::input::TermRead;
 use termion::raw::IntoRawMode;
-use termion::screen::AlternateScreen;
 
 fn main() {
     let stdin = stdin();
@@ -31,19 +28,6 @@ fn main() {
 
     threads.init(&mut screen);
 
-    let mut state = State::Threads(threads);
-
-    for c in stdin.keys() {
-        let c = c.unwrap();
-        // Global keybinds
-        match c {
-            Key::Char('q') => break,
-            _ => ()
-        }
-
-        // Pass to current state
-        state = match state {
-            State::Threads(s) => s.tick(&mut screen, c),
-        };
-    }
+    let state = State::Threads(threads);
+    state.run(screen, stdin);
 }
diff --git a/src/state/mod.rs b/src/state/mod.rs
index 5313b4f..7b498f2 100644
--- a/src/state/mod.rs
+++ b/src/state/mod.rs
@@ -1,7 +1,29 @@
 mod threads;
 
+use std::io::{Stdin, Write};
+use termion::event::Key;
+use termion::input::TermRead;
+
 pub use threads::Threads;
 
 pub enum State {
     Threads(Threads),
 }
+
+impl State {
+    pub fn run<W: Write>(mut self, mut screen: W, stdin: Stdin) {
+        for c in stdin.keys() {
+            let c = c.unwrap();
+            // Global keybinds
+            match c {
+                Key::Char('q') => break,
+                _ => ()
+            }
+
+            // Pass to current state
+            self = match self {
+                State::Threads(s) => s.tick(&mut screen, c),
+            };
+        }
+    }
+}
diff --git a/src/state/threads.rs b/src/state/threads.rs
index 78dc795..6ffd7e7 100644
--- a/src/state/threads.rs
+++ b/src/state/threads.rs
@@ -85,8 +85,8 @@ impl<'d, 'q> Threads {
         }
         if !self.threads.is_empty() {
             self.i = self.i.rem_euclid(self.threads.len());
-            draw(&self, out);
         }
+        draw(&self, out);
         State::Threads(self)
     }
 }