src/zones.rs 19 K · 590 lines · raw · history

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