AOC Day 15 Rust

use rayon::prelude::*;
use indicatif::{ProgressBar, ProgressStyle};

use std::{ptr, vec};
use opencl3::{device::*, command_queue::*, context::*, program::*, kernel::*, memory::*, types::*};

#[derive(Debug, PartialEq, Eq, Hash)]
struct Sensor {
    x: isize,
    y: isize,
    closest_beacon: Beacon,
}

impl Sensor {
    fn _new(x: isize, y: isize, beacon_x: isize, beacon_y: isize) -> Self {
        Self {
            x, y,
            closest_beacon: Beacon::new(beacon_x, beacon_y),
        }
    }

    fn from_str(sensor_line: &str) -> Self {
        let mut tokens = sensor_line.split_whitespace();
        let x: isize =
            tokens.nth(2).unwrap()
            .replace(',', "")
            .chars().skip(2)
            .collect::<String>()
            .parse().unwrap();

        let y: isize =
            tokens.next().unwrap()
            .replace(':', "")
            .chars().skip(2)
            .collect::<String>()
            .parse().unwrap();

        let beacon_x: isize =
            tokens.nth(4).unwrap()
            .replace(',', "")
            .chars().skip(2)
            .collect::<String>()
            .parse().unwrap();

        let beacon_y: isize =
            tokens.next().unwrap()
            .chars().skip(2)
            .collect::<String>()
            .parse().unwrap();

        Self {
            x, y,
            closest_beacon: Beacon::new(beacon_x, beacon_y)
        }
    }

    fn beacon_cannot_exist(&self, x: isize, y: isize) -> bool {
        // If we're on top of the known beacon, then a beacon definitely can exist here
        if x == self.closest_beacon.x && y == self.closest_beacon.y {
            return false;
        }

        // A beacon cannot exist if it's within the same distance to the known closest beacon
        let distance_to_beacon = self.distance_to(self.closest_beacon.x, self.closest_beacon.y);
        let distance_to_this_point = self.distance_to(x, y);
        distance_to_this_point <= distance_to_beacon
    }

    fn distance_to(&self, x: isize, y: isize) -> usize {
        self.x.abs_diff(x) + self.y.abs_diff(y)
    }
}

#[derive(Debug, PartialEq, Eq, Hash)]
struct Beacon {
    x: isize,
    y: isize,
}

impl Beacon {
    fn new(x: isize, y: isize) -> Self {
        Self { x, y }
    }
}

