diff --git a/core/src/coordinates.rs b/core/src/coordinates.rs index b5862fb..d8cb202 100644 --- a/core/src/coordinates.rs +++ b/core/src/coordinates.rs @@ -308,9 +308,12 @@ impl Square { } #[must_use] - pub fn neighbor(self, direction: Direction) -> Option { + pub fn neighbor(self, direction: Direction, n: Option) -> Option { let index: u8 = self as u8; - let dir: i8 = direction.to_offset(); + + let n = n.unwrap_or(1); + let dir: i8 = direction.to_offset() * n; + match direction { Direction::North | Direction::NorthEast => { Square::try_from(index.wrapping_add_signed(dir)).ok() @@ -558,37 +561,37 @@ mod tests { fn valid_neighbors() { let sq = Square::E4; - assert_eq!(sq.neighbor(Direction::North), Some(Square::E5)); - assert_eq!(sq.neighbor(Direction::NorthEast), Some(Square::F5)); - assert_eq!(sq.neighbor(Direction::East), Some(Square::F4)); - assert_eq!(sq.neighbor(Direction::SouthEast), Some(Square::F3)); - assert_eq!(sq.neighbor(Direction::South), Some(Square::E3)); - assert_eq!(sq.neighbor(Direction::SouthWest), Some(Square::D3)); - assert_eq!(sq.neighbor(Direction::West), Some(Square::D4)); - assert_eq!(sq.neighbor(Direction::NorthWest), Some(Square::D5)); + assert_eq!(sq.neighbor(Direction::North, None), Some(Square::E5)); + assert_eq!(sq.neighbor(Direction::NorthEast, None), Some(Square::F5)); + assert_eq!(sq.neighbor(Direction::East, None), Some(Square::F4)); + assert_eq!(sq.neighbor(Direction::SouthEast, None), Some(Square::F3)); + assert_eq!(sq.neighbor(Direction::South, None), Some(Square::E3)); + assert_eq!(sq.neighbor(Direction::SouthWest, None), Some(Square::D3)); + assert_eq!(sq.neighbor(Direction::West, None), Some(Square::D4)); + assert_eq!(sq.neighbor(Direction::NorthWest, None), Some(Square::D5)); } #[test] fn invalid_neighbors() { let sq = Square::A1; - assert!(sq.neighbor(Direction::West).is_none()); - assert!(sq.neighbor(Direction::SouthWest).is_none()); - assert!(sq.neighbor(Direction::South).is_none()); + assert!(sq.neighbor(Direction::West, None).is_none()); + assert!(sq.neighbor(Direction::SouthWest, None).is_none()); + assert!(sq.neighbor(Direction::South, None).is_none()); let sq = Square::H1; - assert!(sq.neighbor(Direction::East).is_none()); - assert!(sq.neighbor(Direction::SouthEast).is_none()); - assert!(sq.neighbor(Direction::South).is_none()); + assert!(sq.neighbor(Direction::East, None).is_none()); + assert!(sq.neighbor(Direction::SouthEast, None).is_none()); + assert!(sq.neighbor(Direction::South, None).is_none()); let sq = Square::A8; - assert!(sq.neighbor(Direction::North).is_none()); - assert!(sq.neighbor(Direction::NorthWest).is_none()); - assert!(sq.neighbor(Direction::West).is_none()); + assert!(sq.neighbor(Direction::North, None).is_none()); + assert!(sq.neighbor(Direction::NorthWest, None).is_none()); + assert!(sq.neighbor(Direction::West, None).is_none()); let sq = Square::H8; - assert!(sq.neighbor(Direction::North).is_none()); - assert!(sq.neighbor(Direction::NorthEast).is_none()); - assert!(sq.neighbor(Direction::East).is_none()); + assert!(sq.neighbor(Direction::North, None).is_none()); + assert!(sq.neighbor(Direction::NorthEast, None).is_none()); + assert!(sq.neighbor(Direction::East, None).is_none()); } #[test]