Remove most allocations

2022/src/day07.rs

@@ -1,5 +1,3 @@
-use std::fmt::Debug;
-
 use anyhow::Context;
 use anyhow::Result;
 use nom::branch::alt;
@@ -8,7 +6,7 @@ use nom::bytes::complete::take_until;
 use nom::character::complete::newline;
 use nom::combinator::map;
 use nom::combinator::opt;
-use nom::multi::many0;
+use nom::multi::fold_many0;
 use nom::sequence::delimited;
 use nom::sequence::preceded;
 use nom::sequence::separated_pair;
@@ -18,21 +16,22 @@ use nom::IResult;
 
 use crate::common::parse_input;
 
-type Listing<'a> = Vec<(&'a [u8], ListEntry<'a>)>;
-type Slicing<'a> = &'a [(&'a [u8], ListEntry<'a>)];
-
-#[derive(Debug)]
-enum ListEntry<'a> {
-    File(u32),
-    Dir(Listing<'a>),
-}
-
-fn parse_dir(input: &[u8]) -> IResult<&[u8], Listing> {
+fn parse_dir<'a>(
+    input: &'a [u8],
+    dirs: &mut Vec<u32>,
+    dir_stack: &mut Vec<&'a [u8]>,
+) -> IResult<&'a [u8], u32> {
     use nom::character::complete::u32;
 
-    let (mut input, mut entries) = preceded(
+    enum Entry<'a> {
+        File(u32),
+        Dir(&'a [u8]),
+    }
+    let initial_len = dir_stack.len();
+
+    let (mut input, mut size) = preceded(
         tag("$ ls\n"),
-        many0(
+        fold_many0(
             // Map many newline-terminated entries
             terminated(
                 // of either
@@ -40,67 +39,57 @@ fn parse_dir(input: &[u8]) -> IResult<&[u8], Listing> {
                     // A size followed by a name
                     map(
                         separated_pair(u32, tag(" "), take_until("\n")),
-                        |(size, name)| (name, ListEntry::File(size)),
+                        |(size, _)| Entry::File(size),
                     ),
                     // Or the word "dir" followed by a name
                     map(preceded(tag("dir "), take_until("\n")), |name| {
-                        (name, ListEntry::Dir(Vec::new()))
+                        Entry::Dir(name)
                     }),
                 )),
                 newline,
             ),
+            || 0u32,
+            |files_sum, entry| match entry {
+                Entry::File(size) => files_sum + size,
+                Entry::Dir(name) => {
+                    dir_stack.push(name);
+                    files_sum
+                }
+            },
         ),
     )(input)?;
 
-    // Assumption: directory entries are queried in the order they are listed
-    //
-    // This assumption appears to hold for my input so it's fine
-    for (name, entry) in &mut entries {
-        if let ListEntry::Dir(placeholder) = entry {
-            let (new_input, contents) = delimited(
-                tuple((tag("$ cd "), tag(*name), newline)),
-                parse_dir,
-                // Optional cd'ing out because the last directory is never exited.
-                opt(tag("$ cd ..\n")),
-            )(input)?;
+    for i in initial_len..dir_stack.len() {
+        let (new_input, content_size) = delimited(
+            tuple((tag("$ cd "), tag(dir_stack[i]), newline)),
+            |input| parse_dir(input, dirs, dir_stack),
+            // Optional cd'ing out because the last directory is never exited.
+            opt(tag("$ cd ..\n")),
+        )(input)?;
 
-            input = new_input;
-            *placeholder = contents;
-        }
+        input = new_input;
+        size += content_size;
     }
 
-    Ok((input, entries))
+    dirs.push(size);
+
+    Ok((input, size))
 }
 
-fn parse_program(input: &[u8]) -> IResult<&[u8], Listing> {
-    preceded(tag("$ cd /\n"), parse_dir)(input)
-}
+fn parse_program(input: &[u8]) -> IResult<&[u8], (u32, Vec<u32>)> {
+    let mut dirs = Vec::new();
+    let mut dirstack = Vec::new();
+    let (input, size) = preceded(tag("$ cd /\n"), |input| {
+        parse_dir(input, &mut dirs, &mut dirstack)
+    })(input)?;
 
-fn sum_sizes(listing: Slicing<'_>, acc: &mut impl FnMut(u32)) -> u32 {
-    let mut total_size = 0;
-
-    for (_, entry) in listing {
-        match entry {
-            ListEntry::File(size) => total_size += size,
-            ListEntry::Dir(listing) => total_size += sum_sizes(listing, acc),
-        }
-    }
-
-    acc(total_size);
-
-    total_size
+    Ok((input, (size, dirs)))
 }
 
 pub fn part1(input: &[u8]) -> Result<String> {
-    let root = parse_input(input, parse_program)?;
+    let (_, sizes) = parse_input(input, parse_program)?;
 
-    let mut searched_size = 0;
-
-    sum_sizes(&root, &mut |size| {
-        if size <= 100000 {
-            searched_size += size
-        }
-    });
+    let searched_size: u32 = sizes.into_iter().filter(|&size| size <= 100000).sum();
 
     Ok(searched_size.to_string())
 }
@@ -109,11 +98,7 @@ pub fn part2(input: &[u8]) -> Result<String> {
     const TARGET: u32 = 30000000;
     const TOTAL: u32 = 70000000;
 
-    let root = parse_input(input, parse_program)?;
-
-    let mut sizes = Vec::new();
-
-    let used = sum_sizes(&root, &mut |size| sizes.push(size));
+    let (used, sizes) = parse_input(input, parse_program)?;
 
     let required = TARGET - (TOTAL - used);