src/zones.rs 19.9 K · 620 lines · raw · history

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