screeps_utils/
large_cost_matrix.rs

1use std::ops::{Index, IndexMut};
2
3use serde::{Deserialize, Serialize};
4use serde_with::serde_as;
5
6use screeps::{
7    constants::ROOM_SIZE,
8    objects::CostMatrix,
9    traits::{CostMatrixGet, CostMatrixSet},
10};
11
12use screeps::local::{linear_index_to_xy, xy_to_linear_index, LocalCostMatrix, Position, RoomXY};
13
14pub const ROOM_AREA: usize = ROOM_SIZE as usize * ROOM_SIZE as usize;
15
16/// A matrix of pathing costs for a room, stored in Rust memory. Stores
17/// u16 values, so it can hold significantly more data in any particular
18/// position compared to the default `CostMatrix`.
19#[serde_as]
20#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(transparent)]
22pub struct LargeCostMatrix {
23    #[serde_as(as = "[_; ROOM_AREA]")]
24    bits: [u16; ROOM_AREA],
25}
26
27impl Default for LargeCostMatrix {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl LargeCostMatrix {
34    #[inline]
35    pub const fn new() -> Self {
36        LargeCostMatrix {
37            bits: [0; ROOM_AREA],
38        }
39    }
40    #[inline]
41    pub const fn new_with_default(default: u16) -> Self {
42        LargeCostMatrix {
43            bits: [default; ROOM_AREA],
44        }
45    }
46
47    // # Notes
48    // This method does no bounds checking for the passed-in `RoomXY`, you may use
49    // `RoomXY::unchecked_new` to skip all bounds checking.
50    #[inline]
51    pub fn set(&mut self, xy: RoomXY, val: u16) {
52        self[xy] = val;
53    }
54
55    // # Notes
56    // This method does no bounds checking for the passed-in `RoomXY`, you may use
57    // `RoomXY::unchecked_new` to skip all bounds checking.
58    #[inline]
59    pub fn get(&self, xy: RoomXY) -> u16 {
60        self[xy]
61    }
62
63    pub const fn get_bits(&self) -> &[u16; ROOM_AREA] {
64        &self.bits
65    }
66
67    pub fn iter(&self) -> impl Iterator<Item = (RoomXY, u16)> + '_ {
68        self.bits
69            .iter()
70            .enumerate()
71            .map(|(idx, &val)| (linear_index_to_xy(idx), val))
72    }
73
74    pub fn iter_mut(&mut self) -> impl Iterator<Item = (RoomXY, &mut u16)> {
75        self.bits
76            .iter_mut()
77            .enumerate()
78            .map(|(idx, val)| (linear_index_to_xy(idx), val))
79    }
80}
81
82impl From<LargeCostMatrix> for Vec<u16> {
83    /// Returns a vector of bits length ROOM_AREA, where each position is
84    /// `idx = ((x * ROOM_SIZE) + y)`.
85    #[inline]
86    fn from(lcm: LargeCostMatrix) -> Vec<u16> {
87        lcm.bits.into()
88    }
89}
90
91impl From<&LargeCostMatrix> for Vec<u16> {
92    fn from(lcm: &LargeCostMatrix) -> Vec<u16> {
93        lcm.bits.into()
94    }
95}
96
97impl From<&CostMatrix> for LargeCostMatrix {
98    fn from(js_matrix: &CostMatrix) -> Self {
99        let mut bits: [u16; ROOM_AREA] = [0; ROOM_AREA];
100        js_matrix
101            .get_bits()
102            .to_vec()
103            .iter()
104            .enumerate()
105            .for_each(|(idx, &val)| bits[idx] = val.into());
106
107        LargeCostMatrix { bits }
108    }
109}
110
111impl Index<RoomXY> for LargeCostMatrix {
112    type Output = u16;
113
114    fn index(&self, xy: RoomXY) -> &Self::Output {
115        // SAFETY: RoomXY is always a valid coordinate.
116        unsafe { self.bits.get_unchecked(xy_to_linear_index(xy)) }
117    }
118}
119
120impl IndexMut<RoomXY> for LargeCostMatrix {
121    fn index_mut(&mut self, xy: RoomXY) -> &mut Self::Output {
122        // SAFETY: RoomXY is always a valid coordinate.
123        unsafe { self.bits.get_unchecked_mut(xy_to_linear_index(xy)) }
124    }
125}
126
127impl Index<Position> for LargeCostMatrix {
128    type Output = u16;
129
130    fn index(&self, idx: Position) -> &Self::Output {
131        &self[RoomXY::from(idx)]
132    }
133}
134
135impl IndexMut<Position> for LargeCostMatrix {
136    fn index_mut(&mut self, idx: Position) -> &mut Self::Output {
137        &mut self[RoomXY::from(idx)]
138    }
139}
140
141impl CostMatrixSet for LargeCostMatrix {
142    fn set_xy(&mut self, xy: RoomXY, cost: u8) {
143        LargeCostMatrix::set(self, xy, cost as u16);
144    }
145}
146
147impl CostMatrixGet for LargeCostMatrix {
148    fn get_xy(&mut self, xy: RoomXY) -> u8 {
149        match u8::try_from(LargeCostMatrix::get(self, xy)) {
150            Ok(var) => var,
151            Err(_) => u8::MAX,
152        }
153    }
154}
155
156impl From<LargeCostMatrix> for LocalCostMatrix {
157    fn from(lcm: LargeCostMatrix) -> LocalCostMatrix {
158        let mut ret_lcm = LocalCostMatrix::new();
159
160        lcm.bits
161            .to_vec()
162            .iter()
163            .enumerate()
164            .for_each(|(idx, &val)| {
165                ret_lcm.set(linear_index_to_xy(idx), val.try_into().unwrap_or(u8::MAX))
166            });
167
168        ret_lcm
169    }
170}
171
172impl From<&LargeCostMatrix> for LocalCostMatrix {
173    fn from(lcm: &LargeCostMatrix) -> LocalCostMatrix {
174        let mut ret_lcm = LocalCostMatrix::new();
175
176        lcm.bits
177            .to_vec()
178            .iter()
179            .enumerate()
180            .for_each(|(idx, &val)| {
181                ret_lcm.set(linear_index_to_xy(idx), val.try_into().unwrap_or(u8::MAX))
182            });
183
184        ret_lcm
185    }
186}
187
188impl From<LargeCostMatrix> for CostMatrix {
189    fn from(lcm: LargeCostMatrix) -> CostMatrix {
190        let mut bits: [u8; ROOM_AREA] = [0; ROOM_AREA];
191        lcm.bits
192            .to_vec()
193            .iter()
194            .enumerate()
195            .for_each(|(idx, &val)| bits[idx] = val.try_into().unwrap_or(u8::MAX));
196
197        CostMatrix::new_from_bits(&bits)
198    }
199}
screeps_utils/large_cost_matrix.rs

screeps_utils/
large_cost_matrix.rs
