screeps_utils/algorithms/
distance_transform.rs

1// Heavily based on https://github.com/Screeps-Tutorials/Screeps-Tutorials/blob/Master/basePlanningAlgorithms/distanceTransform.js
2
3use screeps::{
4    constants::extra::ROOM_SIZE,
5    local::{LocalCostMatrix, LocalRoomTerrain, RoomCoordinate, RoomXY},
6};
7
8use crate::room_coordinate::{range_exclusive, range_inclusive};
9
10/// Provides a Cost Matrix with values equal to the Chebyshev distance from any
11/// wall terrain. This does *not* calculate based on constructed walls, only
12/// terrain walls.
13pub fn chebyshev_distance_transform_from_terrain(
14    room_terrain: &LocalRoomTerrain,
15) -> LocalCostMatrix {
16    let mut initial_cm = LocalCostMatrix::new();
17
18    for (xy, cm_val) in initial_cm.iter_mut() {
19        *cm_val = match room_terrain.get_xy(xy) {
20            screeps::constants::Terrain::Wall => 0,
21            _ => u8::MAX,
22        };
23    }
24    chebyshev_distance_transform_from_cost_matrix(initial_cm)
25}
26
27/// Provides a Cost Matrix with values equal to the Chebyshev distance from any
28/// position in the provided initial Cost Matrix with a value set to 0.
29///
30/// This allows for calculating the distance transform from an arbitrary set of
31/// positions. Other position values in the initial Cost Matrix should be
32/// initialized to 255 (u8::MAX) to ensure the calculations work correctly.
33pub fn chebyshev_distance_transform_from_cost_matrix(mut cm: LocalCostMatrix) -> LocalCostMatrix {
34    let zero = RoomCoordinate::new(0).unwrap();
35    let one = RoomCoordinate::new(1).unwrap();
36    let forty_eight = RoomCoordinate::new(ROOM_SIZE - 2).unwrap();
37    let forty_nine = RoomCoordinate::new(ROOM_SIZE - 1).unwrap();
38    // Pass 1: Top-to-Bottom, Left-to-Right
39
40    // Phase A: first column
41    range_inclusive(one, forty_nine)
42        .map(|y| RoomXY { x: zero, y })
43        .fold(cm.get(RoomXY { x: zero, y: zero }), |top, xy| {
44            let current = cm.get(xy);
45            match top.checked_add(1) {
46                Some(top_val) if top_val < current => {
47                    cm.set(xy, top_val);
48                    top_val
49                }
50                _ => current,
51            }
52        });
53
54    // Phase B: the rest
55    range_inclusive(one, forty_nine)
56        .map(|x| (x, unsafe { RoomCoordinate::unchecked_new(x.u8() - 1) }))
57        .for_each(|(current_x, left_x)| {
58            let initial_top = {
59                let current = cm.get(RoomXY {
60                    x: current_x,
61                    y: zero,
62                });
63                match cm
64                    .get(RoomXY { x: left_x, y: zero })
65                    .min(cm.get(RoomXY { x: left_x, y: one }))
66                    .checked_add(1)
67                {
68                    Some(left_val) if left_val < current => {
69                        cm.set(
70                            RoomXY {
71                                x: current_x,
72                                y: zero,
73                            },
74                            left_val,
75                        );
76                        left_val
77                    }
78                    _ => current,
79                }
80            };
81            let final_top = range_exclusive(zero, forty_nine)
82                .map(|y| {
83                    (RoomXY { x: current_x, y }, unsafe {
84                        [
85                            RoomCoordinate::unchecked_new(y.u8() - 1),
86                            y,
87                            RoomCoordinate::unchecked_new(y.u8() + 1),
88                        ]
89                    })
90                })
91                .fold(initial_top, |top, (current_xy, lefts)| {
92                    let current_val = cm.get(current_xy);
93                    match lefts
94                        .into_iter()
95                        .map(|y| RoomXY { x: left_x, y })
96                        .map(|xy| cm.get(xy))
97                        .min()
98                        .unwrap()
99                        .min(top)
100                        .checked_add(1)
101                    {
102                        Some(neighbour_val) if neighbour_val < current_val => {
103                            cm.set(current_xy, neighbour_val);
104                            neighbour_val
105                        }
106                        _ => current_val,
107                    }
108                });
109            match final_top
110                .min(cm.get(RoomXY {
111                    x: left_x,
112                    y: forty_eight,
113                }))
114                .min(cm.get(RoomXY {
115                    x: left_x,
116                    y: forty_nine,
117                }))
118                .checked_add(1)
119            {
120                Some(val)
121                    if val
122                        < cm.get(RoomXY {
123                            x: current_x,
124                            y: forty_nine,
125                        }) =>
126                {
127                    cm.set(
128                        RoomXY {
129                            x: current_x,
130                            y: forty_nine,
131                        },
132                        val,
133                    );
134                }
135                _ => (),
136            };
137        });
138
139    // Pass 2: Bottom-to-Top, Right-to-Left
140
141    // Phase A: last column
142    range_inclusive(zero, forty_eight)
143        .map(|y| RoomXY { x: forty_nine, y })
144        .rfold(
145            cm.get(RoomXY {
146                x: forty_nine,
147                y: forty_nine,
148            }),
149            |bottom, xy| {
150                let current = cm.get(xy);
151                match bottom.checked_add(1) {
152                    Some(bottom_val) if bottom_val < current => {
153                        cm.set(xy, bottom_val);
154                        bottom_val
155                    }
156                    _ => current,
157                }
158            },
159        );
160
161    // Phase B: the rest
162    range_inclusive(zero, forty_eight)
163        .map(|x| (x, unsafe { RoomCoordinate::unchecked_new(x.u8() + 1) }))
164        .rev()
165        .for_each(|(current_x, right_x)| {
166            let initial_bottom = {
167                let current = cm.get(RoomXY {
168                    x: current_x,
169                    y: forty_nine,
170                });
171                match cm
172                    .get(RoomXY {
173                        x: right_x,
174                        y: forty_nine,
175                    })
176                    .min(cm.get(RoomXY {
177                        x: right_x,
178                        y: forty_eight,
179                    }))
180                    .checked_add(1)
181                {
182                    Some(left) if left < current => {
183                        cm.set(
184                            RoomXY {
185                                x: current_x,
186                                y: forty_nine,
187                            },
188                            left,
189                        );
190                        left
191                    }
192                    _ => current,
193                }
194            };
195            let final_bottom = range_exclusive(zero, forty_nine)
196                .map(|y| {
197                    (RoomXY { x: current_x, y }, unsafe {
198                        [
199                            RoomCoordinate::unchecked_new(y.u8() - 1),
200                            y,
201                            RoomCoordinate::unchecked_new(y.u8() + 1),
202                        ]
203                    })
204                })
205                .rfold(initial_bottom, |bottom, (current_xy, rights)| {
206                    let current_val = cm.get(current_xy);
207                    match rights
208                        .into_iter()
209                        .map(|y| RoomXY { x: right_x, y })
210                        .map(|xy| cm.get(xy))
211                        .min()
212                        .unwrap()
213                        .min(bottom)
214                        .checked_add(1)
215                    {
216                        Some(neighbour_val) if neighbour_val < current_val => {
217                            cm.set(current_xy, neighbour_val);
218                            neighbour_val
219                        }
220                        _ => current_val,
221                    }
222                });
223            match final_bottom
224                .min(cm.get(RoomXY { x: right_x, y: one }))
225                .min(cm.get(RoomXY {
226                    x: right_x,
227                    y: zero,
228                }))
229                .checked_add(1)
230            {
231                Some(val)
232                    if val
233                        < cm.get(RoomXY {
234                            x: current_x,
235                            y: zero,
236                        }) =>
237                {
238                    cm.set(
239                        RoomXY {
240                            x: current_x,
241                            y: zero,
242                        },
243                        val,
244                    );
245                }
246                _ => (),
247            };
248        });
249
250    cm
251}
252
253/// Provides a Cost Matrix with values equal to the Manhattan distance from any
254/// wall terrain. This does *not* calculate based on constructed walls, only
255/// terrain walls.
256pub fn manhattan_distance_transform_from_terrain(
257    room_terrain: &LocalRoomTerrain,
258) -> LocalCostMatrix {
259    let mut initial_cm = LocalCostMatrix::new();
260
261    for (xy, cm_val) in initial_cm.iter_mut() {
262        *cm_val = match room_terrain.get_xy(xy) {
263            screeps::constants::Terrain::Wall => 0,
264            _ => u8::MAX,
265        };
266    }
267    manhattan_distance_transform_from_cost_matrix(initial_cm)
268}
269
270/// Provides a Cost Matrix with values equal to the Manhattan distance from any
271/// position in the provided initial Cost Matrix with a value set to 0.
272///
273/// This allows for calculating the distance transform from an arbitrary set of
274/// positions. Other position values in the initial Cost Matrix should be
275/// initialized to 255 (u8::MAX) to ensure the calculations work correctly.
276pub fn manhattan_distance_transform_from_cost_matrix(mut cm: LocalCostMatrix) -> LocalCostMatrix {
277    let zero = RoomCoordinate::new(0).unwrap();
278    let one = RoomCoordinate::new(1).unwrap();
279    let forty_eight = RoomCoordinate::new(ROOM_SIZE - 2).unwrap();
280    let forty_nine = RoomCoordinate::new(ROOM_SIZE - 1).unwrap();
281    // Pass 1: Top-to-Bottom, Left-to-Right
282
283    // Phase A: first column
284    range_inclusive(one, forty_nine)
285        .map(|y| RoomXY { x: zero, y })
286        .fold(cm.get(RoomXY { x: zero, y: zero }), |top, xy| {
287            let val = cm.get(xy).min(top.saturating_add(1));
288            cm.set(xy, val);
289            val
290        });
291
292    // Phase B: the rest
293    range_inclusive(one, forty_nine)
294        .map(|x| (x, unsafe { RoomCoordinate::unchecked_new(x.u8() - 1) }))
295        .for_each(|(current_x, left_x)| {
296            let initial_top = {
297                let current = cm.get(RoomXY {
298                    x: current_x,
299                    y: zero,
300                });
301                match cm.get(RoomXY { x: left_x, y: zero }).checked_add(1) {
302                    Some(left) if left < current => {
303                        cm.set(
304                            RoomXY {
305                                x: current_x,
306                                y: zero,
307                            },
308                            left,
309                        );
310                        left
311                    }
312                    _ => current,
313                }
314            };
315            range_inclusive(one, forty_nine)
316                .map(|y| (RoomXY { x: current_x, y }, RoomXY { x: left_x, y }))
317                .fold(initial_top, |top, (current_xy, left_xy)| {
318                    let current_val = cm.get(current_xy);
319                    match cm.get(left_xy).min(top).checked_add(1) {
320                        Some(other) if other < current_val => {
321                            cm.set(current_xy, other);
322                            other
323                        }
324                        _ => current_val,
325                    }
326                });
327        });
328
329    // Pass 2: Bottom-to-Top, Right-to-Left
330
331    // Phase A: last column
332    range_inclusive(zero, forty_eight)
333        .map(|y| RoomXY { x: forty_nine, y })
334        .rfold(
335            cm.get(RoomXY {
336                x: forty_nine,
337                y: forty_nine,
338            }),
339            |bottom, xy| {
340                let val = cm.get(xy).min(bottom.saturating_add(1));
341                cm.set(xy, val);
342                val
343            },
344        );
345
346    // Phase B: the rest
347    range_inclusive(zero, forty_eight)
348        .map(|x| (x, unsafe { RoomCoordinate::unchecked_new(x.u8() + 1) }))
349        .rev()
350        .for_each(|(current_x, right_x)| {
351            let initial_bottom = {
352                let current = cm.get(RoomXY {
353                    x: current_x,
354                    y: forty_nine,
355                });
356                match cm
357                    .get(RoomXY {
358                        x: right_x,
359                        y: forty_nine,
360                    })
361                    .checked_add(1)
362                {
363                    Some(right) if right < current => {
364                        cm.set(
365                            RoomXY {
366                                x: current_x,
367                                y: forty_nine,
368                            },
369                            right,
370                        );
371                        right
372                    }
373                    _ => current,
374                }
375            };
376            range_inclusive(zero, forty_eight)
377                .map(|y| (RoomXY { x: current_x, y }, RoomXY { x: right_x, y }))
378                .rfold(initial_bottom, |bottom, (current_xy, right_xy)| {
379                    let current_val = cm.get(current_xy);
380                    match cm.get(right_xy).min(bottom).checked_add(1) {
381                        Some(other) if other < current_val => {
382                            cm.set(current_xy, other);
383                            other
384                        }
385                        _ => current_val,
386                    }
387                });
388        });
389
390    cm
391}
screeps_utils/algorithms/distance_transform.rs

screeps_utils/algorithms/
distance_transform.rs
