src/zones.rs 15.2 K · 464 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 => tool.mode = None,
264 _ => continue,
265 }
266 tool.drag_start = None;
267 return;
268 }
269 let Some(kind) = tool.mode else {
270 return;
271 };
272 if buttons.just_pressed(MouseButton::Left) && !over_ui.0 {
273 tool.drag_start = cursor.tile;
274 }
275 if !buttons.just_released(MouseButton::Left) {
276 return;
277 }
278 if let (Some(start), Some(end)) = (tool.drag_start.take(), cursor.tile) {
279 if removing(&keys) {
280 zones.remove_zone(kind, start.as_ivec2(), end.as_ivec2());
281 } else {
282 zones.add_zone(kind, &map, start.as_ivec2(), end.as_ivec2());
283 }
284 }
285 }
286
287 fn draw_zone_preview(
288 keys: Res<ButtonInput<KeyCode>>,
289 cursor: Res<Cursor>,
290 tool: Res<ZoneTool>,
291 mut gizmos: Gizmos,
292 ) {
293 let (Some(start), Some(end)) = (tool.drag_start, cursor.tile) else {
294 return;
295 };
296 let (min, max) = (start.min(end), start.max(end));
297 let center = (WorldMap::tile_center(min) + WorldMap::tile_center(max)) / 2.0;
298 let size = (max - min + 1).as_vec2() * TILE_SIZE;
299 let color = if removing(&keys) {
300 PREVIEW_REMOVE_COLOR
301 } else {
302 PREVIEW_ADD_COLOR
303 };
304 gizmos.rect_2d(center, size, color);
305 }
306
307 #[cfg(test)]
308 mod tests {
309 use super::*;
310 use crate::world::Terrain;
311
312 #[test]
313 fn stockpiles_only_cover_walkable_dry_land() {
314 let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Grass);
315 map.set(UVec2::new(1, 0), Terrain::Stone);
316 map.set(UVec2::new(2, 0), Terrain::ShallowWater);
317 let mut zones = Zones::default();
318 zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(-1, 0));
319 let mut tiles: Vec<_> = zones.stockpile_tiles().collect();
320 tiles.sort_by_key(|t| t.x);
321 assert_eq!(tiles, vec![UVec2::new(0, 0), UVec2::new(3, 0)]);
322 }
323
324 #[test]
325 fn expansion_grows_compactly_around_the_stockpile() {
326 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
327 let mut zones = Zones::default();
328 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(5, 5));
329 let plan = plan_expansion(&map, &zones, &HashSet::new(), 8);
330 assert_eq!(plan.len(), 8);
331 assert!(plan.iter().all(|t| !zones.is_stockpile(*t)));
332 let beyond_one_ring = plan_expansion(&map, &zones, &HashSet::new(), 20);
333 assert_eq!(beyond_one_ring.len(), 20);
334 assert!(
335 plan.iter()
336 .all(|t| (3..=6).contains(&t.x) && (3..=6).contains(&t.y))
337 );
338 }
339
340 // Trees can be zoned, but stockpiling on open ground keeps forests intact and avoids
341 // colonists cutting trees just to make room, so expansion only takes trees as a last resort.
342 #[test]
343 fn expansion_prefers_open_ground_over_trees_and_skips_unzonable_tiles() {
344 let mut map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
345 map.set(UVec2::new(3, 4), Terrain::Stone);
346 map.set_plant(UVec2::new(5, 4), Some(crate::world::Plant::Tree));
347 let mut zones = Zones::default();
348 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(4, 4));
349 let plan = plan_expansion(&map, &zones, &HashSet::new(), 2);
350 assert_eq!(plan.len(), 2);
351 assert!(!plan.contains(&UVec2::new(3, 4)));
352 assert!(!plan.contains(&UVec2::new(5, 4)));
353 }
354
355 #[test]
356 fn no_expansion_without_a_stockpile() {
357 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
358 assert!(plan_expansion(&map, &Zones::default(), &HashSet::new(), 8).is_empty());
359 }
360
361 #[test]
362 fn farms_need_fertile_ground_and_never_overlap_stockpiles() {
363 let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Soil);
364 map.set(UVec2::new(1, 0), Terrain::Stone);
365 let mut zones = Zones::default();
366 zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(3, 0));
367 zones.add_zone(ZoneKind::Farm, &map, IVec2::ZERO, IVec2::new(3, 0));
368 let mut farm: Vec<_> = zones.farm_tiles().collect();
369 farm.sort_by_key(|t| t.x);
370 assert_eq!(farm, vec![UVec2::new(0, 0), UVec2::new(2, 0)]);
371 zones.add_stockpile(&map, IVec2::ZERO, IVec2::ZERO);
372 assert!(!zones.is_stockpile(UVec2::ZERO));
373 }
374
375 #[test]
376 fn removing_clears_a_rectangle() {
377 let map = WorldMap::filled(UVec2::new(4, 4), Terrain::Grass);
378 let mut zones = Zones::default();
379 zones.add_stockpile(&map, IVec2::ZERO, IVec2::splat(3));
380 zones.remove_stockpile(IVec2::new(1, 1), IVec2::new(2, 2));
381 assert_eq!(zones.stockpile_tiles().count(), 12);
382 assert!(!zones.is_stockpile(UVec2::new(1, 2)));
383 assert!(zones.is_stockpile(UVec2::new(0, 2)));
384 }
385
386 fn tool_app() -> App {
387 let mut app = App::new();
388 app.insert_resource(WorldMap::filled(UVec2::splat(8), Terrain::Grass))
389 .init_resource::<Zones>()
390 .init_resource::<ButtonInput<KeyCode>>()
391 .init_resource::<ButtonInput<MouseButton>>()
392 .insert_resource(ZoneTool {
393 mode: Some(ZoneKind::Stockpile),
394 drag_start: None,
395 })
396 .init_resource::<Cursor>()
397 .init_resource::<PointerOverUi>()
398 .add_message::<Action>()
399 .add_systems(Update, zone_tool);
400 app
401 }
402
403 fn drag(app: &mut App, from: UVec2, to: UVec2) {
404 app.world_mut().resource_mut::<Cursor>().tile = Some(from);
405 app.world_mut()
406 .resource_mut::<ButtonInput<MouseButton>>()
407 .press(MouseButton::Left);
408 app.update();
409 app.world_mut().resource_mut::<Cursor>().tile = Some(to);
410 let mut buttons = app.world_mut().resource_mut::<ButtonInput<MouseButton>>();
411 buttons.clear();
412 buttons.release(MouseButton::Left);
413 app.update();
414 app.world_mut()
415 .resource_mut::<ButtonInput<MouseButton>>()
416 .clear();
417 }
418
419 #[test]
420 fn dragging_adds_and_shift_dragging_removes() {
421 let mut app = tool_app();
422 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
423 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 6);
424 app.world_mut()
425 .resource_mut::<ButtonInput<KeyCode>>()
426 .press(KeyCode::ShiftLeft);
427 drag(&mut app, UVec2::new(3, 2), UVec2::new(3, 1));
428 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 4);
429 }
430
431 #[test]
432 fn dragging_does_nothing_when_the_tool_is_off() {
433 let mut app = tool_app();
434 app.world_mut().resource_mut::<ZoneTool>().mode = None;
435 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
436 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 0);
437 }
438
439 #[test]
440 fn new_worlds_start_without_stockpiles() {
441 let map = WorldMap::filled(UVec2::splat(8), Terrain::Grass);
442 let mut app = App::new();
443 app.init_resource::<Zones>()
444 .add_message::<WorldGenerated>()
445 .add_systems(Update, clear_zones.run_if(on_message::<WorldGenerated>));
446 app.world_mut()
447 .resource_mut::<Zones>()
448 .add_stockpile(&map, IVec2::ZERO, IVec2::ONE);
449 app.world_mut().write_message(WorldGenerated);
450 app.update();
451 assert!(!app.world().resource::<Zones>().has_stockpile());
452 }
453
454 #[test]
455 fn new_stockpiles_go_on_the_nearest_fully_open_ground() {
456 let mut map = WorldMap::filled(UVec2::splat(12), Terrain::Grass);
457 for y in 0..12 {
458 map.set_plant(UVec2::new(4, y), Some(crate::world::Plant::Tree));
459 map.set(UVec2::new(7, y), Terrain::DeepWater);
460 }
461 let site = find_stockpile_site(&map, UVec2::new(5, 5)).unwrap();
462 assert!(new_stockpile_tiles(site).all(|t| t.x < 4));
463 }
464 }