aoc/year2016/

day14.rs

//! # One-Time Pad
//!
//! Brute force slog through all possible keys, parallelized as much as possible. An optimization
//! for part two is a quick method to convert `u32` to 8 ASCII digits.
use crate::util::md5::*;
use crate::util::thread::*;
use implementation::*;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::Mutex;

/// Atomics can be safely shared between threads.
struct Shared<'a> {
    input: &'a str,
    part_two: bool,
    iter: AtomicIter,
    mutex: Mutex<Exclusive>,
}

/// Regular data structures need to be protected by a mutex.
struct Exclusive {
    threes: BTreeMap<i32, u32>,
    fives: BTreeMap<i32, u32>,
    found: BTreeSet<i32>,
}

pub fn parse(input: &str) -> &str {
    input.trim()
}

/// Hash each key once.
pub fn part1(input: &str) -> i32 {
    generate_pad(input, false)
}

/// Hash each key an additional 2016 times.
pub fn part2(input: &str) -> i32 {
    generate_pad(input, true)
}

/// Find the first 64 keys that satisfy the rules.
fn generate_pad(input: &str, part_two: bool) -> i32 {
    let step = if cfg!(feature = "simd") { 32 } else { 1 };

    let exclusive =
        Exclusive { threes: BTreeMap::new(), fives: BTreeMap::new(), found: BTreeSet::new() };
    let shared =
        Shared { input, part_two, iter: AtomicIter::new(0, step), mutex: Mutex::new(exclusive) };

    // Use as many cores as possible to parallelize the search.
    spawn(|| worker(&shared));

    let exclusive = shared.mutex.into_inner().unwrap();
    *exclusive.found.iter().nth(63).unwrap()
}

/// Write the salt and integer index as ASCII characters.
fn format_string(prefix: &str, mut n: i32) -> ([u8; 64], usize) {
    let prefix_len = prefix.len();
    let digits = n.max(1).ilog10() as usize + 1;
    let size = prefix_len + digits;

    let mut buffer = [0; 64];
    buffer[..prefix_len].copy_from_slice(prefix.as_bytes());

    for i in (prefix_len..size).rev() {
        buffer[i] = b'0' + (n % 10) as u8;
        n /= 10;
    }

    (buffer, size)
}

/// Quickly convert a `u32` to an array of 8 ASCII values.
#[inline]
fn to_ascii(n: u32) -> [u8; 8] {
    // Spread each nibble into its own byte, for example `1234abcd` becomes `010203040a0b0c0d`.
    let mut n = n as u64;
    n = ((n << 16) & 0x0000ffff00000000) | (n & 0x000000000000ffff);
    n = ((n << 8) & 0x00ff000000ff0000) | (n & 0x000000ff000000ff);
    n = ((n << 4) & 0x0f000f000f000f00) | (n & 0x000f000f000f000f);

    // If a digit is 0 to 9 then we need to add `0x30` to convert to an ASCII digit.
    // For digits from 10 to 15 we need to further add `0x27` to convert to lowercase ASCII.
    // Steps:
    // * Add 6 to each digit
    // * If digit is 10 or higher then the highest bit in each nibble will be set
    // * Shift this bit to create a mask
    // * Multiply mask by 0x27 to get ASCII conversion offset
    // For example, mask of `010203040a0b0c0d` is `0000000001010101`.

    let mask = ((n + 0x0606060606060606) >> 4) & 0x0101010101010101;
    n = n + 0x3030303030303030 + 0x27 * mask;
    n.to_be_bytes()
}

#[cfg(not(feature = "simd"))]
mod implementation {
    use super::*;

    pub(super) fn worker(shared: &Shared<'_>) {
        while let Some(n) = shared.iter.next() {
            // Get the next key to check.
            let n = n as i32;

            // Calculate the hash.
            let (mut buffer, size) = format_string(shared.input, n);
            let mut result = hash(&mut buffer, size);

            if shared.part_two {
                for _ in 0..2016 {
                    buffer[0..8].copy_from_slice(&to_ascii(result[0]));
                    buffer[8..16].copy_from_slice(&to_ascii(result[1]));
                    buffer[16..24].copy_from_slice(&to_ascii(result[2]));
                    buffer[24..32].copy_from_slice(&to_ascii(result[3]));
                    result = hash(&mut buffer, 32);
                }
            }

            check(shared, n, result);
        }
    }

