1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
//! A stack-allocated map backed by a array that can be indexed by `Direction`.

use std::ops::{Index, IndexMut};

use crate::directions::Direction;

/// A stack-allocated map backed by a array that can be indexed by `Direction`.
///
/// This type is only available with the `collections` feature.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug, Default)]
#[cfg_attr(feature = "serde-1", derive(Serialize, Deserialize))]
pub struct DirectionMap<T> {
    arr: [T; 8],
}

impl<T: Default> DirectionMap<T> {
    pub fn new() -> Self {
        Self::default()
    }
}

impl<T> DirectionMap<T> {
    /// Create a `DirectionMap` from an array.
    pub fn from_array(arr: [T; 8]) -> Self {
        DirectionMap { arr }
    }

    /// Returns the underlying array in this `DirectionMap`.
    pub fn into_array(self) -> [T; 8] {
        self.arr
    }

    /// Returns a reference to the underlying array.
    pub fn as_array(&self) -> &[T; 8] {
        &self.arr
    }

    /// Returns a mutable reference to the underlying array.
    pub fn as_mut_array(&mut self) -> &mut [T; 8] {
        &mut self.arr
    }

    /// Returns a reference to the underlying array as a slice.
    pub fn as_slice(&self) -> &[T] {
        &self.arr
    }

    /// Returns a mutable reference to the underlying array as a slice.
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        &mut self.arr
    }

    /// Returns a reference to the value corresponding to `dir`.
    pub fn get(&self, dir: Direction) -> &T {
        let idx = dir as u8 as usize;
        debug_assert!(idx < 8);

        // SAFETY: It's not possible to normally create a Direction with a bit representation >= 8.
        unsafe { self.arr.get_unchecked(idx) }
    }

    /// Returns a mutable reference to the value corresponding to `dir`.
    pub fn get_mut(&mut self, dir: Direction) -> &mut T {
        let idx = dir as u8 as usize;
        debug_assert!(idx < 8);

        // SAFETY: It's not possible to normally create a Direction with a bit representation >= 8.
        unsafe { self.arr.get_unchecked_mut(idx) }
    }

    /// Construct a `DirectionMap<T>` with values provided from a `FnMut(Direction) -> T`.
    pub fn with_initializer<F>(mut op: F) -> Self
    where
        F: FnMut(Direction) -> T,
    {
        #![allow(clippy::cast_sign_loss)]

        use std::mem::MaybeUninit;

        let mut arr = MaybeUninit::<[T; 8]>::uninit();
        let ptr = arr.as_mut_ptr() as *mut T;

        // SAFETY: integers in 0..8 are valid values for Direction
        let arr = unsafe {
            for i in 0..8 {
                let dir = std::mem::transmute::<u8, Direction>(i as u8);
                ptr.offset(i).write(op(dir));
            }

            arr.assume_init()
        };

        DirectionMap { arr }
    }
}

impl<T> Index<Direction> for DirectionMap<T> {
    type Output = T;

    fn index(&self, d: Direction) -> &T {
        &self.arr[d as u8 as usize]
    }
}

impl<T> IndexMut<Direction> for DirectionMap<T> {
    fn index_mut(&mut self, d: Direction) -> &mut T {
        &mut self.arr[d as u8 as usize]
    }
}