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}