src/economy.rs 7.2 K · 251 lines · raw · history

1 use std::collections::HashMap;
2
3 use bevy::prelude::*;
4
5 use crate::{
6 building::Blueprints,
7 items::ItemKind,
8 sim::{SimSystems, SimTick, TICKS_PER_HOUR},
9 };
10
11 pub struct EconomyPlugin;
12
13 impl Plugin for EconomyPlugin {
14 fn build(&self, app: &mut App) {
15 app.init_resource::<ColonyGoals>()
16 .add_systems(FixedUpdate, learn_needs.in_set(SimSystems::Decide));
17 }
18 }
19
20 const HOURS_OF_MEMORY: f32 = 24.0;
21 const MAX_GOAL: f32 = 300.0;
22
23 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
24 pub enum Stock {
25 Wood,
26 Stone,
27 Iron,
28 Gold,
29 RawFood,
30 Meals,
31 }
32
33 impl Stock {
34 pub const ALL: [Stock; 6] = [
35 Stock::Wood,
36 Stock::Stone,
37 Stock::Iron,
38 Stock::Gold,
39 Stock::RawFood,
40 Stock::Meals,
41 ];
42
43 pub fn of(kind: ItemKind) -> Stock {
44 match kind {
45 ItemKind::Wood => Stock::Wood,
46 ItemKind::Stone => Stock::Stone,
47 ItemKind::Iron => Stock::Iron,
48 ItemKind::Gold => Stock::Gold,
49 ItemKind::Berries | ItemKind::Meat | ItemKind::Vegetables => Stock::RawFood,
50 ItemKind::Meal => Stock::Meals,
51 }
52 }
53
54 pub fn members(self) -> impl Iterator<Item = ItemKind> {
55 ItemKind::ALL
56 .into_iter()
57 .filter(move |&k| Stock::of(k) == self)
58 }
59
60 pub fn is_food(self) -> bool {
61 matches!(self, Stock::RawFood | Stock::Meals)
62 }
63
64 pub fn label(self) -> &'static str {
65 match self {
66 Stock::Wood => "Wood",
67 Stock::Stone => "Stone",
68 Stock::Iron => "Iron",
69 Stock::Gold => "Gold",
70 Stock::RawFood => "Raw food",
71 Stock::Meals => "Meals",
72 }
73 }
74
75 fn buffer_days(self) -> f32 {
76 if self.is_food() { 1.5 } else { 2.0 }
77 }
78
79 fn minimum(self) -> f32 {
80 match self {
81 Stock::Wood => 15.0,
82 Stock::Stone => 5.0,
83 Stock::Iron | Stock::Gold => 0.0,
84 Stock::RawFood => 10.0,
85 Stock::Meals => 3.0,
86 }
87 }
88
89 fn initial_use_per_day(self) -> f32 {
90 match self {
91 Stock::Wood => 20.0,
92 Stock::Stone => 10.0,
93 Stock::Iron | Stock::Gold => 0.0,
94 Stock::RawFood => 30.0,
95 Stock::Meals => 3.0,
96 }
97 }
98 }
99
100 #[derive(Resource, Debug, Clone)]
101 pub struct ColonyGoals {
102 goals: HashMap<Stock, u32>,
103 use_per_day: HashMap<Stock, f32>,
104 used_this_hour: HashMap<Stock, u32>,
105 }
106
107 impl Default for ColonyGoals {
108 fn default() -> Self {
109 let mut goals = Self {
110 goals: HashMap::new(),
111 use_per_day: Stock::ALL
112 .iter()
113 .map(|&s| (s, s.initial_use_per_day()))
114 .collect(),
115 used_this_hour: HashMap::new(),
116 };
117 goals.recompute(&HashMap::new());
118 goals
119 }
120 }
121
122 impl ColonyGoals {
123 #[cfg(test)]
124 pub fn fixed(goals: &[(ItemKind, u32)]) -> Self {
125 Self {
126 goals: goals.iter().map(|&(k, n)| (Stock::of(k), n)).collect(),
127 use_per_day: HashMap::new(),
128 used_this_hour: HashMap::new(),
129 }
130 }
131
132 pub fn goal(&self, stock: Stock) -> u32 {
133 self.goals.get(&stock).copied().unwrap_or(0)
134 }
135
136 pub fn use_per_day(&self, stock: Stock) -> f32 {
137 self.use_per_day.get(&stock).copied().unwrap_or(0.0)
138 }
139
140 pub fn record_use(&mut self, kind: ItemKind, count: u32) {
141 *self.used_this_hour.entry(Stock::of(kind)).or_insert(0) += count;
142 }
143
144 pub fn wanted(
145 &self,
146 kind: ItemKind,
147 stocks: &HashMap<ItemKind, u32>,
148 on_the_way: &HashMap<ItemKind, u32>,
149 ) -> bool {
150 let stock = Stock::of(kind);
151 let have: u32 = stock
152 .members()
153 .map(|k| {
154 stocks.get(&k).copied().unwrap_or(0) + on_the_way.get(&k).copied().unwrap_or(0)
155 })
156 .sum();
157 have < self.goal(stock)
158 }
159
160 pub fn learn_hour(&mut self, planned: &HashMap<Stock, u32>) {
161 for stock in Stock::ALL {
162 let used = self.used_this_hour.remove(&stock).unwrap_or(0) as f32;
163 let rate = self.use_per_day.entry(stock).or_insert(0.0);
164 *rate += (used * 24.0 - *rate) / HOURS_OF_MEMORY;
165 }
166 self.recompute(planned);
167 }
168
169 fn recompute(&mut self, planned: &HashMap<Stock, u32>) {
170 for stock in Stock::ALL {
171 let learned = (self.use_per_day(stock) * stock.buffer_days())
172 .clamp(stock.minimum(), MAX_GOAL)
173 .ceil() as u32;
174 let goal = learned + planned.get(&stock).copied().unwrap_or(0);
175 self.goals.insert(stock, goal);
176 }
177 }
178 }
179
180 pub fn planned_needs(blueprints: &Blueprints) -> HashMap<Stock, u32> {
181 let mut needs = HashMap::new();
182 for kind in blueprints.0.values() {
183 let (item, count) = kind.cost();
184 *needs.entry(Stock::of(item)).or_insert(0) += count;
185 }
186 needs
187 }
188
189 fn learn_needs(tick: Res<SimTick>, blueprints: Res<Blueprints>, mut goals: ResMut<ColonyGoals>) {
190 if tick.0 == 0 || !tick.0.is_multiple_of(TICKS_PER_HOUR) {
191 return;
192 }
193 goals.learn_hour(&planned_needs(&blueprints));
194 }
195
196 #[cfg(test)]
197 mod tests {
198 use super::*;
199 use crate::building::BuildKind;
200
201 fn learn_days(goals: &mut ColonyGoals, kind: ItemKind, per_hour: u32, days: u32) {
202 for _ in 0..days * 24 {
203 goals.record_use(kind, per_hour);
204 goals.learn_hour(&HashMap::new());
205 }
206 }
207
208 #[test]
209 fn targets_follow_what_the_colony_actually_uses() {
210 let mut goals = ColonyGoals::default();
211 learn_days(&mut goals, ItemKind::Meal, 1, 5);
212 let meals = goals.goal(Stock::Meals);
213 assert!((32..=40).contains(&meals), "{meals}");
214 learn_days(&mut goals, ItemKind::Meal, 0, 5);
215 assert_eq!(goals.goal(Stock::Meals), 3);
216 }
217
218 #[test]
219 fn unused_materials_are_not_stockpiled() {
220 let mut goals = ColonyGoals::default();
221 learn_days(&mut goals, ItemKind::Wood, 0, 10);
222 assert_eq!(goals.goal(Stock::Wood), 15);
223 assert_eq!(goals.goal(Stock::Iron), 0);
224 }
225
226 // Raw foods stand in for each other (cooking accepts any), so the colony targets them as one
227 // group and is free to forage, hunt or farm whichever is at hand.
228 #[test]
229 fn raw_foods_share_one_target() {
230 let goals = ColonyGoals::fixed(&[(ItemKind::Berries, 20)]);
231 let stocks = HashMap::from([(ItemKind::Meat, 15), (ItemKind::Vegetables, 4)]);
232 assert!(goals.wanted(ItemKind::Berries, &stocks, &HashMap::new()));
233 let on_the_way = HashMap::from([(ItemKind::Berries, 1)]);
234 assert!(!goals.wanted(ItemKind::Meat, &stocks, &on_the_way));
235 }
236
237 #[test]
238 fn planned_construction_raises_material_targets() {
239 let mut blueprints = Blueprints::default();
240 for x in 0..10 {
241 blueprints.0.insert(
242 UVec2::new(x, 0),
243 BuildKind::Structure(crate::structures::StructureKind::Wall),
244 );
245 }
246 let mut goals = ColonyGoals::default();
247 let before = goals.goal(Stock::Stone);
248 goals.learn_hour(&planned_needs(&blueprints));
249 assert!(goals.goal(Stock::Stone) >= before + 19);
250 }
251 }