pub fn solve() {
    let input = std::fs::read_to_string("inputs/day15.txt")
        .expect("Error reading input!");

    let mut sensors = Vec::<Sensor>::new();
    let mut x_bounds = (isize::MAX, isize::MIN);
    for line in input.lines() {
        let sensor = Sensor::from_str(line);
        let distance = sensor.distance_to(sensor.closest_beacon.x, sensor.closest_beacon.y) as isize;
        x_bounds.0 = *([x_bounds.0, sensor.x - distance, sensor.x + distance, sensor.closest_beacon.x].iter().min().unwrap());
        x_bounds.1 = *([x_bounds.1, sensor.x - distance, sensor.x + distance, sensor.closest_beacon.x].iter().max().unwrap());
        sensors.push(sensor);
    }

    // Part 1
    let cannot_exist =
        (x_bounds.0..=x_bounds.1).into_par_iter()
        .filter(|x| sensors.iter().any(|sensor| sensor.beacon_cannot_exist(*x, 2_000_000)))
        .count();
    println!("Part 1: {cannot_exist}");

    // Part 2
    println!("Attempting Part 2...");

    // Setup OpenCL Device
    let device_id = *get_all_devices(CL_DEVICE_TYPE_GPU).unwrap().first().unwrap();
    let device = Device::new(device_id);
    let context = Context::from_device(&device).unwrap();
    let queue = CommandQueue::create_default_with_properties(&context, CL_QUEUE_PROFILING_ENABLE, 0).unwrap();

    // Read and Build the Kernel
    let kernel_source = std::fs::read_to_string("src/day15.cl")
        .expect("Error reading OpenCL source!");
    let program = Program::create_and_build_from_source(&context, &kernel_source, "").unwrap();
    let kernel = Kernel::create(&program, "can_contain_new_beacon").unwrap();

    // Setup arguments
    let num_of_sensors: cl_int = sensors.len() as i32;
    let host_sensor_x: Vec<cl_long> = sensors.iter().map(|sensor| sensor.x as i64).collect();
    let host_sensor_y: Vec<cl_long> = sensors.iter().map(|sensor| sensor.y as i64).collect();
    let host_beacon_distance: Vec<cl_long> =
        sensors.iter()
        .map(|sensor| sensor.distance_to(sensor.closest_beacon.x, sensor.closest_beacon.y) as i64)
        .collect();
    let mut host_beacon_location: [cl_long; 2] = [0, 0];

    // Create and write to buffers
    let mut dev_sensor_x = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_READ_ONLY, host_sensor_x.len(), ptr::null_mut()).unwrap() };
    let mut dev_sensor_y = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_READ_ONLY, host_sensor_y.len(), ptr::null_mut()).unwrap() };
    let mut dev_beacon_distance = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_READ_ONLY, host_beacon_distance.len(), ptr::null_mut()).unwrap() };
    let mut dev_beacon_location = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_WRITE_ONLY, 2, ptr::null_mut()).unwrap() };
    unsafe {
        queue.enqueue_write_buffer(&mut dev_sensor_x, CL_BLOCKING, 0, &host_sensor_x, &[]).unwrap();
        queue.enqueue_write_buffer(&mut dev_sensor_y, CL_BLOCKING, 0, &host_sensor_y, &[]).unwrap();
        queue.enqueue_write_buffer(&mut dev_beacon_distance, CL_BLOCKING, 0, &host_beacon_distance, &[]).unwrap();
        queue.enqueue_write_buffer(&mut dev_beacon_location, CL_BLOCKING, 0, &host_beacon_location, &[]).unwrap();
    };

    // Step through search space in chunks using OpenCL
    let step_size: usize = 10_000;
    let progress = ProgressBar::new(4_000_001);
    progress.set_style(ProgressStyle::with_template("[{elapsed_precise}] {bar:80} {human_pos:>9}/{human_len:9} ({eta_precise})").unwrap());
    for base_x in (0..=4_000_000 as cl_long).step_by(step_size) {
        progress.inc(step_size as u64);

        // Execute Kernel
        let candidate_beacon_location = unsafe {
            let kernel_event =
                ExecuteKernel::new(&kernel)
                .set_arg(&base_x)
                .set_arg(&num_of_sensors)
                .set_arg(&dev_sensor_x)
                .set_arg(&dev_sensor_y)
                .set_arg(&dev_beacon_distance)
                .set_arg(&dev_beacon_location)
                .set_global_work_sizes(&[step_size, 4_000_001])
                .enqueue_nd_range(&queue)
                .unwrap();

            let events: Vec<cl_event> = vec![kernel_event.get()];
            let read_event = queue.enqueue_read_buffer(&dev_beacon_location, CL_BLOCKING, 0, &mut host_beacon_location, &events).unwrap();
            read_event.wait().unwrap();

            host_beacon_location
        };

        if candidate_beacon_location[0] != 0 || candidate_beacon_location[1] != 0 {
            progress.finish();
            println!("Part 2: {} {candidate_beacon_location:?}", candidate_beacon_location[0] * 4_000_000 + candidate_beacon_location[1]);
            break;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{ptr, vec, iter::repeat};
    use super::*;

    #[test]
    fn part1_data_test() {
        let input = std::fs::read_to_string("inputs/day15_test.txt")
            .expect("Error reading input!");

        let mut sensors = Vec::<Sensor>::new();
        let mut x_bounds = (isize::MAX, isize::MIN);
        for line in input.lines() {
            let sensor = Sensor::from_str(line);
            let distance = sensor.distance_to(sensor.closest_beacon.x, sensor.closest_beacon.y) as isize;
            x_bounds.0 = *([x_bounds.0, sensor.x - distance, sensor.x + distance, sensor.closest_beacon.x].iter().min().unwrap());
            x_bounds.1 = *([x_bounds.1, sensor.x - distance, sensor.x + distance, sensor.closest_beacon.x].iter().max().unwrap());
            sensors.push(sensor);
        }

        let cannot_exist =
            (x_bounds.0..=x_bounds.1).into_par_iter()
            .filter(|x| sensors.iter().any(|sensor| sensor.beacon_cannot_exist(*x, 10)))
            .count();
        assert_eq!(cannot_exist, 26);
    }

    #[test]
    fn part2_data_test() {
        let input = std::fs::read_to_string("inputs/day15_test.txt")
            .expect("Error reading input!");
        let mut sensors = Vec::<Sensor>::new();
        for line in input.lines() {
            let sensor = Sensor::from_str(line);
            sensors.push(sensor);
        }

        let search_space: Vec<(isize, isize)> =
            (0..=20).into_iter()
            .flat_map(|x| repeat(x).zip(0..=20))
            .collect();

        let solutions: Vec<(isize, isize)> =
            search_space.into_par_iter()
            .filter(|(x, y)| sensors.iter().all(|sensor| !sensor.beacon_cannot_exist(*x, *y)))
            .filter(|(x, y)| sensors.iter().all(|sensor| *x != sensor.closest_beacon.x || *y != sensor.closest_beacon.y))
            .collect();

        let tuning_freq: isize =  4_000_000 * solutions[0].0 + solutions[0].1;
        assert_eq!(tuning_freq, 56_000_011);
    }

    #[test]
    fn part2_opencl_test() {
        let input = std::fs::read_to_string("inputs/day15_test.txt")
            .expect("Error reading input!");

        // Parse Sensor Data
        let mut sensors = Vec::<Sensor>::new();
        for line in input.lines() {
            let sensor = Sensor::from_str(line);
            sensors.push(sensor);
        }

        // Setup OpenCL Context
        let device_id = *get_all_devices(CL_DEVICE_TYPE_GPU).unwrap().first().unwrap();
        let device = Device::new(device_id);
        let context = Context::from_device(&device).unwrap();
        let queue = CommandQueue::create_default_with_properties(&context, CL_QUEUE_PROFILING_ENABLE, 0).unwrap();

        // Read and Build Kernel
        let kernel_source = std::fs::read_to_string("src/day15.cl")
            .expect("Error reading OpenCL source!");
        let program = Program::create_and_build_from_source(&context, &kernel_source, "").unwrap();
        let kernel = Kernel::create(&program, "can_contain_new_beacon").unwrap();

        // Setup arguments
        let num_of_sensors: cl_int = sensors.len() as i32;
        let host_sensor_x: Vec<cl_long> = sensors.iter().map(|sensor| sensor.x as i64).collect();
        let host_sensor_y: Vec<cl_long> = sensors.iter().map(|sensor| sensor.y as i64).collect();
        let host_beacon_distance: Vec<cl_long> =
            sensors.iter()
            .map(|sensor| sensor.distance_to(sensor.closest_beacon.x, sensor.closest_beacon.y) as i64)
            .collect();
        let mut host_beacon_location: [cl_long; 2] = [0, 0];

        // Create and write to buffers
        let mut dev_sensor_x = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_READ_ONLY, sensors.len(), ptr::null_mut()).unwrap() };
        let mut dev_sensor_y = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_READ_ONLY, sensors.len(), ptr::null_mut()).unwrap() };
        let mut dev_beacon_distance = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_READ_ONLY, sensors.len(), ptr::null_mut()).unwrap() };
        let mut dev_beacon_location = unsafe { Buffer::<cl_long>::create(&context, CL_MEM_WRITE_ONLY, 2, ptr::null_mut()).unwrap() };
        unsafe {
            queue.enqueue_write_buffer(&mut dev_sensor_x, CL_BLOCKING, 0, &host_sensor_x, &[]).unwrap();
            queue.enqueue_write_buffer(&mut dev_sensor_y, CL_BLOCKING, 0, &host_sensor_y, &[]).unwrap();
            queue.enqueue_write_buffer(&mut dev_beacon_distance, CL_BLOCKING, 0, &host_beacon_distance, &[]).unwrap();
            queue.enqueue_write_buffer(&mut dev_beacon_location, CL_BLOCKING, 0, &host_beacon_location, &[]).unwrap();
        };

        // Step through search space in chunks using OpenCL
        let step_size: usize = 2;
        for base_x in (0..=20 as cl_long).step_by(step_size) {
            let candidate_beacon_location = unsafe {
                let kernel_event =
                    ExecuteKernel::new(&kernel)
                    .set_arg(&base_x)
                    .set_arg(&num_of_sensors)
                    .set_arg(&dev_sensor_x)
                    .set_arg(&dev_sensor_y)
                    .set_arg(&dev_beacon_distance)
                    .set_arg(&dev_beacon_location)
                    .set_global_work_sizes(&[step_size, 21])
                    .enqueue_nd_range(&queue)
                    .unwrap();

                let events: Vec<cl_event> = vec![kernel_event.get()];
                let read_event = queue.enqueue_read_buffer(&dev_beacon_location, CL_BLOCKING, 0, &mut host_beacon_location, &events).unwrap();
                read_event.wait().unwrap();

                host_beacon_location
            };

            if candidate_beacon_location[0] != 0 || candidate_beacon_location[1] != 0 {
                assert_eq!(4_000_000 * candidate_beacon_location[0] + candidate_beacon_location[1], 56_000_011);
                break;
            }
        }
    }
}