steel_core/enchantment_helper/
selection.rs

1//! Enchanting table costs and weighted enchantment selection.
2
3use steel_math::round_f32;
4use steel_registry::data_components::vanilla_components::ENCHANTABLE;
5use steel_registry::enchantment::{Enchantment, EnchantmentRef};
6use steel_registry::item_stack::ItemStack;
7use steel_registry::vanilla_items;
8use steel_utils::random::{Random, weighted_random};
9
10const MAX_BOOKSHELVES: i32 = 15;
11const BASE_COST_RANDOM_BOUND: i32 = 8;
12const ENCHANTABILITY_RANDOM_DIVISOR: i32 = 4;
13const COST_VARIATION: f32 = 0.15;
14const ADDITIONAL_ENCHANTMENT_RANDOM_BOUND: i32 = 50;
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub(crate) struct EnchantmentInstance {
18    pub enchantment: EnchantmentRef,
19    pub level: u32,
20}
21
22impl EnchantmentInstance {
23    #[must_use]
24    pub(crate) const fn weight(&self) -> i32 {
25        self.enchantment.weight as i32
26    }
27}
28
29/// Calculates the required level for an enchanting table offer
30pub(crate) fn get_enchantment_cost<R: Random>(
31    random: &mut R,
32    slot: usize,
33    bookcases: i32,
34    stack: &ItemStack,
35) -> i32 {
36    if !stack.has(ENCHANTABLE) {
37        return 0;
38    }
39    let bookcases = bookcases.min(MAX_BOOKSHELVES);
40    let selected = random.next_i32_bounded(BASE_COST_RANDOM_BOUND)
41        + 1
42        + (bookcases >> 1)
43        + random.next_i32_bounded(bookcases + 1);
44    match slot {
45        0 => (selected / 3).max(1),
46        1 => selected * 2 / 3 + 1,
47        _ => selected.max(bookcases * 2),
48    }
49}
50
51/// Selects enchantments for an enchanting table offer
52pub(crate) fn select_enchantment<R: Random>(
53    random: &mut R,
54    stack: &ItemStack,
55    enchantment_cost: i32,
56    source: impl Iterator<Item = EnchantmentRef>,
57) -> Vec<EnchantmentInstance> {
58    let mut results = Vec::new();
59    let Some(enchantable) = stack.get(ENCHANTABLE) else {
60        return results;
61    };
62    let enchantability = enchantable.value();
63    let mut cost = enchantment_cost
64        + 1
65        + random.next_i32_bounded(enchantability / ENCHANTABILITY_RANDOM_DIVISOR + 1)
66        + random.next_i32_bounded(enchantability / ENCHANTABILITY_RANDOM_DIVISOR + 1);
67    let random_span = (random.next_f32() + random.next_f32() - 1.0) * COST_VARIATION;
68    cost = round_f32(cost as f32 + cost as f32 * random_span).max(1);
69
70    let mut candidates = get_available_enchantment_results(cost, stack, source);
71    if candidates.is_empty() {
72        return results;
73    }
74    if let Some(first) =
75        weighted_random::get_random_item(random, &candidates, EnchantmentInstance::weight)
76    {
77        results.push(*first);
78    }
79    while random.next_i32_bounded(ADDITIONAL_ENCHANTMENT_RANDOM_BOUND) <= cost {
80        if let Some(last) = results.last() {
81            filter_compatible_enchantments(&mut candidates, last);
82        }
83        if candidates.is_empty() {
84            break;
85        }
86        if let Some(pick) =
87            weighted_random::get_random_item(random, &candidates, EnchantmentInstance::weight)
88        {
89            results.push(*pick);
90        }
91        cost /= 2;
92    }
93    results
94}
95
96fn get_available_enchantment_results(
97    value: i32,
98    stack: &ItemStack,
99    source: impl Iterator<Item = EnchantmentRef>,
100) -> Vec<EnchantmentInstance> {
101    let is_book = stack.is(&vanilla_items::BOOK);
102    let mut results = Vec::new();
103    for enchantment in source {
104        if !is_book && !enchantment.is_primary_item(stack.item()) {
105            continue;
106        }
107        for level in (1..=enchantment.max_level).rev() {
108            let level_cost = level as i32;
109            if value >= enchantment.get_min_cost(level_cost)
110                && value <= enchantment.get_max_cost(level_cost)
111            {
112                results.push(EnchantmentInstance { enchantment, level });
113                break;
114            }
115        }
116    }
117    results
118}
119
120fn filter_compatible_enchantments(
121    candidates: &mut Vec<EnchantmentInstance>,
122    target: &EnchantmentInstance,
123) {
124    candidates
125        .retain(|candidate| Enchantment::are_compatible(target.enchantment, candidate.enchantment));
126}
127
128#[cfg(test)]
129mod tests {
130    use std::iter::once;
131
132    use steel_registry::{init_vanilla_registry, vanilla_enchantments, vanilla_items};
133    use steel_utils::random::legacy_random::LegacyRandom;
134
135    use super::*;
136
137    const FULL_POWER_BOOKSHELVES: i32 = 15;
138    const HIGH_OFFER_MIN_COST: i32 = 30;
139    const SHARPNESS_OVERLAP_COST: i32 = 30;
140
141    fn instance(enchantment: EnchantmentRef, level: u32) -> EnchantmentInstance {
142        EnchantmentInstance { enchantment, level }
143    }
144
145    #[test]
146    fn third_offer_costs_at_least_twice_the_bookcases() {
147        init_vanilla_registry();
148        let sword = ItemStack::new(&vanilla_items::DIAMOND_SWORD);
149        for seed in 0..64 {
150            let mut random = LegacyRandom::from_seed(seed);
151            assert!(
152                get_enchantment_cost(&mut random, 2, FULL_POWER_BOOKSHELVES, &sword)
153                    >= HIGH_OFFER_MIN_COST
154            );
155        }
156    }
157
158    #[test]
159    fn available_results_pick_the_highest_level_whose_window_contains_the_value() {
160        init_vanilla_registry();
161        let sword = ItemStack::new(&vanilla_items::DIAMOND_SWORD);
162        // Sharpness level n spans [1 + 11(n - 1), 21 + 11(n - 1)]; 30 fits levels 2 and 3.
163        let results = get_available_enchantment_results(
164            SHARPNESS_OVERLAP_COST,
165            &sword,
166            once(&vanilla_enchantments::SHARPNESS as EnchantmentRef),
167        );
168        assert_eq!(results, vec![instance(&vanilla_enchantments::SHARPNESS, 3)]);
169    }
170
171    #[test]
172    fn books_accept_enchantments_that_are_not_primary_for_other_items() {
173        init_vanilla_registry();
174        let sharpness = once(&vanilla_enchantments::SHARPNESS as EnchantmentRef);
175        let book = ItemStack::new(&vanilla_items::BOOK);
176        assert_eq!(
177            get_available_enchantment_results(SHARPNESS_OVERLAP_COST, &book, sharpness.clone())
178                .len(),
179            1
180        );
181        let axe = ItemStack::new(&vanilla_items::DIAMOND_AXE);
182        assert_eq!(
183            get_available_enchantment_results(SHARPNESS_OVERLAP_COST, &axe, sharpness),
184            Vec::new()
185        );
186    }
187}
steel_core/enchantment_helper/selection.rs

steel_core/enchantment_helper/
selection.rs