    /// Check for sequences of 3 or 5 consecutive matching digits.
    fn check(shared: &Shared<'_>, n: i32, hash: [u32; 4]) {
        let [a, b, c, d] = hash;

        let mut prev = u32::MAX;
        let mut same = 1;
        let mut three = 0;
        let mut five = 0;

        for mut word in [d, c, b, a] {
            for _ in 0..8 {
                let next = word & 0xf;

                if next == prev {
                    same += 1;
                } else {
                    same = 1;
                }

                if same == 3 {
                    three = 1 << next;
                }
                if same == 5 {
                    five |= 1 << next;
                }

                word >>= 4;
                prev = next;
            }
        }

        if three != 0 || five != 0 {
            let mut exclusive = shared.mutex.lock().unwrap();
            let mut candidates = Vec::new();

            // Compare against all 5 digit sequences.
            if three != 0 {
                exclusive.threes.insert(n, three);

                for (_, mask) in exclusive.fives.range(n + 1..n + 1001) {
                    if three & mask != 0 {
                        candidates.push(n);
                    }
                }
            }

            // Compare against all 3 digit sequences.
            if five != 0 {
                exclusive.fives.insert(n, five);

                for (&index, &mask) in exclusive.threes.range(n - 1000..n) {
                    if five & mask != 0 {
                        candidates.push(index);
                    }
                }
            }

            // Add any matching keys found, finishing once we have at least 64 keys.
            exclusive.found.extend(candidates);

            if exclusive.found.len() >= 64 {
                shared.iter.stop();
            }
        }
    }
}

#[cfg(feature = "simd")]
mod implementation {
    use super::*;
    use crate::util::bitset::*;
    use crate::util::md5::simd::hash_fixed;
    use std::simd::cmp::SimdPartialEq as _;
    use std::simd::*;

    /// Use SIMD to compute hashes in parallel in blocks of 32.
    #[expect(clippy::needless_range_loop)]
    pub(super) fn worker(shared: &Shared<'_>) {
        let mut result = [Simd::splat(0); 4];
        let mut buffers = [[0; 64]; 32];

        while let Some(start) = shared.iter.next() {
            // Get the next key to check.
            let start = start as i32;

            // Calculate the hash.
            for i in 0..32 {
                let (mut buffer, size) = format_string(shared.input, start + i as i32);
                let [a, b, c, d] = hash(&mut buffer, size);

                result[0][i] = a;
                result[1][i] = b;
                result[2][i] = c;
                result[3][i] = d;
            }

            if shared.part_two {
                for _ in 0..2016 {
                    for i in 0..32 {
                        buffers[i][0..8].copy_from_slice(&to_ascii(result[0][i]));
                        buffers[i][8..16].copy_from_slice(&to_ascii(result[1][i]));
                        buffers[i][16..24].copy_from_slice(&to_ascii(result[2][i]));
                        buffers[i][24..32].copy_from_slice(&to_ascii(result[3][i]));
                    }
                    result = hash_fixed(&mut buffers, 32);
                }
            }

            check(shared, start, &result);
        }
    }

    /// Check for sequences of 3 or 5 consecutive matching digits.
    #[inline]
    fn check(shared: &Shared<'_>, start: i32, hash: &[Simd<u32, 32>; 4]) {
        let &[a, b, c, d] = hash;

        let mut prev: Simd<u32, 32> = Simd::splat(u32::MAX);
        let mut same: Simd<u32, 32> = Simd::splat(1);
        let mut three: Simd<u32, 32> = Simd::splat(0);
        let mut five: Simd<u32, 32> = Simd::splat(0);

        for mut word in [d, c, b, a] {
            for _ in 0..8 {
                let next = word & Simd::splat(0xf);
                same = next.simd_eq(prev).select(same + Simd::splat(1), Simd::splat(1));

                three = same.simd_eq(Simd::splat(3)).select(Simd::splat(1) << next, three);
                five |= same.simd_eq(Simd::splat(5)).select(Simd::splat(1) << next, Simd::splat(0));

                word >>= 4;
                prev = next;
            }
        }

        let three_mask = three.simd_ne(Simd::splat(0)).to_bitmask();
        let five_mask = five.simd_ne(Simd::splat(0)).to_bitmask();

        if three_mask != 0 || five_mask != 0 {
            let mut exclusive = shared.mutex.lock().unwrap();
            let mut candidates = Vec::new();

            for i in three_mask.biterator() {
                let three = three[i];
                let n = start + i as i32;

                // Compare against all 5 digit sequences.
                if three != 0 {
                    exclusive.threes.insert(n, three);

                    for (_, mask) in exclusive.fives.range(n + 1..n + 1001) {
                        if three & mask != 0 {
                            candidates.push(n);
                        }
                    }
                }
            }

            for i in five_mask.biterator() {
                let five = five[i];
                let n = start + i as i32;

                // Compare against all 3 digit sequences.
                if five != 0 {
                    exclusive.fives.insert(n, five);

                    for (&index, &mask) in exclusive.threes.range(n - 1000..n) {
                        if five & mask != 0 {
                            candidates.push(index);
                        }
                    }
                }
            }

            // Add any matching keys found, finishing once we have at least 64 keys.
            exclusive.found.extend(candidates);

            if exclusive.found.len() >= 64 {
                shared.iter.stop();
            }
        }
    }
}