src/zones.rs 12.7 K · 385 lines · raw · history

1 use std::collections::HashSet;
2
3 use bevy::prelude::*;
4
5 use crate::{
6 actions::Action,
7 pathfinding::nearest_tile,
8 selection::{Cursor, update_cursor},
9 ui::PointerOverUi,
10 world::{TILE_SIZE, WorldMap},
11 worldgen::WorldGenerated,
12 };
13
14 pub struct ZonesPlugin;
15
16 impl Plugin for ZonesPlugin {
17 fn build(&self, app: &mut App) {
18 app.init_resource::<Zones>()
19 .init_resource::<ZoneTool>()
20 .add_systems(Update, clear_zones.run_if(on_message::<WorldGenerated>))
21 .add_systems(
22 PostUpdate,
23 (zone_tool, draw_zone_preview).chain().after(update_cursor),
24 );
25 }
26 }
27
28 pub const NEW_STOCKPILE_SIZE: u32 = 4;
29 const STOCKPILE_SITE_SEARCH_TILES: usize = 4000;
30 const PREVIEW_ADD_COLOR: Color = Color::srgba(0.4, 0.7, 1.0, 0.9);
31 const PREVIEW_REMOVE_COLOR: Color = Color::srgba(1.0, 0.4, 0.4, 0.9);
32
33 #[derive(Resource, Default, Debug, Clone, PartialEq)]
34 pub struct Zones {
35 stockpile: HashSet<UVec2>,
36 }
37
38 impl Zones {
39 pub fn is_stockpile(&self, tile: UVec2) -> bool {
40 self.stockpile.contains(&tile)
41 }
42
43 pub fn stockpile_tiles(&self) -> impl Iterator<Item = UVec2> + '_ {
44 self.stockpile.iter().copied()
45 }
46
47 pub fn home(&self, fallback: UVec2) -> UVec2 {
48 let centroid = self.stockpile.iter().map(|t| t.as_vec2()).sum::<Vec2>()
49 / self.stockpile.len().max(1) as f32;
50 self.stockpile
51 .iter()
52 .copied()
53 .min_by(|a, b| {
54 a.as_vec2()
55 .distance_squared(centroid)
56 .total_cmp(&b.as_vec2().distance_squared(centroid))
57 .then((a.y, a.x).cmp(&(b.y, b.x)))
58 })
59 .unwrap_or(fallback)
60 }
61
62 pub fn remove_tile(&mut self, tile: UVec2) {
63 self.stockpile.remove(&tile);
64 }
65
66 pub fn has_stockpile(&self) -> bool {
67 !self.stockpile.is_empty()
68 }
69
70 pub fn empty_stockpile_tiles(&self, occupied: &HashSet<UVec2>) -> usize {
71 self.stockpile.difference(occupied).count()
72 }
73
74 pub fn add_stockpile(&mut self, map: &WorldMap, a: IVec2, b: IVec2) {
75 self.add_tiles(map, rect_tiles(a, b).filter_map(|t| t.try_into().ok()));
76 }
77
78 pub fn add_tiles(&mut self, map: &WorldMap, tiles: impl IntoIterator<Item = UVec2>) {
79 for tile in tiles {
80 if can_zone(map, tile.as_ivec2()) {
81 self.stockpile.insert(tile);
82 }
83 }
84 }
85
86 pub fn remove_stockpile(&mut self, a: IVec2, b: IVec2) {
87 for tile in rect_tiles(a, b) {
88 if tile.x >= 0 && tile.y >= 0 {
89 self.stockpile.remove(&tile.as_uvec2());
90 }
91 }
92 }
93 }
94
95 pub fn plan_expansion(
96 map: &WorldMap,
97 zones: &Zones,
98 excluded: &HashSet<UVec2>,
99 count: usize,
100 ) -> Vec<UVec2> {
101 if zones.stockpile.is_empty() {
102 return Vec::new();
103 }
104 let centroid =
105 zones.stockpile.iter().map(|t| t.as_vec2()).sum::<Vec2>() / zones.stockpile.len() as f32;
106 let preference = |t: UVec2| {
107 (
108 map.plant(t.as_ivec2()).is_some(),
109 t.as_vec2().distance_squared(centroid),
110 t.y,
111 t.x,
112 )
113 };
114 let mut planned: Vec<UVec2> = Vec::new();
115 while planned.len() < count {
116 let taken =
117 |t: UVec2| zones.is_stockpile(t) || planned.contains(&t) || excluded.contains(&t);
118 let best = zones
119 .stockpile
120 .iter()
121 .chain(planned.iter())
122 .flat_map(|t| {
123 [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y].map(|d| t.as_ivec2() + d)
124 })
125 .filter(|&t| can_zone(map, t) && !taken(t.as_uvec2()))
126 .map(|t| t.as_uvec2())
127 .min_by(|&a, &b| {
128 let (ka, kb) = (preference(a), preference(b));
129 ka.0.cmp(&kb.0)
130 .then(ka.1.total_cmp(&kb.1))
131 .then((ka.2, ka.3).cmp(&(kb.2, kb.3)))
132 });
133 match best {
134 Some(tile) => planned.push(tile),
135 None => break,
136 }
137 }
138 planned
139 }
140
141 pub fn can_zone(map: &WorldMap, tile: IVec2) -> bool {
142 map.move_cost(tile).is_some() && !map.get(tile).is_some_and(|t| t.is_water())
143 }
144
145 fn rect_tiles(a: IVec2, b: IVec2) -> impl Iterator<Item = IVec2> {
146 let (min, max) = (a.min(b), a.max(b));
147 (min.y..=max.y).flat_map(move |y| (min.x..=max.x).map(move |x| IVec2::new(x, y)))
148 }
149
150 #[derive(Resource, Default, Debug)]
151 pub struct ZoneTool {
152 pub active: bool,
153 drag_start: Option<UVec2>,
154 }
155
156 fn clear_zones(mut zones: ResMut<Zones>) {
157 *zones = Zones::default();
158 }
159
160 pub fn new_stockpile_tiles(corner: UVec2) -> impl Iterator<Item = UVec2> {
161 (0..NEW_STOCKPILE_SIZE)
162 .flat_map(move |y| (0..NEW_STOCKPILE_SIZE).map(move |x| corner + UVec2::new(x, y)))
163 }
164
165 fn is_open_square(map: &WorldMap, corner: UVec2) -> bool {
166 new_stockpile_tiles(corner)
167 .all(|t| can_zone(map, t.as_ivec2()) && map.plant(t.as_ivec2()).is_none())
168 }
169
170 pub fn find_stockpile_site(map: &WorldMap, near: UVec2) -> Option<UVec2> {
171 nearest_tile(map, near, STOCKPILE_SITE_SEARCH_TILES, |p| {
172 is_open_square(map, p)
173 })
174 .or_else(|| {
175 nearest_tile(map, near, STOCKPILE_SITE_SEARCH_TILES, |p| {
176 can_zone(map, p.as_ivec2())
177 })
178 })
179 }
180
181 fn removing(keys: &ButtonInput<KeyCode>) -> bool {
182 keys.any_pressed([KeyCode::ShiftLeft, KeyCode::ShiftRight])
183 }
184
185 fn zone_tool(
186 mut actions: MessageReader<Action>,
187 keys: Res<ButtonInput<KeyCode>>,
188 buttons: Res<ButtonInput<MouseButton>>,
189 over_ui: Res<PointerOverUi>,
190 cursor: Res<Cursor>,
191 map: Res<WorldMap>,
192 mut tool: ResMut<ZoneTool>,
193 mut zones: ResMut<Zones>,
194 ) {
195 for action in actions.read() {
196 match action {
197 Action::ToggleStockpileTool => tool.active = !tool.active,
198 Action::Cancel => tool.active = false,
199 _ => continue,
200 }
201 tool.drag_start = None;
202 return;
203 }
204 if !tool.active {
205 return;
206 }
207 if buttons.just_pressed(MouseButton::Left) && !over_ui.0 {
208 tool.drag_start = cursor.tile;
209 }
210 if !buttons.just_released(MouseButton::Left) {
211 return;
212 }
213 if let (Some(start), Some(end)) = (tool.drag_start.take(), cursor.tile) {
214 if removing(&keys) {
215 zones.remove_stockpile(start.as_ivec2(), end.as_ivec2());
216 } else {
217 zones.add_stockpile(&map, start.as_ivec2(), end.as_ivec2());
218 }
219 }
220 }
221
222 fn draw_zone_preview(
223 keys: Res<ButtonInput<KeyCode>>,
224 cursor: Res<Cursor>,
225 tool: Res<ZoneTool>,
226 mut gizmos: Gizmos,
227 ) {
228 let (Some(start), Some(end)) = (tool.drag_start, cursor.tile) else {
229 return;
230 };
231 let (min, max) = (start.min(end), start.max(end));
232 let center = (WorldMap::tile_center(min) + WorldMap::tile_center(max)) / 2.0;
233 let size = (max - min + 1).as_vec2() * TILE_SIZE;
234 let color = if removing(&keys) {
235 PREVIEW_REMOVE_COLOR
236 } else {
237 PREVIEW_ADD_COLOR
238 };
239 gizmos.rect_2d(center, size, color);
240 }
241
242 #[cfg(test)]
243 mod tests {
244 use super::*;
245 use crate::world::Terrain;
246
247 #[test]
248 fn stockpiles_only_cover_walkable_dry_land() {
249 let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Grass);
250 map.set(UVec2::new(1, 0), Terrain::Stone);
251 map.set(UVec2::new(2, 0), Terrain::ShallowWater);
252 let mut zones = Zones::default();
253 zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(-1, 0));
254 let mut tiles: Vec<_> = zones.stockpile_tiles().collect();
255 tiles.sort_by_key(|t| t.x);
256 assert_eq!(tiles, vec![UVec2::new(0, 0), UVec2::new(3, 0)]);
257 }
258
259 #[test]
260 fn expansion_grows_compactly_around_the_stockpile() {
261 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
262 let mut zones = Zones::default();
263 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(5, 5));
264 let plan = plan_expansion(&map, &zones, &HashSet::new(), 8);
265 assert_eq!(plan.len(), 8);
266 assert!(plan.iter().all(|t| !zones.is_stockpile(*t)));
267 let beyond_one_ring = plan_expansion(&map, &zones, &HashSet::new(), 20);
268 assert_eq!(beyond_one_ring.len(), 20);
269 assert!(
270 plan.iter()
271 .all(|t| (3..=6).contains(&t.x) && (3..=6).contains(&t.y))
272 );
273 }
274
275 // Trees can be zoned, but stockpiling on open ground keeps forests intact and avoids
276 // colonists cutting trees just to make room, so expansion only takes trees as a last resort.
277 #[test]
278 fn expansion_prefers_open_ground_over_trees_and_skips_unzonable_tiles() {
279 let mut map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
280 map.set(UVec2::new(3, 4), Terrain::Stone);
281 map.set_plant(UVec2::new(5, 4), Some(crate::world::Plant::Tree));
282 let mut zones = Zones::default();
283 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(4, 4));
284 let plan = plan_expansion(&map, &zones, &HashSet::new(), 2);
285 assert_eq!(plan.len(), 2);
286 assert!(!plan.contains(&UVec2::new(3, 4)));
287 assert!(!plan.contains(&UVec2::new(5, 4)));
288 }
289
290 #[test]
291 fn no_expansion_without_a_stockpile() {
292 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
293 assert!(plan_expansion(&map, &Zones::default(), &HashSet::new(), 8).is_empty());
294 }
295
296 #[test]
297 fn removing_clears_a_rectangle() {
298 let map = WorldMap::filled(UVec2::new(4, 4), Terrain::Grass);
299 let mut zones = Zones::default();
300 zones.add_stockpile(&map, IVec2::ZERO, IVec2::splat(3));
301 zones.remove_stockpile(IVec2::new(1, 1), IVec2::new(2, 2));
302 assert_eq!(zones.stockpile_tiles().count(), 12);
303 assert!(!zones.is_stockpile(UVec2::new(1, 2)));
304 assert!(zones.is_stockpile(UVec2::new(0, 2)));
305 }
306
307 fn tool_app() -> App {
308 let mut app = App::new();
309 app.insert_resource(WorldMap::filled(UVec2::splat(8), Terrain::Grass))
310 .init_resource::<Zones>()
311 .init_resource::<ButtonInput<KeyCode>>()
312 .init_resource::<ButtonInput<MouseButton>>()
313 .insert_resource(ZoneTool {
314 active: true,
315 drag_start: None,
316 })
317 .init_resource::<Cursor>()
318 .init_resource::<PointerOverUi>()
319 .add_message::<Action>()
320 .add_systems(Update, zone_tool);
321 app
322 }
323
324 fn drag(app: &mut App, from: UVec2, to: UVec2) {
325 app.world_mut().resource_mut::<Cursor>().tile = Some(from);
326 app.world_mut()
327 .resource_mut::<ButtonInput<MouseButton>>()
328 .press(MouseButton::Left);
329 app.update();
330 app.world_mut().resource_mut::<Cursor>().tile = Some(to);
331 let mut buttons = app.world_mut().resource_mut::<ButtonInput<MouseButton>>();
332 buttons.clear();
333 buttons.release(MouseButton::Left);
334 app.update();
335 app.world_mut()
336 .resource_mut::<ButtonInput<MouseButton>>()
337 .clear();
338 }
339
340 #[test]
341 fn dragging_adds_and_shift_dragging_removes() {
342 let mut app = tool_app();
343 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
344 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 6);
345 app.world_mut()
346 .resource_mut::<ButtonInput<KeyCode>>()
347 .press(KeyCode::ShiftLeft);
348 drag(&mut app, UVec2::new(3, 2), UVec2::new(3, 1));
349 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 4);
350 }
351
352 #[test]
353 fn dragging_does_nothing_when_the_tool_is_off() {
354 let mut app = tool_app();
355 app.world_mut().resource_mut::<ZoneTool>().active = false;
356 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
357 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 0);
358 }
359
360 #[test]
361 fn new_worlds_start_without_stockpiles() {
362 let map = WorldMap::filled(UVec2::splat(8), Terrain::Grass);
363 let mut app = App::new();
364 app.init_resource::<Zones>()
365 .add_message::<WorldGenerated>()
366 .add_systems(Update, clear_zones.run_if(on_message::<WorldGenerated>));
367 app.world_mut()
368 .resource_mut::<Zones>()
369 .add_stockpile(&map, IVec2::ZERO, IVec2::ONE);
370 app.world_mut().write_message(WorldGenerated);
371 app.update();
372 assert!(!app.world().resource::<Zones>().has_stockpile());
373 }
374
375 #[test]
376 fn new_stockpiles_go_on_the_nearest_fully_open_ground() {
377 let mut map = WorldMap::filled(UVec2::splat(12), Terrain::Grass);
378 for y in 0..12 {
379 map.set_plant(UVec2::new(4, y), Some(crate::world::Plant::Tree));
380 map.set(UVec2::new(7, y), Terrain::DeepWater);
381 }
382 let site = find_stockpile_site(&map, UVec2::new(5, 5)).unwrap();
383 assert!(new_stockpile_tiles(site).all(|t| t.x < 4));
384 }
385 }