advent-of-code/common/rust/src/vec2.rs

use std::{
    fmt,
    ops::{Add, AddAssign, Mul, Sub, SubAssign},
    str::FromStr,
};

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Vec2 {
    pub x: i32, // increases toward right
    pub y: i32, // increases toward top
}

impl Vec2 {
    #[must_use]
    pub const fn new(x: i32, y: i32) -> Self {
        Self { x, y }
    }

    #[must_use]
    pub fn map(self, mut f: impl FnMut(i32) -> i32) -> Self {
        Self {
            x: f(self.x),
            y: f(self.y),
        }
    }

    #[must_use]
    pub fn map2(self, other: Vec2, mut f: impl FnMut(i32, i32) -> i32) -> Self {
        Self {
            x: f(self.x, other.x),
            y: f(self.y, other.y),
        }
    }

    #[must_use]
    pub fn len(self) -> f64 {
        (f64::from(self.x).powi(2) + f64::from(self.y).powi(2)).sqrt()
    }

    /// Calculates the manhattan distance between `self` and `other`.
    #[must_use]
    pub fn manhattan_dist(self, other: Vec2) -> u32 {
        let x = self.x.abs_diff(other.x);
        let y = self.y.abs_diff(other.y);
        x + y
    }

    #[must_use]
    pub fn on_x_with_manhattan_dist(self, x: i32, dist: u32) -> Option<(Vec2, Vec2)> {
        let dx = x.abs_diff(self.x);

        let dy = dist.checked_sub(dx)?;
        let dy = i32::try_from(dy).ok()?;

        let p1 = Vec2::new(x, self.y - dy);
        let p2 = Vec2::new(x, self.y + dy);

        debug_assert!(p1 <= p2);
        debug_assert_eq!(self.manhattan_dist(p1), dist);
        debug_assert_eq!(self.manhattan_dist(p2), dist);

        Some((p1, p2))
    }

    #[must_use]
    pub fn on_y_with_manhattan_dist(self, y: i32, dist: u32) -> Option<(Vec2, Vec2)> {
        let dy = y.abs_diff(self.y);

        let dx = dist.checked_sub(dy)?;
        let dx = i32::try_from(dx).ok()?;

        let p1 = Vec2::new(self.x - dx, y);
        let p2 = Vec2::new(self.x + dx, y);

        debug_assert!(p1 <= p2);
        debug_assert_eq!(self.manhattan_dist(p1), dist);
        debug_assert_eq!(self.manhattan_dist(p2), dist);

        Some((p1, p2))
    }
}

impl fmt::Display for Vec2 {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{:?}", (self.x, self.y))
    }
}

impl Add for Vec2 {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
        }
    }
}

impl AddAssign for Vec2 {
    fn add_assign(&mut self, rhs: Self) {
        self.x += rhs.x;
        self.y += rhs.y;
    }
}

impl Sub for Vec2 {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self {
            x: self.x - rhs.x,
            y: self.y - rhs.y,
        }
    }
}

impl SubAssign for Vec2 {
    fn sub_assign(&mut self, rhs: Self) {
        self.x -= rhs.x;
        self.y -= rhs.y;
    }
}

impl Mul<i32> for Vec2 {
    type Output = Self;

    fn mul(self, n: i32) -> Self::Output {
        Vec2 {
            x: self.x * n,
            y: self.y * n,
        }
    }
}

impl Mul<Vec2> for i32 {
    type Output = Vec2;

    fn mul(self, Vec2 { x, y }: Vec2) -> Self::Output {
        Vec2 {
            x: x * self,
            y: y * self,
        }
    }
}

impl From<(i32, i32)> for Vec2 {
    fn from((x, y): (i32, i32)) -> Self {
        Self { x, y }
    }
}

impl From<Vec2> for (i32, i32) {
    fn from(v: Vec2) -> Self {
        (v.x, v.y)
    }
}

impl From<Direction> for Vec2 {
    /// Creates a unit vector in the given direction.
    fn from(dir: Direction) -> Self {
        match dir {
            Direction::Up => (0, 1),
            Direction::Down => (0, -1),
            Direction::Left => (-1, 0),
            Direction::Right => (1, 0),
        }
        .into()
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Direction {
    Up,
    Down,
    Left,
    Right,
}

impl FromStr for Direction {
    type Err = ();

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "U" => Ok(Direction::Up),
            "D" => Ok(Direction::Down),
            "L" => Ok(Direction::Left),
            "R" => Ok(Direction::Right),
            _ => Err(()),
        }
    }
}

/// A line from one point to another.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Line {
    pub start: Vec2,
    pub end: Vec2,
}

impl Line {
    /// Returns an iterator over the line's points represented as [`Vec2`], including the start and
    /// end points.
    ///
    /// Only works for orthogonal lines.
    #[must_use]
    pub fn points(self) -> Option<LinePoints> {
        LinePoints::try_from(self).ok()
    }
}

/// An iterator over a line's points. See [`Line::points()`] for details.
pub struct LinePoints {
    line: Option<Line>,
}

impl TryFrom<Line> for LinePoints {
    type Error = ();

    fn try_from(line: Line) -> Result<Self, Self::Error> {
        if line.start.x == line.end.x || line.start.y == line.end.y {
            Ok(Self { line: Some(line) })
        } else {
            Err(())
        }
    }
}

impl Iterator for LinePoints {
    type Item = Vec2;

    fn next(&mut self) -> Option<Self::Item> {
        let Some(line) = self.line.as_mut() else {
            return None;
        };

        let delta = (line.end - line.start).map(i32::signum);

        let p = line.start;
        line.start += delta;

        if delta == Vec2::new(0, 0) {
            self.line = None;
        }

        Some(p)
    }
}