Day 14, puzzle 2

src/bin/d14p1.rs

@@ -1,7 +1,7 @@
 use aoc22::{day14, util};
 
 pub fn main() {
-    let mut cave = day14::parse_cave(&util::parse_input());
+    let mut cave = day14::parse_cave(&util::parse_input(), false);
 
     println!("Initial cave:");
     cave.print();

src/bin/d14p2.rs

@@ -0,0 +1,17 @@
+use aoc22::{day14, util};
+
+pub fn main() {
+    let mut cave = day14::parse_cave(&util::parse_input(), true);
+
+    println!("Initial cave:");
+    cave.print();
+
+    let mut sand = 0;
+    while cave.drop_sand() {
+        sand += 1;
+    }
+
+    println!("Final cave:");
+    cave.print();
+    println!("Total sand dropped: {}", sand);
+}

src/day14.rs

@@ -34,7 +34,11 @@ impl Cave {
         let mut y = y as isize;
         let mut x = x as isize;
 
-        while y < self.rows as isize && x > 0 && x < self.cols as isize {
+        while y < self.rows as isize
+            && x > 0
+            && x < self.cols as isize
+            && !self.is_blocked(self.sand_source.0 as isize, self.sand_source.1 as isize)
+        {
             if !self.is_blocked(y + 1, x) {
                 y += 1;
             } else if !self.is_blocked(y + 1, x - 1) {
@@ -58,7 +62,7 @@ impl Cave {
     }
 }
 
-pub fn parse_cave(input: &String) -> Cave {
+pub fn parse_cave(input: &String, do_floor: bool) -> Cave {
     let mut paths = HashSet::new();
 
     for line in input.lines() {
@@ -90,13 +94,23 @@ pub fn parse_cave(input: &String) -> Cave {
 
     assert_ne!(paths.len(), 0);
 
-    let (y_min, y_max) = paths.iter().map(|c| c.0).minmax().into_option().unwrap();
-    let (x_min, x_max) = paths.iter().map(|c| c.1).minmax().into_option().unwrap();
+    let (mut y_min, mut y_max) = paths.iter().map(|c| c.0).minmax().into_option().unwrap();
+    let (mut x_min, mut x_max) = paths.iter().map(|c| c.1).minmax().into_option().unwrap();
 
     // Sand source is always at y = 0, x = 500
-    let y_min = y_min.min(SAND_SOURCE.0);
-    let x_min = x_min.min(SAND_SOURCE.1);
-    let x_max = x_max.max(SAND_SOURCE.1);
+    y_min = y_min.min(SAND_SOURCE.0);
+    x_min = x_min.min(SAND_SOURCE.1);
+    x_max = x_max.max(SAND_SOURCE.1);
+
+    if do_floor {
+        y_max += 2;
+        let dy = y_max - y_min;
+        x_min = x_min.min(SAND_SOURCE.1 - dy - 1);
+        x_max = x_max.max(SAND_SOURCE.1 + dy + 1);
+        for x in x_min..x_max {
+            paths.insert((y_max, x));
+        }
+    }
 
     let rows = y_max - y_min + 1;
     let cols = x_max - x_min + 1;