un-nest Block enum

day09/src/main.rs

@@ -1,6 +1,12 @@
+use std::{
+    collections::HashSet,
+    fmt::{Debug, Display, Write},
+};
+
 fn main() {
     let input = std::fs::read_to_string("input").unwrap();
     println!("{}", pt1(&input));
+    println!("{}", pt2(&input));
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -10,7 +16,6 @@ enum Cell {
 }
 
 fn pt1(input: &str) -> usize {
-    let mut next_id = 0;
     let mut map = Vec::new();
     let mut frees = Vec::new();
     for (i, c) in input.chars().enumerate() {
@@ -21,11 +26,10 @@ fn pt1(input: &str) -> usize {
             }
         };
         if i % 2 == 0 {
-            let cell = Cell::Block(next_id);
+            let cell = Cell::Block(i / 2);
             for _ in 0..n {
                 map.push(cell);
             }
-            next_id += 1;
         } else {
             for _ in 0..n {
                 frees.push(map.len());
@@ -47,6 +51,88 @@ fn pt1(input: &str) -> usize {
         map.swap(free, i);
     }
 
+    map.iter()
+        .enumerate()
+        .map(|(i, &b)| match b {
+            Cell::Block(v) => i * v,
+            _ => 0,
+        })
+        .sum()
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+enum Block {
+    Free { size: usize },
+    File { size: usize, id: usize },
+}
+
+impl Block {
+    fn is_free(&self) -> Option<usize> {
+        match self {
+            Block::Free { size } => Some(*size),
+            _ => None,
+        }
+    }
+
+    fn size(&self) -> usize {
+        match self {
+            Block::Free { size } => *size,
+            Block::File { size, .. } => *size,
+        }
+    }
+
+    fn set_size(&mut self, s: usize) {
+        match self {
+            Block::Free { size } => *size = s,
+            Block::File { size, .. } => *size = s,
+        }
+    }
+}
+
+fn pt2(input: &str) -> usize {
+    let mut map = Vec::new();
+
+    for (i, c) in input.chars().enumerate() {
+        let n = match c.to_digit(10) {
+            Some(n) => n as usize,
+            _ => {
+                continue;
+            }
+        };
+        if i % 2 == 0 {
+            let cell = Block::File { size: n, id: i / 2 };
+            map.push(cell);
+        } else {
+            map.push(Block::Free { size: n });
+        }
+    }
+
+    let mut processed = HashSet::new();
+    let mut done = false;
+    while !done {
+        for i in (0..map.len()).rev() {
+            if i == 0 {
+                done = true;
+            }
+            match map[i] {
+                Block::Free { .. } => {
+                    continue;
+                }
+                Block::File { id, .. } => {
+                    if processed.contains(&id) {
+                        continue;
+                    }
+                    processed.insert(id);
+                    if move_file(&mut map, i) {
+                        map = compact(&map);
+                        break;
+                    }
+                }
+            }
+        }
+    }
+    let map = render_map(&map);
+
     //println!("{map:?}");
 
     map.iter()
@@ -58,14 +144,102 @@ fn pt1(input: &str) -> usize {
         .sum()
 }
 
-enum Block {
-    Free(usize),
-    // size, id
-    Block(usize, usize),
+fn move_file(map: &mut Vec<Block>, file_index: usize) -> bool {
+    let file_size = map[file_index].size();
+    for j in 0..file_index {
+        if let Some(free_size) = map[j].is_free() {
+            if map[file_index].size() <= free_size {
+                map.swap(file_index, j);
+                let diff = free_size - file_size;
+                if diff > 0 {
+                    map[file_index].set_size(file_size);
+                    map.insert(j + 1, Block::Free { size: diff });
+                    return true;
+                }
+                break;
+            }
+        }
+    }
+    false
 }
 
-fn pt2(input: &str) -> usize {
-    todo!()
+fn render_map(map: &[Block]) -> Vec<Cell> {
+    let mut out = Vec::new();
+
+    for block in map {
+        match block {
+            Block::Free { size } => {
+                for _ in 0..*size {
+                    out.push(Cell::Free);
+                }
+            }
+            Block::File { size, id } => {
+                for _ in 0..*size {
+                    out.push(Cell::Block(*id));
+                }
+            }
+        }
+    }
+
+    out
+}
+
+impl Display for Block {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Block::Free { size: s } => {
+                let out = (0..*s).fold(String::new(), |mut out, _| {
+                    let _ = write!(out, ".");
+                    out
+                });
+                write!(f, "{out}")
+            }
+            Block::File { size: sz, id } => {
+                let out = (0..*sz).fold(String::new(), |mut out, _| {
+                    if out.is_empty() {
+                        let _ = write!(out, "{id}");
+                    } else {
+                        let _ = write!(out, "|{id}");
+                    }
+                    out
+                });
+                write!(f, "{out}")
+            }
+        }
+    }
+}
+
+impl Debug for Block {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Block::Free { size } => write!(f, "free: {}", size),
+            Block::File { size, id } => write!(f, "{}: {}", id, size),
+        }
+    }
+}
+
+fn compact(v: &[Block]) -> Vec<Block> {
+    let mut out = Vec::new();
+    let mut i = 0;
+
+    while i < v.len() {
+        match v[i] {
+            Block::File { size, id } => out.push(Block::File { size, id }),
+            Block::Free { size } => {
+                let mut total = size;
+                while let Some(next_free_size) = v.get(i + 1).and_then(|b| b.is_free()) {
+                    total += next_free_size;
+                    i += 1;
+                }
+                if total > 0 {
+                    out.push(Block::Free { size: total });
+                }
+            }
+        }
+        i += 1;
+    }
+
+    out
 }
 
 #[cfg(test)]
@@ -79,4 +253,9 @@ mod test {
         let v = pt1(INPUT);
         assert_eq!(v, 1928);
     }
+    #[test]
+    fn p2() {
+        let v = pt2(INPUT);
+        assert_eq!(v, 2858);
+    }
 }