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