src/zones.rs 19.8 K · 617 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, 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() && !map.get(tile).is_some_and(|t| t.is_water())
323 }
324
325 pub fn can_farm(map: &WorldMap, tile: IVec2) -> bool {
326 can_zone(map, tile) && map.get(tile).is_some_and(|t| fertility(t) > 0.0)
327 }
328
329 fn rect_tiles(a: IVec2, b: IVec2) -> impl Iterator<Item = IVec2> {
330 let (min, max) = (a.min(b), a.max(b));
331 (min.y..=max.y).flat_map(move |y| (min.x..=max.x).map(move |x| IVec2::new(x, y)))
332 }
333
334 #[derive(Resource, Default, Debug)]
335 pub struct ZoneTool {
336 pub mode: Option<ZoneKind>,
337 drag_start: Option<UVec2>,
338 }
339
340 impl ZoneTool {
341 pub fn active(&self) -> bool {
342 self.mode.is_some()
343 }
344
345 fn toggle(&mut self, kind: ZoneKind) {
346 self.mode = (self.mode != Some(kind)).then_some(kind);
347 }
348 }
349
350 fn clear_zones(mut zones: ResMut<Zones>) {
351 *zones = Zones::default();
352 }
353
354 pub fn new_stockpile_tiles(corner: UVec2) -> impl Iterator<Item = UVec2> {
355 (0..NEW_STOCKPILE_SIZE)
356 .flat_map(move |y| (0..NEW_STOCKPILE_SIZE).map(move |x| corner + UVec2::new(x, y)))
357 }
358
359 fn is_open_square(map: &WorldMap, corner: UVec2) -> bool {
360 new_stockpile_tiles(corner)
361 .all(|t| can_zone(map, t.as_ivec2()) && map.plant(t.as_ivec2()).is_none())
362 }
363
364 pub fn find_stockpile_site(map: &WorldMap, near: UVec2) -> Option<UVec2> {
365 nearest_tile(map, near, STOCKPILE_SITE_SEARCH_TILES, |p| {
366 is_open_square(map, p)
367 })
368 .or_else(|| {
369 nearest_tile(map, near, STOCKPILE_SITE_SEARCH_TILES, |p| {
370 can_zone(map, p.as_ivec2())
371 })
372 })
373 }
374
375 pub fn removing(keys: &ButtonInput<KeyCode>) -> bool {
376 keys.any_pressed([KeyCode::ShiftLeft, KeyCode::ShiftRight])
377 }
378
379 fn zone_tool(
380 mut actions: MessageReader<Action>,
381 keys: Res<ButtonInput<KeyCode>>,
382 buttons: Res<ButtonInput<MouseButton>>,
383 over_ui: Res<PointerOverUi>,
384 cursor: Res<Cursor>,
385 map: Res<WorldMap>,
386 mut tool: ResMut<ZoneTool>,
387 mut zones: ResMut<Zones>,
388 ) {
389 for action in actions.read() {
390 match action {
391 Action::ToggleStockpileTool => tool.toggle(ZoneKind::Stockpile),
392 Action::ToggleFarmTool => tool.toggle(ZoneKind::Farm),
393 Action::Cancel
394 | Action::CycleBuildTool
395 | Action::ToggleDesignateTool
396 | Action::ToggleDemolishTool => tool.mode = None,
397 _ => continue,
398 }
399 tool.drag_start = None;
400 return;
401 }
402 let Some(kind) = tool.mode else {
403 return;
404 };
405 if buttons.just_pressed(MouseButton::Left) && !over_ui.0 {
406 tool.drag_start = cursor.tile;
407 }
408 if !buttons.just_released(MouseButton::Left) {
409 return;
410 }
411 if let (Some(start), Some(end)) = (tool.drag_start.take(), cursor.tile) {
412 if removing(&keys) {
413 zones.remove_zone(kind, start.as_ivec2(), end.as_ivec2());
414 } else {
415 zones.add_zone(kind, &map, start.as_ivec2(), end.as_ivec2());
416 }
417 }
418 }
419
420 fn draw_zone_preview(
421 keys: Res<ButtonInput<KeyCode>>,
422 cursor: Res<Cursor>,
423 tool: Res<ZoneTool>,
424 mut gizmos: Gizmos,
425 ) {
426 let (Some(start), Some(end)) = (tool.drag_start, cursor.tile) else {
427 return;
428 };
429 let (min, max) = (start.min(end), start.max(end));
430 let center = (WorldMap::tile_center(min) + WorldMap::tile_center(max)) / 2.0;
431 let size = (max - min + 1).as_vec2() * TILE_SIZE;
432 let color = if removing(&keys) {
433 PREVIEW_REMOVE_COLOR
434 } else {
435 PREVIEW_ADD_COLOR
436 };
437 gizmos.rect_2d(center, size, color);
438 }
439
440 #[cfg(test)]
441 mod tests {
442 use super::*;
443 use crate::world::Terrain;
444
445 #[test]
446 fn stockpiles_only_cover_walkable_dry_land() {
447 let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Grass);
448 map.set(UVec2::new(1, 0), Terrain::Stone);
449 map.set(UVec2::new(2, 0), Terrain::ShallowWater);
450 let mut zones = Zones::default();
451 zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(-1, 0));
452 let mut tiles: Vec<_> = zones.stockpile_tiles().collect();
453 tiles.sort_by_key(|t| t.x);
454 assert_eq!(tiles, vec![UVec2::new(0, 0), UVec2::new(3, 0)]);
455 }
456
457 #[test]
458 fn expansion_grows_compactly_around_the_stockpile() {
459 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
460 let mut zones = Zones::default();
461 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(5, 5));
462 let plan = plan_expansion(&map, &zones, &HashSet::new(), 8);
463 assert_eq!(plan.len(), 8);
464 assert!(plan.iter().all(|t| !zones.is_stockpile(*t)));
465 let beyond_one_ring = plan_expansion(&map, &zones, &HashSet::new(), 20);
466 assert_eq!(beyond_one_ring.len(), 20);
467 assert!(
468 plan.iter()
469 .all(|t| (3..=6).contains(&t.x) && (3..=6).contains(&t.y))
470 );
471 }
472
473 // Trees can be zoned, but stockpiling on open ground keeps forests intact and avoids
474 // colonists cutting trees just to make room, so expansion only takes trees as a last resort.
475 #[test]
476 fn expansion_prefers_open_ground_over_trees_and_skips_unzonable_tiles() {
477 let mut map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
478 map.set(UVec2::new(3, 4), Terrain::Stone);
479 map.set_plant(UVec2::new(5, 4), Some(crate::world::Plant::Tree));
480 let mut zones = Zones::default();
481 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(4, 4));
482 let plan = plan_expansion(&map, &zones, &HashSet::new(), 2);
483 assert_eq!(plan.len(), 2);
484 assert!(!plan.contains(&UVec2::new(3, 4)));
485 assert!(!plan.contains(&UVec2::new(5, 4)));
486 }
487
488 #[test]
489 fn no_expansion_without_a_stockpile() {
490 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
491 assert!(plan_expansion(&map, &Zones::default(), &HashSet::new(), 8).is_empty());
492 }
493
494 #[test]
495 fn farms_need_fertile_ground_and_never_overlap_stockpiles() {
496 let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Soil);
497 map.set(UVec2::new(1, 0), Terrain::Stone);
498 let mut zones = Zones::default();
499 zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(3, 0));
500 zones.add_zone(ZoneKind::Farm, &map, IVec2::ZERO, IVec2::new(3, 0));
501 let mut farm: Vec<_> = zones.farm_tiles().collect();
502 farm.sort_by_key(|t| t.x);
503 assert_eq!(farm, vec![UVec2::new(0, 0), UVec2::new(2, 0)]);
504 zones.add_stockpile(&map, IVec2::ZERO, IVec2::ZERO);
505 assert!(!zones.is_stockpile(UVec2::ZERO));
506 }
507
508 #[test]
509 fn storage_rules_apply_to_one_connected_stockpile() {
510 let map = WorldMap::filled(UVec2::new(10, 4), Terrain::Grass);
511 let mut zones = Zones::default();
512 zones.add_stockpile(&map, IVec2::ZERO, IVec2::new(2, 1));
513 zones.add_stockpile(&map, IVec2::new(6, 0), IVec2::new(8, 1));
514 zones.update_region_rule(UVec2::ZERO, |rule| {
515 rule.toggle(Stock::Stone);
516 rule.priority = StoragePriority::High;
517 });
518 assert!(!zones.accepts(UVec2::new(2, 1), ItemKind::Stone));
519 assert!(zones.accepts(UVec2::new(2, 1), ItemKind::Wood));
520 assert_eq!(zones.rule(UVec2::new(1, 0)).priority, StoragePriority::High);
521 assert!(zones.accepts(UVec2::new(7, 0), ItemKind::Stone));
522 assert_eq!(
523 zones.rule(UVec2::new(7, 0)).priority,
524 StoragePriority::Normal
525 );
526 }
527
528 #[test]
529 fn removing_clears_a_rectangle() {
530 let map = WorldMap::filled(UVec2::new(4, 4), Terrain::Grass);
531 let mut zones = Zones::default();
532 zones.add_stockpile(&map, IVec2::ZERO, IVec2::splat(3));
533 zones.remove_stockpile(IVec2::new(1, 1), IVec2::new(2, 2));
534 assert_eq!(zones.stockpile_tiles().count(), 12);
535 assert!(!zones.is_stockpile(UVec2::new(1, 2)));
536 assert!(zones.is_stockpile(UVec2::new(0, 2)));
537 }
538
539 fn tool_app() -> App {
540 let mut app = App::new();
541 app.insert_resource(WorldMap::filled(UVec2::splat(8), Terrain::Grass))
542 .init_resource::<Zones>()
543 .init_resource::<ButtonInput<KeyCode>>()
544 .init_resource::<ButtonInput<MouseButton>>()
545 .insert_resource(ZoneTool {
546 mode: Some(ZoneKind::Stockpile),
547 drag_start: None,
548 })
549 .init_resource::<Cursor>()
550 .init_resource::<PointerOverUi>()
551 .add_message::<Action>()
552 .add_systems(Update, zone_tool);
553 app
554 }
555
556 fn drag(app: &mut App, from: UVec2, to: UVec2) {
557 app.world_mut().resource_mut::<Cursor>().tile = Some(from);
558 app.world_mut()
559 .resource_mut::<ButtonInput<MouseButton>>()
560 .press(MouseButton::Left);
561 app.update();
562 app.world_mut().resource_mut::<Cursor>().tile = Some(to);
563 let mut buttons = app.world_mut().resource_mut::<ButtonInput<MouseButton>>();
564 buttons.clear();
565 buttons.release(MouseButton::Left);
566 app.update();
567 app.world_mut()
568 .resource_mut::<ButtonInput<MouseButton>>()
569 .clear();
570 }
571
572 #[test]
573 fn dragging_adds_and_shift_dragging_removes() {
574 let mut app = tool_app();
575 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
576 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 6);
577 app.world_mut()
578 .resource_mut::<ButtonInput<KeyCode>>()
579 .press(KeyCode::ShiftLeft);
580 drag(&mut app, UVec2::new(3, 2), UVec2::new(3, 1));
581 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 4);
582 }
583
584 #[test]
585 fn dragging_does_nothing_when_the_tool_is_off() {
586 let mut app = tool_app();
587 app.world_mut().resource_mut::<ZoneTool>().mode = None;
588 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
589 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 0);
590 }
591
592 #[test]
593 fn new_worlds_start_without_stockpiles() {
594 let map = WorldMap::filled(UVec2::splat(8), Terrain::Grass);
595 let mut app = App::new();
596 app.init_resource::<Zones>()
597 .add_message::<WorldGenerated>()
598 .add_systems(Update, clear_zones.run_if(on_message::<WorldGenerated>));
599 app.world_mut()
600 .resource_mut::<Zones>()
601 .add_stockpile(&map, IVec2::ZERO, IVec2::ONE);
602 app.world_mut().write_message(WorldGenerated);
603 app.update();
604 assert!(!app.world().resource::<Zones>().has_stockpile());
605 }
606
607 #[test]
608 fn new_stockpiles_go_on_the_nearest_fully_open_ground() {
609 let mut map = WorldMap::filled(UVec2::splat(12), Terrain::Grass);
610 for y in 0..12 {
611 map.set_plant(UVec2::new(4, y), Some(crate::world::Plant::Tree));
612 map.set(UVec2::new(7, y), Terrain::DeepWater);
613 }
614 let site = find_stockpile_site(&map, UVec2::new(5, 5)).unwrap();
615 assert!(new_stockpile_tiles(site).all(|t| t.x < 4));
616 }
617 }