chore: refactor for cargo-aoc / codspeed

src/day9.rs

@@ -0,0 +1,166 @@
+use aoc_runner_derive::{aoc, aoc_generator};
+use itertools::Itertools;
+use std::fmt::{Display, Write};
+use std::io::{BufRead, Lines};
+
+#[aoc_generator(day9)]
+pub fn get_input(input: &[u8]) -> DiskMap {
+    DiskMap::from(input.lines())
+}
+
+#[derive(Debug, Eq, PartialEq, Clone, Copy)]
+enum Unit {
+    File(usize),
+    Free,
+}
+
+#[derive(Clone)]
+struct Inode {
+    id: usize,
+    pos: usize,
+    len: u8,
+}
+
+#[derive(Clone)]
+pub struct DiskMap {
+    map: Vec<Unit>,
+    files: Vec<Inode>,
+    frees: Vec<Inode>,
+}
+
+impl<T: BufRead> From<Lines<T>> for DiskMap {
+    fn from(mut input: Lines<T>) -> Self {
+        let line_s = input.next().unwrap().unwrap();
+        let line = line_s.as_bytes();
+        let mut file_id = 0;
+        let mut map = Vec::new();
+        let mut files = Vec::new();
+        let mut frees = Vec::new();
+        for (i, c) in line.iter().enumerate() {
+            let len = c - b'0';
+            if i % 2 == 0 {
+                // file
+                files.push(Inode {
+                    id: file_id,
+                    pos: map.len(),
+                    len,
+                });
+                for _ in 0..len {
+                    map.push(Unit::File(file_id))
+                }
+                file_id += 1;
+            } else {
+                // free
+                frees.push(Inode {
+                    id: 0,
+                    pos: map.len(),
+                    len,
+                });
+                for _ in 0..len {
+                    map.push(Unit::Free)
+                }
+            }
+        }
+        Self { map, files, frees }
+    }
+}
+
+impl Display for DiskMap {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        for unit in &self.map {
+            match unit {
+                Unit::File(i) => f.write_char((b'0' + (*i % 10) as u8) as char)?,
+                Unit::Free => f.write_char('.')?,
+            }
+        }
+        Ok(())
+    }
+}
+
+impl DiskMap {
+    fn checksum(&self) -> u64 {
+        self.map
+            .iter()
+            .enumerate()
+            .map(|(i, u)| match u {
+                Unit::File(id) => i * id,
+                Unit::Free => 0,
+            })
+            .sum::<usize>() as u64
+    }
+}
+
+// PROBLEM 1 solution
+#[aoc(day9, part1)]
+pub fn part1(map: &DiskMap) -> u64 {
+    let mut map = map.to_owned();
+    let mut last_free = 0;
+    for file in map.files.iter().rev() {
+        let frees = map
+            .map
+            .iter()
+            .enumerate()
+            .skip(last_free) // we greedy fill, so no need to check for free space before the last one we used
+            .take(file.pos + file.len as usize) // and we only need to search until the end of the current file
+            .filter(|(_i, u)| **u == Unit::Free || **u == Unit::File(file.id)) // look for free space or our existing space
+            .map(|(i, _u)| i)
+            .take(file.len as usize) // get the first file.len free blocks
+            .collect_vec();
+        // Note: no need to test for too small frees list here, since we are guaranteed at worst to find our current position
+        if frees[0] >= file.pos {
+            // if the first available free is > file.pos, it's fully packed, job done
+            break;
+        }
+        #[allow(clippy::needless_range_loop)]
+        for j in 0..file.len as usize {
+            map.map.swap(frees[j], file.pos + j);
+        }
+        last_free = frees[file.len as usize - 1]
+    }
+    map.checksum()
+}
+
+// PROBLEM 2 solution
+#[aoc(day9, part2)]
+pub fn part2(map: &DiskMap) -> u64 {
+    let mut map = map.to_owned();
+    for file in map.files.iter().rev() {
+        let free = map.frees.iter_mut().find(|inode| inode.len >= file.len); // find the first entry in the free space map large enough
+        if let Some(free) = free {
+            if free.pos >= file.pos {
+                // if it's past our position, continue, but can't break since there might be free space for future files
+                continue;
+            }
+            for j in 0..file.len {
+                map.map.swap(free.pos + j as usize, file.pos + j as usize);
+            }
+            // Note: It is slightly faster to keep these hanging around in the free map with size = 0 then to remove them from the vec
+            free.len -= file.len;
+            free.pos += file.len as usize;
+
+            map.frees.push(Inode {
+                id: 0,
+                pos: file.pos,
+                len: file.len,
+            });
+        }
+    }
+    map.checksum()
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::day9::*;
+
+    const EXAMPLE: &[u8] = b"2333133121414131402";
+
+    #[test]
+    fn part1_example() {
+        assert_eq!(part1(&get_input(EXAMPLE)), 1928);
+    }
+
+    #[test]
+    fn part2_example() {
+        assert_eq!(part2(&get_input(EXAMPLE)), 2858);
+    }
+}