diff --git a/src/day4/main.rs b/src/day4/main.rs
index d005e6e..f4148b2 100644
--- a/src/day4/main.rs
+++ b/src/day4/main.rs
@@ -8,47 +8,7 @@ enum Position {
     CanMove,
 }
 
-trait NicePrint {
-    fn nice_print(&self);
-}
-
-impl NicePrint for Vec<Vec<Position>> {
-    fn nice_print(&self) {
-        self.iter().for_each(|a| {
-            a.iter().for_each(|p| match *p {
-                Position::Roll => print!("@"),
-                Position::Empty => print!("."),
-                Position::CanMove => print!("x"),
-            });
-            print!("\n")
-        });
-    }
-}
-
-fn sol1(text: &str) -> u64 {
-    let pos: Vec<Vec<Position>> =
-        text.split("\n")
-            .filter(|s| s.len() > 0)
-            .fold(Vec::new(), |mut sum, sline| {
-                sum.push(
-                    sline
-                        .chars()
-                        .map(|a| {
-                            if a == '@' {
-                                Position::Roll
-                            } else {
-                                Position::Empty
-                            }
-                        })
-                        .collect(),
-                );
-                sum
-            });
-    let (sum, _) = sol(pos);
-    sum
-}
-
-fn sol(pos: Vec<Vec<Position>>) -> (u64, Vec<Vec<Position>>) {
+fn step(pos: Vec<Vec<Position>>) -> (u64, Vec<Vec<Position>>) {
     let y_max = pos.len() as i64;
     let x_max = pos[0].len() as i64;
     let mut new_pos: Vec<Vec<Position>> = pos.iter().cloned().collect();
@@ -79,12 +39,10 @@ fn sol(pos: Vec<Vec<Position>>) -> (u64, Vec<Vec<Position>>) {
                 }
             })
     });
-
-    //new_pos.nice_print();
     (new_sum, new_pos)
 }
 
-fn sol2(text: &String) -> u64 {
+fn sol(text: &String, stop_after: Option<usize>) -> u64 {
     let mut pos: Vec<Vec<Position>> =
         text.split("\n")
             .filter(|s| s.len() > 0)
@@ -104,9 +62,14 @@ fn sol2(text: &String) -> u64 {
                 sum
             });
     let mut count = 0;
+    let mut step_num = 0;
     loop {
-        let (c, new_pos) = sol(pos);
-        println!("Removed {} rolls ", c);
+        if let Some(stop_after) = stop_after
+            && stop_after == step_num
+        {
+            break;
+        }
+        let (c, new_pos) = step(pos);
         if c == 0 {
             break;
         }
@@ -125,6 +88,7 @@ fn sol2(text: &String) -> u64 {
                     .collect()
             })
             .collect();
+        step_num += 1;
     }
     count
 }
@@ -132,8 +96,7 @@ fn sol2(text: &String) -> u64 {
 fn main() -> Result<()> {
     let text = fs::read_to_string("src/day4/act.txt")?;
 
-    println!("sol1: {}", sol1(&text));
-    println!("sol2: {}", sol2(&text));
+    println!("sol1: {}; sol2: {}", sol(&text, Some(1)), sol(&text, None));
 
     Ok(())
 }