screeps_utils/algorithms/floodfill.rs
1use std::collections::VecDeque;
2
3use crate::large_cost_matrix::LargeCostMatrix;
4use screeps::{
5 constants::Direction,
6 local::{LocalCostMatrix, LocalRoomTerrain, RoomXY},
7};
8
9/// Creates a LocalCostMatrix from LocalRoomTerrain, such that all the positions
10/// which are Walls will have the value u8::MAX, and all other positions will
11/// have a value of 0.
12pub fn get_obstacles_lcm_from_terrain(room_terrain: &LocalRoomTerrain) -> LocalCostMatrix {
13 let mut obstacles = LocalCostMatrix::new();
14 for (xy, cm_val) in obstacles.iter_mut() {
15 *cm_val = match room_terrain.get_xy(xy) {
16 screeps::constants::Terrain::Wall => u8::MAX,
17 _ => 0,
18 };
19 }
20
21 obstacles
22}
23
24/// Generates a distance floodfill from origin positions.
25///
26/// Takes a Vec of origin locations to start the floodfill from, and a Cost
27/// Matrix of obstacles, and produces a `LargeCostMatrix` with distances for all
28/// positions that can be reached from the origin points.
29///
30/// The obstacles Cost Matrix should have u8::MAX set on all positions that are
31/// obstacles, and 0 everywhere else.
32pub fn numerical_floodfill(
33 origins: &Vec<RoomXY>,
34 obstacles: &LocalCostMatrix,
35 max_distance: u16,
36) -> LargeCostMatrix {
37 let mut output_cm = LargeCostMatrix::new_with_default(u16::MAX);
38
39 let mut queue: VecDeque<RoomXY> = VecDeque::new();
40 let mut seen = LocalCostMatrix::new();
41
42 // Add all the valid neighbors of the origin positions to the queue to be
43 // visited
44 for current_position in origins {
45 // The current origin position is trivially reachable from the set of origin
46 // positions
47 output_cm.set(*current_position, 0);
48
49 // We've visited this origin position
50 seen.set(*current_position, 1);
51
52 let neighbor_distance = 1;
53
54 // Check each neighbor of the current origin position for validity and add it to
55 // the queue to be checked
56 Direction::iter()
57 .filter_map(|dir| current_position.checked_add_direction(*dir))
58 .filter(|position| obstacles.get(*position) == 0)
59 .for_each(|position| {
60 // Only process neighbors that haven't been seen yet
61 if seen.get(position) == 0 && neighbor_distance <= max_distance {
62 queue.push_back(position);
63 seen.set(position, 1);
64 output_cm.set(position, neighbor_distance);
65 }
66 });
67 }
68
69 // Process all entries in the queue
70 let mut max_queue_length = 0;
71 while !queue.is_empty() {
72 let queue_length = queue.len();
73 if queue_length > max_queue_length {
74 max_queue_length = queue_length;
75 }
76
77 // Pop the next entry off the queue
78 if let Some(current_position) = queue.pop_front() {
79 // Mark current position as visited
80 seen.set(current_position, 1);
81
82 // Get the current distance value to increment for unvisited neighbors
83 let current_distance = output_cm.get(current_position);
84 let neighbor_distance = current_distance + 1;
85
86 // Get list of valid neighbors and add them to the queue to be visited
87 Direction::iter()
88 .filter_map(|dir| current_position.checked_add_direction(*dir))
89 .filter(|position| obstacles.get(*position) == 0)
90 .for_each(|position| {
91 // Only process neighbors that haven't been seen yet
92 if seen.get(position) == 0 && neighbor_distance <= max_distance {
93 queue.push_back(position);
94 seen.set(position, 1);
95
96 // Only update the position distance if it hasn't already been set
97 if output_cm.get(position) == u16::MAX {
98 output_cm.set(position, neighbor_distance);
99 }
100 }
101 });
102 };
103 }
104
105 output_cm
106}
107
108/// Generates a reachability floodfill from origin positions.
109///
110/// Takes a Vec of origin locations to start the floodfill from, and a Cost
111/// Matrix of obstacles, and produces a Cost Matrix with 1 values for all
112/// positions that can be reached from the origin points, and 0 values
113/// everywhere else.
114///
115/// The obstacles Cost Matrix should have u8::MAX set on all positions that are
116/// obstacles, and 0 everywhere else.
117pub fn reachability_floodfill(
118 origins: &Vec<RoomXY>,
119 obstacles: &LocalCostMatrix,
120) -> LocalCostMatrix {
121 let ff_lg_cm = numerical_floodfill(origins, obstacles, u16::MAX);
122
123 let mut ret_cm = LocalCostMatrix::new();
124
125 for (xy, cm_val) in ret_cm.iter_mut() {
126 *cm_val = match ff_lg_cm.get(xy) {
127 u16::MAX => 0,
128 _ => 1,
129 };
130 }
131
132 ret_cm
133}