adventofcode23/src/util/grid.rs

use std::ops::Range;

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Pos {
    pub x: i64,
    pub y: i64,
}

const DIRECTIONS_CROSS: [Direction; 4] = [
    Direction::North,
    Direction::East,
    Direction::South,
    Direction::West,
];

const DIRECTIONS_DIAG: [Direction; 8] = [
    Direction::East,
    Direction::Southeast,
    Direction::South,
    Direction::Southwest,
    Direction::West,
    Direction::Northwest,
    Direction::North,
    Direction::Northeast,
];

const DIRECTIONS_DIAG2: [Direction; 8] = [
    Direction::North,
    Direction::Northeast,
    Direction::East,
    Direction::Southeast,
    Direction::South,
    Direction::Southwest,
    Direction::West,
    Direction::Northwest,
];

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum Direction {
    North,
    Northeast,
    East,
    Southeast,
    South,
    Southwest,
    West,
    Northwest,
}

impl std::fmt::Debug for Pos {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("({}|{})", self.x, self.y))
    }
}

impl std::ops::Add for Pos {
    type Output = Pos;

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

impl std::ops::Sub for Pos {
    type Output = Pos;

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

impl Pos {
    pub const fn usize(x: usize, y: usize) -> Self {
        Self {
            x: x as i64,
            y: y as i64,
        }
    }

    /// Create a new position offset from itself
    pub fn offset(self, dir: Direction) -> Self {
        match dir {
            Direction::North => Self {
                x: self.x,
                y: self.y - 1,
            },
            Direction::South => Self {
                x: self.x,
                y: self.y + 1,
            },
            Direction::West => Self {
                x: self.x - 1,
                y: self.y,
            },
            Direction::East => Self {
                x: self.x + 1,
                y: self.y,
            },
            Direction::Northeast => Self {
                x: self.x + 1,
                y: self.y - 1,
            },
            Direction::Southeast => Self {
                x: self.x + 1,
                y: self.y + 1,
            },
            Direction::Southwest => Self {
                x: self.x - 1,
                y: self.y + 1,
            },
            Direction::Northwest => Self {
                x: self.x - 1,
                y: self.y - 1,
            },
        }
    }

    pub fn direction(self, p2: Pos) -> Option<Direction> {
        if p2 == self.offset(Direction::North) {
            Some(Direction::North)
        } else if p2 == self.offset(Direction::East) {
            Some(Direction::East)
        } else if p2 == self.offset(Direction::South) {
            Some(Direction::South)
        } else if p2 == self.offset(Direction::West) {
            Some(Direction::West)
        } else {
            None
        }
    }

    /// Get an iterator over all neighbors (horizontal and diagonal)
    pub fn neighbors(self, cfg: &NeighborsCfg) -> Neighbors {
        Neighbors {
            pos: self,
            i: 0,
            diag: cfg.diag,
            tl: cfg.tl,
            br: cfg.br,
        }
    }

    /// Range between 2 X coordinates
    pub fn xrange(self, p2: Pos) -> Range<i64> {
        self.x.min(p2.x)..self.x.max(p2.x)
    }

    /// Range between 2 Y coordinates
    pub fn yrange(self, p2: Pos) -> Range<i64> {
        self.y.min(p2.y)..self.y.max(p2.y)
    }
}

impl Direction {
    pub fn rotate_45(self, n: i8) -> Direction {
        let dl = DIRECTIONS_DIAG2.len() as i8;
        let mut i = (self as i8 + n) % dl;
        if i < 0 {
            i += dl;
        }
        DIRECTIONS_DIAG2[i as usize]
    }

    pub fn rotate_90(self, n: i8) -> Direction {
        let dl = DIRECTIONS_DIAG2.len() as i16;
        let mut i = (self as i16 + 2 * i16::from(n)) % dl;
        if i < 0 {
            i += dl;
        }
        DIRECTIONS_DIAG2[i as usize]
    }
}

#[derive(Default, Debug)]
pub struct NeighborsCfg {
    // output diagonal neighbors, too
    pub diag: bool,
    // top left grid position
    pub tl: Option<Pos>,
    // bottom right grid position
    pub br: Option<Pos>,
}

pub struct Neighbors {
    // starting position
    pos: Pos,
    // index
    i: u8,
    // output diagonal neighbors, too
    diag: bool,
    // top left grid position
    tl: Option<Pos>,
    // bottom right grid position
    br: Option<Pos>,
}

impl Neighbors {
    fn check_lim(&self, pos: Pos) -> bool {
        if let Some(tl) = self.tl {
            if pos.x < tl.x || pos.y < tl.y {
                return false;
            }
        }
        if let Some(br) = self.br {
            if pos.x > br.x || pos.y > br.y {
                return false;
            }
        }
        true
    }
}

impl Iterator for Neighbors {
    type Item = Pos;

    fn next(&mut self) -> Option<Self::Item> {
        let dlist = if self.diag {
            DIRECTIONS_DIAG.as_slice()
        } else {
            DIRECTIONS_CROSS.as_slice()
        };
        let pos = dlist.get(self.i as usize).map(|d| self.pos.offset(*d));
        if let Some(pos) = pos {
            self.i += 1;
            // Skip to next neighbor if out of limits
            if !self.check_lim(pos) {
                return self.next();
            }
        }
        pos
    }
}

pub struct Griterator {
    tl: Pos,
    br: Pos,
    p: Pos,
}

impl Griterator {
    pub fn new(tl: Pos, br: Pos) -> Self {
        Self {
            tl,
            br,
            p: Pos {
                x: tl.x - 1,
                y: tl.y,
            },
        }
    }
}

impl Iterator for Griterator {
    type Item = Pos;

    fn next(&mut self) -> Option<Self::Item> {
        self.p.x += 1;
        if self.p.x > self.br.x {
            self.p.x = self.tl.x;
            self.p.y += 1;
            if self.p.y > self.br.y {
                return None;
            }
        }
        Some(self.p)
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use itertools::Itertools;

    use super::*;

    #[test]
    fn neighbors() {
        let p1 = Pos::usize(1, 1);
        let nx = p1
            .neighbors(&NeighborsCfg {
                diag: false,
                ..Default::default()
            })
            .collect::<HashSet<_>>();
        assert_eq!(nx.len(), 4);
    }

    #[test]
    fn neighbors_diag() {
        let p1 = Pos::usize(1, 1);
        let nx = p1
            .neighbors(&NeighborsCfg {
                diag: true,
                ..Default::default()
            })
            .collect::<HashSet<_>>();
        assert_eq!(nx.len(), 8);
    }

    #[test]
    fn griterator() {
        let positions = Griterator::new(Pos::usize(0, 0), Pos::usize(2, 2)).collect_vec();
        assert_eq!(positions[0], Pos::usize(0, 0));
        assert_eq!(positions[8], Pos::usize(2, 2));
        assert_eq!(positions.len(), 9);
    }

    #[test]
    fn dir_rotate() {
        assert_eq!(Direction::North.rotate_90(0), Direction::North);
        assert_eq!(Direction::North.rotate_90(1), Direction::East);
        assert_eq!(Direction::North.rotate_90(2), Direction::South);
        assert_eq!(Direction::North.rotate_90(-1), Direction::West);
        assert_eq!(Direction::North.rotate_90(-2), Direction::South);
        assert_eq!(Direction::North.rotate_90(-121), Direction::West);
    }
}