1use std::{borrow::Cow, hash::Hash};
30
31use crate::{
32 assets::{AssetExt, BoxedError, FileAsset, load_ron},
33 comp::{Item, inventory::item},
34};
35use rand::prelude::*;
36use serde::{Deserialize, Serialize, de::DeserializeOwned};
37use tracing::warn;
38
39#[derive(Clone, Debug, PartialEq, Deserialize)]
40pub struct Lottery<T> {
41 items: Vec<(f32, T)>,
42 total: f32,
43}
44
45impl<T: DeserializeOwned + Send + Sync + 'static> FileAsset for Lottery<T> {
46 const EXTENSION: &'static str = "ron";
47
48 fn from_bytes(bytes: Cow<[u8]>) -> Result<Self, BoxedError> {
49 load_ron::<Vec<(f32, T)>>(&bytes).map(Vec::into)
50 }
51}
52
53impl<T> From<Vec<(f32, T)>> for Lottery<T> {
54 fn from(mut items: Vec<(f32, T)>) -> Lottery<T> {
55 let mut total = 0.0;
56
57 for (rate, _) in &mut items {
58 total += *rate;
59 *rate = total - *rate;
60 }
61
62 Self { items, total }
63 }
64}
65
66impl<T> Lottery<T> {
67 pub fn choose_seeded(&self, seed: u32) -> &T {
68 let x = ((seed % 65536) as f32 / 65536.0) * self.total;
69 &self.items[self
70 .items
71 .binary_search_by(|(y, _)| y.partial_cmp(&x).unwrap())
72 .unwrap_or_else(|i| i.saturating_sub(1))]
73 .1
74 }
75
76 pub fn choose(&self) -> &T { self.choose_seeded(rand::rng().random()) }
77
78 pub fn iter(&self) -> impl Iterator<Item = &(f32, T)> { self.items.iter() }
79
80 pub fn total(&self) -> f32 { self.total }
81}
82
83pub fn distribute_many<T: Copy + Eq + Hash, I>(
85 participants: impl IntoIterator<Item = (f32, T)>,
86 rng: &mut impl Rng,
87 items: &[I],
88 mut get_amount: impl FnMut(&I) -> u32,
89 mut exec_item: impl FnMut(&I, T, u32),
90) {
91 struct Participant<T> {
92 weight: f32,
94 sorted_weight: f32,
95 data: T,
96 recieved_count: u32,
97 current_recieved_count: u32,
98 }
99
100 impl<T> Participant<T> {
101 fn give(&mut self, amount: u32) {
102 self.current_recieved_count += amount;
103 self.recieved_count += amount;
104 }
105 }
106
107 if items.is_empty() {
109 return;
110 }
111
112 let mut total_weight = 0.0;
113
114 let mut participants = participants
115 .into_iter()
116 .map(|(weight, participant)| Participant {
117 weight,
118 sorted_weight: {
119 total_weight += weight;
120 total_weight - weight
121 },
122 data: participant,
123 recieved_count: 0,
124 current_recieved_count: 0,
125 })
126 .collect::<Vec<_>>();
127
128 let total_item_amount = items.iter().map(&mut get_amount).sum::<u32>();
129
130 let mut current_total_weight = total_weight;
131
132 for item in items.iter() {
133 let amount = get_amount(item);
134 let mut distributed = 0;
135
136 let Some(mut give) = participants
137 .iter()
138 .map(|participant| {
139 (total_item_amount as f32 * participant.weight / total_weight).ceil() as u32
140 - participant.recieved_count
141 })
142 .min()
143 else {
144 tracing::error!("Tried to distribute items to no participants.");
145 return;
146 };
147
148 while distributed < amount {
149 let max_give = (amount / participants.len() as u32).clamp(1, amount - distributed);
152 give = give.clamp(1, max_give);
153 let x = rng.random_range(0.0..=current_total_weight);
154
155 let index = participants
156 .binary_search_by(|item| item.sorted_weight.partial_cmp(&x).unwrap())
157 .unwrap_or_else(|i| i.saturating_sub(1));
158
159 let participant_count = participants.len();
160
161 let Some(winner) = participants.get_mut(index) else {
162 tracing::error!("Tried to distribute items to no participants.");
163 return;
164 };
165
166 winner.give(give);
167 distributed += give;
168
169 if participant_count > 1
171 && winner.recieved_count as f32 / total_item_amount as f32
172 >= winner.weight / total_weight
173 {
174 current_total_weight = index
175 .checked_sub(1)
176 .and_then(|i| Some(participants.get(i)?.sorted_weight))
177 .unwrap_or(0.0);
178 let winner = participants.swap_remove(index);
179 exec_item(item, winner.data, winner.current_recieved_count);
180
181 for participant in &mut participants[index..] {
183 current_total_weight += participant.weight;
184 participant.sorted_weight = current_total_weight - participant.weight;
185 }
186
187 give = participants
189 .iter()
190 .map(|participant| {
191 (total_item_amount as f32 * participant.weight / total_weight).ceil() as u32
192 - participant.recieved_count
193 })
194 .min()
195 .unwrap_or(0);
196 } else {
197 give = give.min(
198 (total_item_amount as f32 * winner.weight / total_weight).ceil() as u32
199 - winner.recieved_count,
200 );
201 }
202 }
203 for participant in participants.iter_mut() {
204 if participant.current_recieved_count != 0 {
205 exec_item(item, participant.data, participant.current_recieved_count);
206 participant.current_recieved_count = 0;
207 }
208 }
209 }
210}
211
212#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
213#[rustfmt::skip] #[derive(Default)]
215pub enum LootSpec<T: AsRef<str>> {
216 Item(T),
218 LootTable(T),
220 #[default]
222 Nothing,
223 ModularWeapon {
225 tool: item::tool::ToolKind,
226 material: item::Material,
227 hands: Option<item::tool::Hands>,
228 },
229 ModularWeaponPrimaryComponent {
232 tool: item::tool::ToolKind,
233 material: item::Material,
234 hands: Option<item::tool::Hands>,
235 },
236 MultiDrop(Box<LootSpec<T>>, u32, u32),
250 All(Vec<LootSpec<T>>),
274 Lottery(Vec<(f32, LootSpec<T>)>),
296}
297
298impl<T: AsRef<str>> LootSpec<T> {
299 fn to_items_inner(
300 &self,
301 rng: &mut rand::rngs::ThreadRng,
302 amount: u32,
303 items: &mut Vec<(u32, Item)>,
304 ) {
305 let convert_item = |item: &T| {
306 Item::new_from_asset(item.as_ref()).map_or_else(
307 |e| {
308 warn!(?e, "error while loading item: {}", item.as_ref());
309 None
310 },
311 Some,
312 )
313 };
314 let mut push_item = |mut item: Item, count: u32| {
315 let count = item.amount().saturating_mul(count);
316 item.set_amount(1).expect("1 is always a valid amount.");
317 let hash = item.item_hash();
318 match items.binary_search_by_key(&hash, |(_, item)| item.item_hash()) {
319 Ok(i) => {
320 let has_same_hash = |i: &usize| items[*i].1.item_hash() == hash;
327 if let Some(i) = (i..items.len())
328 .take_while(has_same_hash)
329 .chain((0..i).rev().take_while(has_same_hash))
330 .find(|i| items[*i].1 == item)
331 {
332 items[i].0 = items[i].0.saturating_add(count);
335 } else {
336 items.insert(i, (count, item));
337 }
338 },
339 Err(i) => items.insert(i, (count, item)),
340 }
341 };
342
343 match self {
344 Self::Item(item) => {
345 if let Some(item) = convert_item(item) {
346 push_item(item, amount);
347 }
348 },
349 Self::LootTable(table) => {
350 let loot_spec = Lottery::<LootSpec<String>>::load_expect(table.as_ref()).read();
351 for _ in 0..amount {
352 loot_spec.choose().to_items_inner(rng, 1, items)
353 }
354 },
355 Self::Lottery(table) => {
356 let lottery = Lottery::from(
357 table
358 .iter()
359 .map(|(weight, spec)| (*weight, spec))
360 .collect::<Vec<_>>(),
361 );
362
363 for _ in 0..amount {
364 lottery.choose().to_items_inner(rng, 1, items)
365 }
366 },
367 Self::Nothing => {},
368 Self::ModularWeapon {
369 tool,
370 material,
371 hands,
372 } => {
373 for _ in 0..amount {
374 match item::modular::random_weapon(*tool, *material, *hands, rng) {
375 Ok(item) => push_item(item, 1),
376 Err(e) => {
377 warn!(
378 ?e,
379 "error while creating modular weapon. Toolkind: {:?}, Material: \
380 {:?}, Hands: {:?}",
381 tool,
382 material,
383 hands,
384 );
385 },
386 }
387 }
388 },
389 Self::ModularWeaponPrimaryComponent {
390 tool,
391 material,
392 hands,
393 } => {
394 for _ in 0..amount {
395 match item::modular::random_weapon(*tool, *material, *hands, rng) {
396 Ok(item) => push_item(item, 1),
397 Err(e) => {
398 warn!(
399 ?e,
400 "error while creating modular weapon primary component. Toolkind: \
401 {:?}, Material: {:?}, Hands: {:?}",
402 tool,
403 material,
404 hands,
405 );
406 },
407 }
408 }
409 },
410 Self::MultiDrop(loot_spec, lower, upper) => {
411 let sub_amount = rng.random_range(*lower..=*upper);
412 loot_spec.to_items_inner(rng, sub_amount.saturating_mul(amount), items);
415 },
416 Self::All(loot_specs) => {
417 for loot_spec in loot_specs {
418 loot_spec.to_items_inner(rng, amount, items);
419 }
420 },
421 }
422 }
423
424 pub fn to_items(&self) -> Option<Vec<(u32, Item)>> {
425 let mut items = Vec::new();
426 self.to_items_inner(&mut rand::rng(), 1, &mut items);
427
428 if !items.is_empty() {
429 items.sort_unstable_by_key(|(amount, _)| *amount);
430
431 Some(items)
432 } else {
433 None
434 }
435 }
436}
437
438#[cfg(test)]
439pub mod tests {
440 use std::borrow::Borrow;
441
442 use super::*;
443 use crate::{assets, comp::Item};
444 use assets::AssetExt;
445
446 #[cfg(test)]
447 pub fn validate_loot_spec(item: &LootSpec<String>) {
448 let mut rng = rand::rng();
449 match item {
450 LootSpec::Item(item) => {
451 Item::new_from_asset_expect(item);
452 },
453 LootSpec::LootTable(loot_table) => {
454 let loot_table = Lottery::<LootSpec<String>>::load_expect(loot_table).read();
455 validate_table_contents(&loot_table);
456 },
457 LootSpec::Nothing => {},
458 LootSpec::ModularWeapon {
459 tool,
460 material,
461 hands,
462 } => {
463 item::modular::random_weapon(*tool, *material, *hands, &mut rng).unwrap_or_else(
464 |_| {
465 panic!(
466 "Failed to synthesize a modular {tool:?} made of {material:?} that \
467 had a hand restriction of {hands:?}."
468 )
469 },
470 );
471 },
472 LootSpec::ModularWeaponPrimaryComponent {
473 tool,
474 material,
475 hands,
476 } => {
477 item::modular::random_weapon_primary_component(*tool, *material, *hands, &mut rng)
478 .unwrap_or_else(|_| {
479 panic!(
480 "Failed to synthesize a modular weapon primary component: {tool:?} \
481 made of {material:?} that had a hand restriction of {hands:?}."
482 )
483 });
484 },
485 LootSpec::MultiDrop(loot_spec, lower, upper) => {
486 assert!(
487 upper >= lower,
488 "Upper quantity must be at least the value of lower quantity. Upper value: \
489 {}, low value: {}.",
490 upper,
491 lower
492 );
493 validate_loot_spec(loot_spec);
494 },
495 LootSpec::All(loot_specs) => {
496 for loot_spec in loot_specs {
497 validate_loot_spec(loot_spec);
498 }
499 },
500 LootSpec::Lottery(table) => {
501 let lottery = Lottery::from(
502 table
503 .iter()
504 .map(|(weight, spec)| (*weight, spec))
505 .collect::<Vec<_>>(),
506 );
507
508 validate_table_contents(&lottery);
509 },
510 }
511 }
512
513 fn validate_table_contents<T: Borrow<LootSpec<String>>>(table: &Lottery<T>) {
514 for (_, item) in table.iter() {
515 validate_loot_spec(item.borrow());
516 }
517 }
518
519 #[test]
520 fn test_loot_tables() {
521 let loot_tables = assets::load_rec_dir::<Lottery<LootSpec<String>>>("common.loot_tables")
522 .expect("load loot_tables");
523 for loot_table in loot_tables.read().ids() {
524 let loot_table = Lottery::<LootSpec<String>>::load_expect(loot_table);
525 validate_table_contents(&loot_table.read());
526 }
527 }
528
529 #[test]
530 fn test_distribute_many() {
531 let mut rng = rand::rng();
532
533 for _ in 0..10 {
535 distribute_many(
536 vec![(0.4f32, "a"), (0.4, "b"), (0.2, "c")],
537 &mut rng,
538 &[("item", 10)],
539 |(_, m)| *m,
540 |_item, winner, count| match winner {
541 "a" | "b" => assert_eq!(count, 4),
542 "c" => assert_eq!(count, 2),
543 _ => unreachable!(),
544 },
545 );
546 }
547 }
548}