From 2ac9103fe2d4c51e3dcd3bcd72434d14679fae4f Mon Sep 17 00:00:00 2001
From: Matej Focko <mfocko@redhat.com>
Date: Wed, 14 Dec 2022 15:47:56 +0100
Subject: [PATCH] day(14): refactor

Signed-off-by: Matej Focko <mfocko@redhat.com>
---
 src/bin/day14.rs | 89 +++++++++++++++++++++++-------------------------
 1 file changed, 42 insertions(+), 47 deletions(-)

diff --git a/src/bin/day14.rs b/src/bin/day14.rs
index 8c47a74..ce80891 100644
--- a/src/bin/day14.rs
+++ b/src/bin/day14.rs
@@ -103,6 +103,35 @@ fn can_move(grain: &Coord, lines: &[Line], set: &BTreeSet<Coord>) -> Option<Coor
     None
 }
 
+fn simulate_falling<F>(lines: &Input, terminate: F) -> usize
+where  
+    F: Fn(&Coord, &Input, &BTreeSet<Coord>, bool) -> bool
+{
+    let mut grains: VecDeque<Coord> = VecDeque::new();
+    let mut set_grains: BTreeSet<Coord> = BTreeSet::new();
+
+    loop {
+        for i in (0..grains.len()).rev() {
+            let grain = grains[i];
+
+            let mut is_set = false;
+            if let Some(new_grain) = can_move(&grain, lines, &set_grains) {
+                grains[i] = new_grain;
+            } else {
+                let grain = grains.pop_back().unwrap();
+                set_grains.insert(grain);
+                is_set = true;
+            }
+
+            if terminate(&grain, lines, &set_grains, is_set) {
+                return set_grains.len();
+            }
+        }
+
+        grains.push_front(Coord::new(500, 0));
+    }
+}
+
 struct Day14;
 impl Solution<Input, Output> for Day14 {
     fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
@@ -113,65 +142,31 @@ impl Solution<Input, Output> for Day14 {
     }
 
     fn part_1(input: &Input) -> Output {
-        let mut grains: VecDeque<Coord> = VecDeque::new();
-        let mut set_grains: BTreeSet<Coord> = BTreeSet::new();
-
-        loop {
-            for i in (0..grains.len()).rev() {
-                let grain = &mut grains[i];
-
-                if free_falling(grain, input, &set_grains) {
-                    return set_grains.len();
-                }
-
-                if let Some(new_grain) = can_move(grain, input, &set_grains) {
-                    *grain = new_grain;
-                } else {
-                    let grain = grains.pop_back().unwrap();
-                    set_grains.insert(grain);
-                }
-            }
-
-            grains.push_front(Coord::new(500, 0));
+        fn terminate(grain: &Coord, lines: &Input, set_grains: &BTreeSet<Coord>, _is_set: bool) -> bool {
+            free_falling(grain, lines, set_grains)
         }
+
+        simulate_falling(input, terminate)
     }
 
     fn part_2(input: &Input) -> Output {
-        let mut lines = input.clone();
-        let y = lines
+        fn terminate(grain: &Coord, _lines: &Input, _set_grains: &BTreeSet<Coord>, is_set: bool) -> bool {
+            is_set && grain == &Coord::new(500, 0)
+        }
+
+        let y = input
             .iter()
             .map(|l| max(l.from.y(), l.to.y()))
             .max()
             .unwrap();
 
+        let mut lines = input.clone();
         lines.push(Line {
-            from: Coord::new(500 - 2 * (*y + 2), *y + 2),
-            to: Coord::new(500 + 2 * (*y + 2), *y + 2),
+            from: Coord::new(500 - (*y + 2), *y + 2),
+            to: Coord::new(500 + (*y + 2), *y + 2),
         });
 
-        let mut grains: VecDeque<Coord> = VecDeque::new();
-        let mut set_grains: BTreeSet<Coord> = BTreeSet::new();
-
-        loop {
-            // debug!("Grains setting: {:?}, set grains: {:?}", grains.len(), set_grains.len());
-            for i in (0..grains.len()).rev() {
-                let grain = &mut grains[i];
-
-                if let Some(new_grain) = can_move(grain, &lines, &set_grains) {
-                    *grain = new_grain;
-                } else {
-                    let grain = grains.pop_back().unwrap();
-                    set_grains.insert(grain);
-
-                    // debug!("Set grain {:?}", grain);
-                    if grain == Coord::new(500, 0) {
-                        return set_grains.len();
-                    }
-                }
-            }
-
-            grains.push_front(Coord::new(500, 0));
-        }
+        simulate_falling(&lines, terminate)
     }
 }