diff --git a/src/day_02.rs b/src/day_02.rs
index 83e3239..fc9f157 100644
--- a/src/day_02.rs
+++ b/src/day_02.rs
@@ -1,6 +1,9 @@
 use std::str::FromStr;
 
 use aoc_runner_derive::{aoc, aoc_generator};
+use yap::{IntoTokens, Tokens};
+
+use crate::parsing::parse_number;
 
 #[aoc_generator(day2)]
 pub fn parse_movements(input: &str) -> Vec<Movement> {
@@ -12,35 +15,37 @@ pub fn parse_movements(input: &str) -> Vec<Movement> {
 
 #[aoc(day2, part1)]
 pub fn part1(input: &[Movement]) -> isize {
-    let mut state = NaiveState::default();
-    input.iter().for_each(|movement| state.movement(movement));
-    state.depth * state.horizontal_pos
+    NaiveState::default().movements(input)
 }
 
 #[aoc(day2, part2)]
 pub fn part2(input: &[Movement]) -> isize {
-    let mut state = State::default();
-    input.iter().for_each(|movement| state.movement(movement));
-    state.depth * state.horizontal_pos
+    CorrectState::default().movements(input)
 }
 
-pub struct State {
+pub trait State {
+    fn movement(&mut self, movement: &Movement);
+    fn depth(&self) -> isize;
+    fn horizontal_pos(&self) -> isize;
+    fn result(&self) -> isize {
+        self.depth() * self.horizontal_pos()
+    }
+    fn movements(&mut self, moves: &[Movement]) -> isize {
+        for movement in moves {
+            self.movement(&movement);
+        }
+        self.result()
+    }
+}
+
+#[derive(Default)]
+pub struct CorrectState {
     horizontal_pos: isize,
     depth: isize,
     aim: isize,
 }
 
-impl Default for State {
-    fn default() -> State {
-        State {
-            horizontal_pos: 0,
-            depth: 0,
-            aim: 0,
-        }
-    }
-}
-
-impl State {
+impl State for CorrectState {
     fn movement(&mut self, movement: &Movement) {
         match movement {
             Movement::Forward(units) => {
@@ -51,23 +56,23 @@ impl State {
             Movement::Up(units) => self.aim -= units,
         }
     }
+
+    fn depth(&self) -> isize {
+        self.depth
+    }
+
+    fn horizontal_pos(&self) -> isize {
+        self.horizontal_pos
+    }
 }
 
+#[derive(Default)]
 pub struct NaiveState {
     horizontal_pos: isize,
     depth: isize,
 }
 
-impl Default for NaiveState {
-    fn default() -> NaiveState {
-        NaiveState {
-            horizontal_pos: 0,
-            depth: 0,
-        }
-    }
-}
-
-impl NaiveState {
+impl State for NaiveState {
     fn movement(&mut self, movement: &Movement) {
         match movement {
             Movement::Forward(units) => self.horizontal_pos += units,
@@ -75,6 +80,14 @@ impl NaiveState {
             Movement::Up(units) => self.depth -= units,
         }
     }
+
+    fn depth(&self) -> isize {
+        self.depth
+    }
+
+    fn horizontal_pos(&self) -> isize {
+        self.horizontal_pos
+    }
 }
 
 pub enum Movement {
@@ -86,15 +99,12 @@ pub enum Movement {
 impl FromStr for Movement {
     type Err = ();
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        if let Some(units) = s.strip_prefix("forward ") {
-            Ok(Self::Forward(units.parse().unwrap()))
-        } else if let Some(units) = s.strip_prefix("down ") {
-            Ok(Self::Down(units.parse().unwrap()))
-        } else if let Some(units) = s.strip_prefix("up ") {
-            Ok(Self::Up(units.parse().unwrap()))
-        } else {
-            Err(())
-        }
+        let mut tokens = s.into_tokens();
+        yap::one_of!(ts from &mut tokens;
+            ts.tokens("forward ".chars()).then(|| parse_number(ts).map(|num| Movement::Forward(num))).flatten(),
+            ts.tokens("down ".chars()).then(|| parse_number(ts).map(|num| Movement::Down(num))).flatten(),
+            ts.tokens("up ".chars()).then(|| parse_number(ts).map(|num| Movement::Up(num))).flatten(),
+            ).ok_or(())
     }
 }