screeps_utils/
sparse_cost_matrix.rs

1use std::collections::HashMap;
2
3use screeps::{
4    constants::ROOM_SIZE,
5    local::{linear_index_to_xy, LocalCostMatrix, Position, RoomXY},
6    objects::CostMatrix,
7    traits::{CostMatrixGet, CostMatrixSet},
8};
9use serde::{Deserialize, Serialize};
10
11pub const ROOM_AREA: usize = ROOM_SIZE as usize * ROOM_SIZE as usize;
12
13#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(transparent)]
15pub struct SparseCostMatrix {
16    inner: HashMap<RoomXY, u8>,
17}
18
19impl Default for SparseCostMatrix {
20    fn default() -> Self {
21        Self::new()
22    }
23}
24
25impl SparseCostMatrix {
26    pub fn new() -> Self {
27        SparseCostMatrix {
28            inner: HashMap::new(),
29        }
30    }
31
32    pub fn get(&self, xy: RoomXY) -> u8 {
33        *self.inner.get(&xy).unwrap_or(&0)
34    }
35
36    pub fn set(&mut self, xy: RoomXY, val: u8) {
37        self.inner.insert(xy, val);
38    }
39
40    pub fn iter(&self) -> impl Iterator<Item = (RoomXY, u8)> + '_ {
41        self.inner.iter().map(|(&pos, &val)| (pos, val))
42    }
43
44    pub fn iter_mut(&mut self) -> impl Iterator<Item = (RoomXY, &mut u8)> {
45        self.inner.iter_mut().map(|(&pos, val)| (pos, val))
46    }
47
48    // Takes all non-zero entries in `src`, and inserts them into `self`.
49    //
50    // If an entry for that position exists already, overwrites it with the new
51    // value.
52    pub fn merge_from_dense(&mut self, src: &LocalCostMatrix) {
53        self.inner.extend(src.iter().filter_map(
54            |(xy, val)| {
55                if val > 0 {
56                    Some((xy, val))
57                } else {
58                    None
59                }
60            },
61        ))
62    }
63
64    // Takes all entries in `src` and merges them into `self`.
65    //
66    // If an entry for that position exists already, overwrites it with the new
67    // value.
68    pub fn merge_from_sparse(&mut self, src: &SparseCostMatrix) {
69        self.inner.extend(src.inner.iter());
70    }
71}
72
73impl From<HashMap<RoomXY, u8>> for SparseCostMatrix {
74    fn from(inner: HashMap<RoomXY, u8>) -> Self {
75        SparseCostMatrix { inner }
76    }
77}
78
79impl From<&HashMap<RoomXY, u8>> for SparseCostMatrix {
80    fn from(map: &HashMap<RoomXY, u8>) -> Self {
81        SparseCostMatrix { inner: map.clone() }
82    }
83}
84
85impl From<&HashMap<Position, u8>> for SparseCostMatrix {
86    fn from(map: &HashMap<Position, u8>) -> Self {
87        SparseCostMatrix {
88            inner: map.iter().map(|(&pos, &val)| (pos.into(), val)).collect(),
89        }
90    }
91}
92
93impl From<&CostMatrix> for SparseCostMatrix {
94    fn from(js_matrix: &CostMatrix) -> Self {
95        let vals: Vec<u8> = js_matrix.get_bits().to_vec();
96        assert!(
97            vals.len() == ROOM_AREA,
98            "JS CostMatrix had length {} instead of {}.",
99            vals.len(),
100            ROOM_AREA
101        );
102
103        SparseCostMatrix {
104            inner: vals
105                .into_iter()
106                .enumerate()
107                .filter_map(|(idx, val)| {
108                    // 0 is the same as unset, so filtering it out
109                    if val > 0 {
110                        Some((linear_index_to_xy(idx), val))
111                    } else {
112                        None
113                    }
114                })
115                .collect(),
116        }
117    }
118}
119
120impl From<&LocalCostMatrix> for SparseCostMatrix {
121    fn from(lcm: &LocalCostMatrix) -> Self {
122        SparseCostMatrix {
123            inner: lcm
124                .iter()
125                .filter_map(|(xy, val)| if val > 0 { Some((xy, val)) } else { None })
126                .collect(),
127        }
128    }
129}
130
131impl From<SparseCostMatrix> for LocalCostMatrix {
132    fn from(mut scm: SparseCostMatrix) -> Self {
133        let mut lcm = LocalCostMatrix::new();
134        for (pos, val) in scm.inner.drain() {
135            lcm[pos] = val;
136        }
137        lcm
138    }
139}
140
141impl From<&SparseCostMatrix> for LocalCostMatrix {
142    fn from(scm: &SparseCostMatrix) -> Self {
143        let mut lcm = LocalCostMatrix::new();
144        for (&pos, &val) in scm.inner.iter() {
145            lcm[pos] = val;
146        }
147        lcm
148    }
149}
150
151impl CostMatrixSet for SparseCostMatrix {
152    fn set_xy(&mut self, xy: RoomXY, cost: u8) {
153        SparseCostMatrix::set(self, xy, cost);
154    }
155}
156
157impl CostMatrixGet for SparseCostMatrix {
158    fn get_xy(&mut self, xy: RoomXY) -> u8 {
159        SparseCostMatrix::get(self, xy)
160    }
161}
screeps_utils/sparse_cost_matrix.rs

screeps_utils/
sparse_cost_matrix.rs
