Cleaned up the day 9 code a little bit

src/day9.rs

@@ -0,0 +1,171 @@
+use crate::day_solver::DaySolver;
+use num::Integer;
+
+pub struct Day9 {
+    disk_map: Vec<u8>
+}
+
+impl Day9 {
+
+    pub fn create(input: String) -> Self {
+        Day9 {
+            disk_map: input.chars().map(|c| (c as u8) - b'0').collect()
+        }
+    }
+
+    // fn print_layout(disk_blocks: &Vec<Block>) {
+    //     for block in disk_blocks {
+    //         match block {
+    //             Block::File(size, idx) => {
+    //                 for _ in 0..*size {
+    //                     print!("{}", idx);
+    //                 }
+    //             }
+    //             Block::Empty(size) => {
+    //                 print!("{}", ".".repeat(*size));
+    //             }
+    //         }
+    //     }
+    //     println!();
+    // }
+}
+
+impl DaySolver for Day9 {
+
+
+    fn solve_part1(&mut self) -> String {
+
+        let mut cur_disk_idx = 0usize;
+        let mut disk_map_idx = 0usize;
+        let mut disk_map_idx_rev = self.disk_map.len() + (self.disk_map.len() % 2);
+        let mut left_to_move = 0u8;
+
+        let mut res = 0usize;
+
+        while disk_map_idx <= disk_map_idx_rev {
+            if disk_map_idx.is_even() {
+                // We're reading a file from the front of the disk map
+                let left_in_file = if disk_map_idx == disk_map_idx_rev { left_to_move } else { self.disk_map[disk_map_idx] };
+                for _ in 0..left_in_file {
+                    // print!("{}({})", disk_map_idx / 2, cur_disk_idx);
+                    res += cur_disk_idx * (disk_map_idx / 2);
+                    cur_disk_idx += 1;
+                }
+            } else {
+                // We're reading empty space
+                for _ in 0..self.disk_map[disk_map_idx] {
+                    // So let's grab some blocks from files at the end of the disk map
+                    while left_to_move == 0 &&  disk_map_idx_rev > disk_map_idx {
+                        disk_map_idx_rev -= 2;
+                        left_to_move = self.disk_map[disk_map_idx_rev];
+                    }
+                    if disk_map_idx_rev < disk_map_idx {
+                        break
+                    }
+                    // print!("{}({})", disk_map_idx_rev / 2, cur_disk_idx);
+                    res += cur_disk_idx * (disk_map_idx_rev / 2);
+                    cur_disk_idx += 1;
+                    left_to_move -= 1;
+                }
+            }
+            disk_map_idx += 1;
+        }
+
+        res.to_string()
+    }
+
+    fn solve_part2(&mut self) -> String {
+
+        // Somewhat nicer data structure
+        let mut disk_blocks: Vec<Block> = Vec::new();
+        for (i, b) in self.disk_map.iter().enumerate() {
+            if i.is_even() {
+                disk_blocks.push(Block::File(b.to_owned().into(), i / 2))
+            } else if b > &0 {
+                disk_blocks.push(Block::Empty(b.to_owned().into()))
+            }
+        }
+
+        // Self::print_layout(&disk_blocks);
+
+        let original_disk = disk_blocks.clone();
+        for block in original_disk.iter().rev() {
+            if let Block::File(size, idx) = block {
+                // Self::print_layout(&disk_blocks);
+                // println!("{:?}", disk_blocks);
+                // Try to move the file
+                let move_target = disk_blocks.iter().enumerate()
+                    .find(|(_, b) | if let Block::Empty(s) = b { s >= size } else { false })
+                    // To prevent double borrowing, we need to clone the block here
+                    .map(|(i, b)| (i, b.to_owned()));
+                if let Some((can_move_to_idx, b)) = move_target.to_owned() {
+                    match b {
+                        Block::Empty(empty_size) => {
+                            let original_index = disk_blocks.iter().position(|b| b == block).unwrap();
+                            // Never move blocks to the right
+                            if can_move_to_idx >= original_index { continue; }
+
+                            // We replace the original block by an empty one
+                            // Note that, because we never move blocks to the right, it doesn't matter that the empty space around the block we're moving is inaccurate.
+                            disk_blocks[original_index] = Block::Empty(size.to_owned());
+                            disk_blocks[can_move_to_idx] = Block::File(size.to_owned(), idx.to_owned());
+                            if &empty_size > size {
+                                // Add additional empty space
+                                disk_blocks.insert(can_move_to_idx + 1, Block::Empty(empty_size.to_owned() - size.to_owned()))
+                            }
+                        }
+                        _ => panic!("uhm?")
+                    }
+                }
+            }
+        }
+
+        // Self::print_layout(&disk_blocks);
+        let mut checksum = 0usize;
+        let mut cur_disk_idx = 0usize;
+        for block in disk_blocks {
+            match block {
+                Block::File(size, idx) => {
+                    for _ in 0..size {
+                        checksum += cur_disk_idx * idx;
+                        cur_disk_idx += 1;
+                    }
+                }
+                Block::Empty(size) => {
+                    cur_disk_idx += size;
+                }
+            }
+        }
+        checksum.to_string()
+    }
+}
+
+#[derive(Copy, Clone, PartialEq, Debug)]
+enum Block {
+    Empty(usize),
+    File(usize, usize)
+}
+
+#[test]
+fn test_part1_basic() {
+    let mut day = Day9::create("12345".to_string());
+    assert_eq!("60", day.solve_part1());
+}
+
+#[test]
+fn test_part1() {
+    let mut day = Day9::create("2333133121414131402".to_string());
+    assert_eq!("1928", day.solve_part1());
+}
+
+#[test]
+fn test_part2() {
+    let mut day = Day9::create("2333133121414131402".to_string());
+    assert_eq!("2858", day.solve_part2());
+}
+
+#[test]
+fn test_part2_bonus() {
+    let mut day = Day9::create("2931514".to_string());
+    assert_eq!("180", day.solve_part2());
+}

src/main.rs

@@ -9,6 +9,7 @@ use crate::day5::Day5;
 use crate::day6::Day6;
 use crate::day7::Day7;
 use crate::day8::Day8;
+use crate::day9::Day9;
 use crate::day_solver::DaySolver;
 use crate::util::read_file;
 
@@ -22,6 +23,7 @@ mod day5;
 mod day6;
 mod day7;
 mod day8;
+mod day9;
 
 const DEFAULT_BENCHMARK_AMOUNT: u32 = 100;
 
@@ -101,7 +103,7 @@ fn build_day_solver(day: u8, input: String) -> Option<Box<dyn DaySolver>> {
         6 => Some(Box::new(Day6::create(input))),
         7 => Some(Box::new(Day7::create(input))),
         8 => Some(Box::new(Day8::create(input))),
-        // 9 => Some(Box::new(Day9::create(input))),
+        9 => Some(Box::new(Day9::create(input))),
         // 10 => Some(Box::new(Day10::create(input))),
         // 11 => Some(Box::new(Day11::create(input))),
         // 12 => Some(Box::new(Day12::create(input))),
@@ -188,7 +190,7 @@ fn solve_all(silent: bool) -> AocBenchResult {
     for day in 1..(max_day + 1) {
         bench_results.push(solve(day, silent));
     }
-    return AocBenchResult {
+    AocBenchResult {
         read_file: bench_results.iter().map(|t| t.read_file).sum(),
         init: bench_results.iter().map(|t| t.init).sum(),
         part1: bench_results.iter().map(|t| t.part1).sum(),