src/zones.rs 11.1 K · 335 lines · raw · history

1 use std::collections::HashSet;
2
3 use bevy::prelude::*;
4
5 use crate::{
6 actions::Action,
7 selection::{Cursor, update_cursor},
8 ui::PointerOverUi,
9 world::{TILE_SIZE, WorldMap},
10 worldgen::{WorldGenerated, spawn_tiles},
11 };
12
13 pub struct ZonesPlugin;
14
15 impl Plugin for ZonesPlugin {
16 fn build(&self, app: &mut App) {
17 app.init_resource::<Zones>()
18 .init_resource::<ZoneTool>()
19 .add_systems(
20 Update,
21 place_starting_stockpile.run_if(on_message::<WorldGenerated>),
22 )
23 .add_systems(
24 PostUpdate,
25 (zone_tool, draw_zone_preview).chain().after(update_cursor),
26 );
27 }
28 }
29
30 const STARTING_STOCKPILE_OFFSET: IVec2 = IVec2::new(2, -2);
31 const STARTING_STOCKPILE_SIZE: IVec2 = IVec2::new(4, 4);
32 const PREVIEW_ADD_COLOR: Color = Color::srgba(0.4, 0.7, 1.0, 0.9);
33 const PREVIEW_REMOVE_COLOR: Color = Color::srgba(1.0, 0.4, 0.4, 0.9);
34
35 #[derive(Resource, Default, Debug, Clone, PartialEq)]
36 pub struct Zones {
37 stockpile: HashSet<UVec2>,
38 }
39
40 impl Zones {
41 pub fn is_stockpile(&self, tile: UVec2) -> bool {
42 self.stockpile.contains(&tile)
43 }
44
45 #[cfg(test)]
46 pub fn stockpile_tiles(&self) -> impl Iterator<Item = UVec2> + '_ {
47 self.stockpile.iter().copied()
48 }
49
50 pub fn has_stockpile(&self) -> bool {
51 !self.stockpile.is_empty()
52 }
53
54 pub fn empty_stockpile_tiles(&self, occupied: &HashSet<UVec2>) -> usize {
55 self.stockpile.difference(occupied).count()
56 }
57
58 pub fn add_stockpile(&mut self, map: &WorldMap, a: IVec2, b: IVec2) {
59 self.add_tiles(map, rect_tiles(a, b).filter_map(|t| t.try_into().ok()));
60 }
61
62 pub fn add_tiles(&mut self, map: &WorldMap, tiles: impl IntoIterator<Item = UVec2>) {
63 for tile in tiles {
64 if can_zone(map, tile.as_ivec2()) {
65 self.stockpile.insert(tile);
66 }
67 }
68 }
69
70 pub fn remove_stockpile(&mut self, a: IVec2, b: IVec2) {
71 for tile in rect_tiles(a, b) {
72 if tile.x >= 0 && tile.y >= 0 {
73 self.stockpile.remove(&tile.as_uvec2());
74 }
75 }
76 }
77 }
78
79 pub fn plan_expansion(map: &WorldMap, zones: &Zones, count: usize) -> Vec<UVec2> {
80 if zones.stockpile.is_empty() {
81 return Vec::new();
82 }
83 let centroid =
84 zones.stockpile.iter().map(|t| t.as_vec2()).sum::<Vec2>() / zones.stockpile.len() as f32;
85 let preference = |t: UVec2| {
86 (
87 map.plant(t.as_ivec2()).is_some(),
88 t.as_vec2().distance_squared(centroid),
89 t.y,
90 t.x,
91 )
92 };
93 let mut planned: Vec<UVec2> = Vec::new();
94 while planned.len() < count {
95 let taken = |t: UVec2| zones.is_stockpile(t) || planned.contains(&t);
96 let best = zones
97 .stockpile
98 .iter()
99 .chain(planned.iter())
100 .flat_map(|t| {
101 [IVec2::X, IVec2::NEG_X, IVec2::Y, IVec2::NEG_Y].map(|d| t.as_ivec2() + d)
102 })
103 .filter(|&t| can_zone(map, t) && !taken(t.as_uvec2()))
104 .map(|t| t.as_uvec2())
105 .min_by(|&a, &b| {
106 let (ka, kb) = (preference(a), preference(b));
107 ka.0.cmp(&kb.0)
108 .then(ka.1.total_cmp(&kb.1))
109 .then((ka.2, ka.3).cmp(&(kb.2, kb.3)))
110 });
111 match best {
112 Some(tile) => planned.push(tile),
113 None => break,
114 }
115 }
116 planned
117 }
118
119 pub fn can_zone(map: &WorldMap, tile: IVec2) -> bool {
120 map.move_cost(tile).is_some() && !map.get(tile).is_some_and(|t| t.is_water())
121 }
122
123 fn rect_tiles(a: IVec2, b: IVec2) -> impl Iterator<Item = IVec2> {
124 let (min, max) = (a.min(b), a.max(b));
125 (min.y..=max.y).flat_map(move |y| (min.x..=max.x).map(move |x| IVec2::new(x, y)))
126 }
127
128 #[derive(Resource, Default, Debug)]
129 pub struct ZoneTool {
130 pub active: bool,
131 drag_start: Option<UVec2>,
132 }
133
134 fn place_starting_stockpile(map: Res<WorldMap>, mut zones: ResMut<Zones>) {
135 *zones = Zones::default();
136 if let Some(&spawn) = spawn_tiles(&map, 1).first() {
137 let corner = spawn.as_ivec2() + STARTING_STOCKPILE_OFFSET;
138 zones.add_stockpile(&map, corner, corner + STARTING_STOCKPILE_SIZE - 1);
139 }
140 }
141
142 fn removing(keys: &ButtonInput<KeyCode>) -> bool {
143 keys.any_pressed([KeyCode::ShiftLeft, KeyCode::ShiftRight])
144 }
145
146 fn zone_tool(
147 mut actions: MessageReader<Action>,
148 keys: Res<ButtonInput<KeyCode>>,
149 buttons: Res<ButtonInput<MouseButton>>,
150 over_ui: Res<PointerOverUi>,
151 cursor: Res<Cursor>,
152 map: Res<WorldMap>,
153 mut tool: ResMut<ZoneTool>,
154 mut zones: ResMut<Zones>,
155 ) {
156 for action in actions.read() {
157 match action {
158 Action::ToggleStockpileTool => tool.active = !tool.active,
159 Action::Cancel => tool.active = false,
160 _ => continue,
161 }
162 tool.drag_start = None;
163 return;
164 }
165 if !tool.active {
166 return;
167 }
168 if buttons.just_pressed(MouseButton::Left) && !over_ui.0 {
169 tool.drag_start = cursor.tile;
170 }
171 if !buttons.just_released(MouseButton::Left) {
172 return;
173 }
174 if let (Some(start), Some(end)) = (tool.drag_start.take(), cursor.tile) {
175 if removing(&keys) {
176 zones.remove_stockpile(start.as_ivec2(), end.as_ivec2());
177 } else {
178 zones.add_stockpile(&map, start.as_ivec2(), end.as_ivec2());
179 }
180 }
181 }
182
183 fn draw_zone_preview(
184 keys: Res<ButtonInput<KeyCode>>,
185 cursor: Res<Cursor>,
186 tool: Res<ZoneTool>,
187 mut gizmos: Gizmos,
188 ) {
189 let (Some(start), Some(end)) = (tool.drag_start, cursor.tile) else {
190 return;
191 };
192 let (min, max) = (start.min(end), start.max(end));
193 let center = (WorldMap::tile_center(min) + WorldMap::tile_center(max)) / 2.0;
194 let size = (max - min + 1).as_vec2() * TILE_SIZE;
195 let color = if removing(&keys) {
196 PREVIEW_REMOVE_COLOR
197 } else {
198 PREVIEW_ADD_COLOR
199 };
200 gizmos.rect_2d(center, size, color);
201 }
202
203 #[cfg(test)]
204 mod tests {
205 use super::*;
206 use crate::world::Terrain;
207
208 #[test]
209 fn stockpiles_only_cover_walkable_dry_land() {
210 let mut map = WorldMap::filled(UVec2::new(4, 1), Terrain::Grass);
211 map.set(UVec2::new(1, 0), Terrain::Stone);
212 map.set(UVec2::new(2, 0), Terrain::ShallowWater);
213 let mut zones = Zones::default();
214 zones.add_stockpile(&map, IVec2::new(3, 0), IVec2::new(-1, 0));
215 let mut tiles: Vec<_> = zones.stockpile_tiles().collect();
216 tiles.sort_by_key(|t| t.x);
217 assert_eq!(tiles, vec![UVec2::new(0, 0), UVec2::new(3, 0)]);
218 }
219
220 #[test]
221 fn expansion_grows_compactly_around_the_stockpile() {
222 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
223 let mut zones = Zones::default();
224 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(5, 5));
225 let plan = plan_expansion(&map, &zones, 8);
226 assert_eq!(plan.len(), 8);
227 assert!(plan.iter().all(|t| !zones.is_stockpile(*t)));
228 let beyond_one_ring = plan_expansion(&map, &zones, 20);
229 assert_eq!(beyond_one_ring.len(), 20);
230 assert!(
231 plan.iter()
232 .all(|t| (3..=6).contains(&t.x) && (3..=6).contains(&t.y))
233 );
234 }
235
236 // Trees can be zoned, but stockpiling on open ground keeps forests intact and avoids
237 // colonists cutting trees just to make room, so expansion only takes trees as a last resort.
238 #[test]
239 fn expansion_prefers_open_ground_over_trees_and_skips_unzonable_tiles() {
240 let mut map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
241 map.set(UVec2::new(3, 4), Terrain::Stone);
242 map.set_plant(UVec2::new(5, 4), Some(crate::world::Plant::Tree));
243 let mut zones = Zones::default();
244 zones.add_stockpile(&map, IVec2::new(4, 4), IVec2::new(4, 4));
245 let plan = plan_expansion(&map, &zones, 2);
246 assert_eq!(plan.len(), 2);
247 assert!(!plan.contains(&UVec2::new(3, 4)));
248 assert!(!plan.contains(&UVec2::new(5, 4)));
249 }
250
251 #[test]
252 fn no_expansion_without_a_stockpile() {
253 let map = WorldMap::filled(UVec2::splat(10), Terrain::Grass);
254 assert!(plan_expansion(&map, &Zones::default(), 8).is_empty());
255 }
256
257 #[test]
258 fn removing_clears_a_rectangle() {
259 let map = WorldMap::filled(UVec2::new(4, 4), Terrain::Grass);
260 let mut zones = Zones::default();
261 zones.add_stockpile(&map, IVec2::ZERO, IVec2::splat(3));
262 zones.remove_stockpile(IVec2::new(1, 1), IVec2::new(2, 2));
263 assert_eq!(zones.stockpile_tiles().count(), 12);
264 assert!(!zones.is_stockpile(UVec2::new(1, 2)));
265 assert!(zones.is_stockpile(UVec2::new(0, 2)));
266 }
267
268 fn tool_app() -> App {
269 let mut app = App::new();
270 app.insert_resource(WorldMap::filled(UVec2::splat(8), Terrain::Grass))
271 .init_resource::<Zones>()
272 .init_resource::<ButtonInput<KeyCode>>()
273 .init_resource::<ButtonInput<MouseButton>>()
274 .insert_resource(ZoneTool {
275 active: true,
276 drag_start: None,
277 })
278 .init_resource::<Cursor>()
279 .init_resource::<PointerOverUi>()
280 .add_message::<Action>()
281 .add_systems(Update, zone_tool);
282 app
283 }
284
285 fn drag(app: &mut App, from: UVec2, to: UVec2) {
286 app.world_mut().resource_mut::<Cursor>().tile = Some(from);
287 app.world_mut()
288 .resource_mut::<ButtonInput<MouseButton>>()
289 .press(MouseButton::Left);
290 app.update();
291 app.world_mut().resource_mut::<Cursor>().tile = Some(to);
292 let mut buttons = app.world_mut().resource_mut::<ButtonInput<MouseButton>>();
293 buttons.clear();
294 buttons.release(MouseButton::Left);
295 app.update();
296 app.world_mut()
297 .resource_mut::<ButtonInput<MouseButton>>()
298 .clear();
299 }
300
301 #[test]
302 fn dragging_adds_and_shift_dragging_removes() {
303 let mut app = tool_app();
304 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
305 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 6);
306 app.world_mut()
307 .resource_mut::<ButtonInput<KeyCode>>()
308 .press(KeyCode::ShiftLeft);
309 drag(&mut app, UVec2::new(3, 2), UVec2::new(3, 1));
310 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 4);
311 }
312
313 #[test]
314 fn dragging_does_nothing_when_the_tool_is_off() {
315 let mut app = tool_app();
316 app.world_mut().resource_mut::<ZoneTool>().active = false;
317 drag(&mut app, UVec2::new(1, 1), UVec2::new(3, 2));
318 assert_eq!(app.world().resource::<Zones>().stockpile_tiles().count(), 0);
319 }
320
321 #[test]
322 fn new_worlds_start_with_a_stockpile() {
323 let mut app = App::new();
324 app.insert_resource(crate::worldgen::generate(UVec2::new(96, 64), 1))
325 .init_resource::<Zones>()
326 .add_message::<WorldGenerated>()
327 .add_systems(
328 Update,
329 place_starting_stockpile.run_if(on_message::<WorldGenerated>),
330 );
331 app.world_mut().write_message(WorldGenerated);
332 app.update();
333 assert!(app.world().resource::<Zones>().stockpile_tiles().count() > 8);
334 }
335 }