1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use std::collections::{BTreeMap, BTreeSet};
use std::ops::DerefMut;
use std::sync::Mutex;
use utils::md5;
use utils::prelude::*;

/// Finding MD5 hashes, part three.
///
/// Similar to [2015 Day 4](../year2015/struct.Day04.html) and [2016 Day 5](crate::Day05), but with
/// even more complex logic to assemble the answer from the matching hashes and key stretching for
/// part 2.
///
/// See [`md5::find_hash_with_appended_count()`].
#[derive(Clone, Debug)]
pub struct Day14<'a> {
    prefix: &'a str,
}

impl<'a> Day14<'a> {
    pub fn new(input: &'a str, _: InputType) -> Result<Self, InputError> {
        Ok(Self { prefix: input })
    }

    #[must_use]
    pub fn part1(&self) -> u32 {
        // TODO using multiple threads is not worth it for ~10,000 hashes
        self.find_64th_key(0)
    }

    #[must_use]
    pub fn part2(&self) -> u32 {
        self.find_64th_key(2016)
    }

    fn find_64th_key(&self, additional: u32) -> u32 {
        let mutex = Mutex::new((BTreeSet::new(), BTreeMap::new(), BTreeMap::new()));

        md5::find_hash_with_appended_count(self.prefix, additional, |i, [a, b, c, d]: [u32; 4]| {
            // Spread each nibble into bytes like utils::md5::u32_to_hex, but without the
            // unnecessary mapping to ASCII hex
            let mut nibble_bytes = [0u8; 32];
            nibble_bytes[0..8].copy_from_slice(&spread_nibbles(a).to_be_bytes());
            nibble_bytes[8..16].copy_from_slice(&spread_nibbles(b).to_be_bytes());
            nibble_bytes[16..24].copy_from_slice(&spread_nibbles(c).to_be_bytes());
            nibble_bytes[24..32].copy_from_slice(&spread_nibbles(d).to_be_bytes());

            let Some(triplet) = nibble_bytes
                .windows(3)
                .find(|&w| w[0] == w[1] && w[0] == w[2])
                .map(|w| 1u16 << w[0])
            else {
                // No triplet means there is also no quintuplet
                return false;
            };

            let quintuplet = nibble_bytes
                .windows(5)
                .filter(|&w| w[0] == w[1] && w[0] == w[2] && w[0] == w[3] && w[0] == w[4])
                .fold(0, |acc, w| acc | 1u16 << w[0]);

            let mut guard = mutex.lock().unwrap();
            let (ref mut keys, ref mut triplets, ref mut quintuplets) = guard.deref_mut();

            triplets.insert(i, triplet);

            // Check if a matching quintuplet has already been found in the following thousand
            if quintuplets
                .range(i + 1..i + 1001)
                .any(|(_, &m)| triplet & m != 0)
            {
                keys.insert(i);
            }

            if quintuplet != 0 {
                quintuplets.insert(i, quintuplet);

                // Check if any matching triplets have already been found in the previous thousand
                triplets
                    .range(i.saturating_sub(1000)..i)
                    .filter(|(_, &m)| quintuplet & m != 0)
                    .for_each(|(&k, _)| {
                        keys.insert(k);
                    });
            }

            keys.len() >= 64
        });

        let (keys, ..) = mutex.into_inner().unwrap();
        *keys.iter().nth(63).unwrap()
    }
}

fn spread_nibbles(n: u32) -> u64 {
    let mut n = u64::from(n);
    // n = 0x0000_0000_1234_5678

    n = ((n & 0x0000_0000_FFFF_0000) << 16) | (n & 0x0000_0000_0000_FFFF);
    // n = 0x0000_1234_0000_5678

    n = ((n & 0x0000_FF00_0000_FF00) << 8) | (n & 0x0000_00FF_0000_00FF);
    // n = 0x0012_0034_0056_0078

    n = ((n & 0x00F0_00F0_00F0_00F0) << 4) | (n & 0x000F_000F_000F_000F);
    // n = 0x0102_0304_0506_0708

    n
}

examples!(Day14<'_> -> (u32, u32) [
    {input: "abc", part1: 22728, part2: 22551},
]);