Skip to main content

steel_registry/loot_table/
functions.rs

1use super::{
2    DyeColor, EquipmentSlotGroup, Identifier, InstrumentRef, ItemPredicate, ItemStack,
3    LootCondition, LootContext, LootContextEntity, LootEntry, NumberProvider, REGISTRY, RngExt,
4    TaggedRegistryExt, context::math_round,
5};
6
7/// Options for selecting enchantments - either a tag reference or explicit list.
8#[derive(Debug, Clone)]
9pub enum EnchantmentOptions {
10    /// Reference to an enchantment tag (e.g., "`on_random_loot`").
11    Tag(Identifier),
12    /// Explicit list of enchantment IDs.
13    List(&'static [Identifier]),
14}
15
16/// Options for selecting an instrument from a registry tag or explicit list.
17#[derive(Debug, Clone)]
18pub enum InstrumentOptions {
19    Tag(Identifier),
20    Direct(&'static [InstrumentRef]),
21}
22
23impl InstrumentOptions {
24    fn get_random<R: rand::Rng>(&self, rng: &mut R) -> Option<InstrumentRef> {
25        match self {
26            Self::Tag(tag) => {
27                let instruments = REGISTRY.instruments.get_tag(tag)?;
28                (!instruments.is_empty()).then(|| {
29                    let index = rng.random_range(0..instruments.len());
30                    instruments[index]
31                })
32            }
33            Self::Direct(instruments) => (!instruments.is_empty()).then(|| {
34                let index = rng.random_range(0..instruments.len());
35                instruments[index]
36            }),
37        }
38    }
39}
40
41/// A function with optional conditions.
42#[derive(Debug, Clone)]
43pub struct ConditionalLootFunction {
44    pub function: LootFunction,
45    pub conditions: &'static [LootCondition],
46}
47
48/// A function that modifies loot items.
49#[derive(Debug, Clone)]
50pub enum LootFunction {
51    /// Set the count of the item.
52    SetCount { count: NumberProvider, add: bool },
53    /// Apply explosion decay - each item has 1/radius chance to survive.
54    ExplosionDecay,
55    /// Apply bonus count based on enchantment level.
56    ApplyBonus {
57        enchantment: Identifier,
58        formula: BonusFormula,
59    },
60    /// Increase count based on enchantment (like looting).
61    EnchantedCountIncrease {
62        enchantment: Identifier,
63        count: NumberProvider,
64        limit: i32,
65    },
66    /// Limit the count to a range.
67    LimitCount { min: Option<i32>, max: Option<i32> },
68    /// Set the damage of the item (0.0 = broken, 1.0 = full durability).
69    SetDamage { damage: NumberProvider, add: bool },
70    /// Enchant the item randomly with enchantments from options.
71    EnchantRandomly { options: EnchantmentOptions },
72    /// Enchant the item as if using an enchanting table at the specified level.
73    EnchantWithLevels {
74        levels: NumberProvider,
75        options: EnchantmentOptions,
76    },
77    /// Copy components from the block entity to the item.
78    CopyComponents {
79        source: CopySource,
80        include: &'static [Identifier],
81    },
82    /// Copy block state properties to the item.
83    CopyState {
84        block: Identifier,
85        properties: &'static [&'static str],
86    },
87    /// Set components on the item.
88    SetComponents { components: &'static str },
89    /// Set custom NBT data on the item (merges with existing `custom_data`).
90    SetCustomData {
91        tag: fn() -> crate::data_components::CustomData,
92    },
93    /// Smelt the item (convert raw to cooked, ore to ingot, etc.).
94    FurnaceSmelt { use_input_count: bool },
95    /// Create an exploration map pointing to a structure.
96    ExplorationMap {
97        destination: Identifier,
98        decoration: Identifier,
99        zoom: i32,
100        skip_existing_chunks: bool,
101    },
102    /// Set the custom name of the item.
103    SetName {
104        name: &'static str,
105        target: NameTarget,
106    },
107    /// Set the ominous bottle amplifier.
108    SetOminousBottleAmplifier { amplifier: NumberProvider },
109    /// Set the potion type.
110    SetPotion { id: Identifier },
111    /// Set the suspicious stew effects.
112    SetStewEffect { effects: &'static [StewEffect] },
113    /// Set the instrument for goat horns.
114    SetInstrument { options: InstrumentOptions },
115    /// Set enchantments on the item.
116    SetEnchantments {
117        enchantments: &'static [(Identifier, NumberProvider)],
118        add: bool,
119    },
120    /// Change the item type entirely.
121    SetItem { item: Identifier },
122    /// Copy name from source entity/block to item.
123    CopyName { source: CopySource },
124    /// Add lore lines to the item.
125    SetLore {
126        lore: &'static [&'static str],
127        mode: ListOperation,
128    },
129    /// Set container inventory contents.
130    SetContents {
131        entries: &'static [LootEntry],
132        component_type: Identifier,
133    },
134    /// Modify existing container contents.
135    ModifyContents {
136        modifier: &'static [ConditionalLootFunction],
137        component_type: Identifier,
138    },
139    /// Set container's loot table reference.
140    SetLootTable {
141        loot_table: Identifier,
142        seed: Option<i64>,
143    },
144    /// Set attribute modifiers on the item.
145    SetAttributes {
146        modifiers: &'static [AttributeModifier],
147        replace: bool,
148    },
149    /// Fill player head with texture from entity.
150    FillPlayerHead { entity: LootContextEntity },
151    /// Copy NBT/custom data from source.
152    CopyCustomData {
153        source: CopySource,
154        operations: &'static [CopyDataOperation],
155    },
156    /// Set banner pattern layers.
157    SetBannerPattern {
158        patterns: &'static [BannerPattern],
159        append: bool,
160    },
161    /// Set firework rocket properties.
162    SetFireworks {
163        explosions: Option<&'static [FireworkExplosion]>,
164        flight_duration: Option<i32>,
165    },
166    /// Set firework star explosion properties.
167    SetFireworkExplosion { explosion: FireworkExplosion },
168    /// Set book cover (title/author for written books).
169    SetBookCover {
170        title: Option<&'static str>,
171        author: Option<&'static str>,
172        generation: Option<i32>,
173    },
174    /// Set written book page contents.
175    SetWrittenBookPages {
176        pages: &'static [&'static str],
177        mode: ListOperation,
178    },
179    /// Set writable book page contents.
180    SetWritableBookPages {
181        pages: &'static [&'static str],
182        mode: ListOperation,
183    },
184    /// Toggle tooltip visibility.
185    ToggleTooltips {
186        toggles: &'static [(Identifier, bool)],
187    },
188    /// Discard/delete the item entirely.
189    Discard,
190    /// Reference to a named function in the registry.
191    Reference(Identifier),
192    /// Apply multiple functions in sequence.
193    Sequence {
194        functions: &'static [ConditionalLootFunction],
195    },
196    /// Apply `on_pass`/`on_fail` depending on whether the item matches `item_filter`.
197    Filtered {
198        item_filter: ItemPredicate,
199        on_pass: Option<&'static ConditionalLootFunction>,
200        on_fail: Option<&'static ConditionalLootFunction>,
201    },
202}
203
204/// Operation mode for list modifications (lore, book pages).
205#[derive(Debug, Clone, Copy)]
206pub enum ListOperation {
207    /// Replace all existing entries.
208    ReplaceAll,
209    /// Replace a section of entries.
210    ReplaceSection { offset: i32, size: Option<i32> },
211    /// Insert before existing entries.
212    InsertBefore { offset: i32 },
213    /// Insert after existing entries.
214    InsertAfter { offset: i32 },
215    /// Append to the end.
216    Append,
217}
218
219/// An attribute modifier for `SetAttributes` function.
220#[derive(Debug, Clone)]
221pub struct AttributeModifier {
222    pub attribute: Identifier,
223    pub operation: AttributeOperation,
224    pub amount: NumberProvider,
225    pub id: Identifier,
226    pub slot: EquipmentSlotGroup,
227}
228
229/// Attribute modifier operation type.
230#[expect(clippy::enum_variant_names, reason = "matches Vanilla naming")]
231#[derive(Debug, Clone, Copy)]
232pub enum AttributeOperation {
233    AddValue,
234    AddMultipliedBase,
235    AddMultipliedTotal,
236}
237
238/// Copy data operation for `CopyCustomData`.
239#[derive(Debug, Clone)]
240pub struct CopyDataOperation {
241    pub source_path: &'static str,
242    pub target_path: &'static str,
243    pub op: CopyDataOp,
244}
245
246/// Operation type for data copying.
247#[derive(Debug, Clone, Copy)]
248pub enum CopyDataOp {
249    Replace,
250    Append,
251    Merge,
252}
253
254/// A banner pattern layer.
255#[derive(Debug, Clone)]
256pub struct BannerPattern {
257    pub pattern: Identifier,
258    pub color: DyeColor,
259}
260
261/// A firework explosion definition.
262#[derive(Debug, Clone)]
263pub struct FireworkExplosion {
264    pub shape: FireworkShape,
265    pub colors: &'static [i32],
266    pub fade_colors: &'static [i32],
267    pub has_trail: bool,
268    pub has_twinkle: bool,
269}
270
271/// Firework explosion shape.
272#[derive(Debug, Clone, Copy)]
273pub enum FireworkShape {
274    SmallBall,
275    LargeBall,
276    Star,
277    Creeper,
278    Burst,
279}
280
281/// Formula types for `apply_bonus` function.
282#[derive(Debug, Clone, Copy)]
283pub enum BonusFormula {
284    /// Ore drops formula: count * (max(0, random(0..level+2) - 1) + 1)
285    OreDrops,
286    /// Uniform bonus: count + random(0..bonusMultiplier * level + 1)
287    UniformBonusCount { bonus_multiplier: i32 },
288    /// Binomial with bonus count: for each of (level + extra) trials, probability p to add 1
289    BinomialWithBonusCount { extra: i32, probability: f32 },
290}
291
292/// Source for copying components.
293#[derive(Debug, Clone, Copy)]
294pub enum CopySource {
295    BlockEntity,
296    This,
297    Attacker,
298    DirectAttacker,
299}
300
301/// Target for `set_name` function.
302#[derive(Debug, Clone, Copy)]
303pub enum NameTarget {
304    CustomName,
305    ItemName,
306}
307
308/// A stew effect for suspicious stew.
309#[derive(Debug, Clone)]
310pub struct StewEffect {
311    pub effect_type: Identifier,
312    pub duration: NumberProvider,
313}
314
315impl LootFunction {
316    /// Apply this function to modify the item stack.
317    ///
318    /// This modifies the item in place. Functions can change:
319    /// - Count (`SetCount`, `ExplosionDecay`, `ApplyBonus`, etc.)
320    /// - Damage/durability (`SetDamage`)
321    /// - Enchantments (`EnchantRandomly`, `EnchantWithLevels`, `SetEnchantments`)
322    /// - Components/NBT (`CopyComponents`, `SetComponents`, `CopyState`)
323    /// - Item type (`FurnaceSmelt`)
324    /// - And more...
325    pub fn apply<R: rand::Rng>(&self, item: &mut ItemStack, ctx: &mut LootContext<'_, R>) {
326        match self {
327            LootFunction::SetCount {
328                count: provider,
329                add,
330            } => {
331                let value = provider.get_int(ctx.rng);
332                if *add {
333                    item.count += value;
334                } else {
335                    item.count = value;
336                }
337            }
338            LootFunction::ExplosionDecay => {
339                if let Some(radius) = ctx.explosion_radius {
340                    // Each item has 1/radius chance to survive
341                    let probability = 1.0 / radius;
342                    let mut result_count = 0;
343                    for _ in 0..item.count {
344                        if ctx.rng.random::<f32>() <= probability {
345                            result_count += 1;
346                        }
347                    }
348                    item.count = result_count;
349                }
350            }
351            LootFunction::ApplyBonus {
352                enchantment,
353                formula,
354            } => {
355                let level = ctx.get_enchantment_level_by_id(enchantment);
356                item.count = formula.apply(item.count, level, ctx.rng);
357            }
358            LootFunction::EnchantedCountIncrease {
359                enchantment,
360                count: provider,
361                limit,
362            } => {
363                let level = ctx.get_attacking_entity_enchantment_level(enchantment);
364                if level > 0 {
365                    item.count += math_round(level as f32 * provider.get_float(ctx));
366                    if *limit > 0 {
367                        item.count = item.count.min(*limit);
368                    }
369                }
370            }
371            LootFunction::LimitCount { min, max } => {
372                if let Some(min_val) = min {
373                    item.count = item.count.max(*min_val);
374                }
375                if let Some(max_val) = max {
376                    item.count = item.count.min(*max_val);
377                }
378            }
379            LootFunction::SetDamage { damage, add } => {
380                item.set_damage_fraction(damage.get_simple(ctx.rng), *add);
381            }
382            LootFunction::EnchantRandomly { options } => {
383                // TODO: Implement when enchantment system is ready
384                item.enchant_randomly(options, ctx.rng);
385            }
386            LootFunction::EnchantWithLevels { levels, options } => {
387                // TODO: Implement when enchantment system is ready
388                let level = levels.get_int(ctx.rng);
389                item.enchant_with_levels(level, options, ctx.rng);
390            }
391            LootFunction::CopyComponents { source, include } => {
392                // TODO: Implement when block entity system is ready
393                item.copy_components(*source, include, ctx);
394            }
395            LootFunction::CopyState { block, properties } => {
396                // TODO: Implement block state copying
397                item.copy_block_state(block, properties, ctx);
398            }
399            LootFunction::SetComponents { components } => {
400                // TODO: Implement component setting from JSON
401                item.set_components_from_json(components);
402            }
403            LootFunction::SetCustomData { tag } => {
404                item.set_custom_data(&tag());
405            }
406            LootFunction::FurnaceSmelt { use_input_count } => {
407                item.apply_furnace_smelt(*use_input_count);
408            }
409            LootFunction::ExplorationMap {
410                destination,
411                decoration,
412                zoom,
413                skip_existing_chunks,
414            } => {
415                // TODO: Implement exploration map creation
416                item.create_exploration_map(destination, decoration, *zoom, *skip_existing_chunks);
417            }
418            LootFunction::SetName { name, target } => {
419                // TODO: Implement name setting
420                item.set_name(name, *target);
421            }
422            LootFunction::SetOminousBottleAmplifier { amplifier } => {
423                let amp = amplifier.get_int(ctx.rng).clamp(
424                    crate::data_components::OminousBottleAmplifier::MIN_AMPLIFIER,
425                    crate::data_components::OminousBottleAmplifier::MAX_AMPLIFIER,
426                );
427                item.set_ominous_bottle_amplifier(amp);
428            }
429            LootFunction::SetPotion { id } => {
430                item.set_potion(id);
431            }
432            LootFunction::SetStewEffect { effects } => {
433                item.set_stew_effects(effects, ctx.rng);
434            }
435            LootFunction::SetInstrument { options } => {
436                if let Some(instrument) = options.get_random(ctx.rng) {
437                    item.set(
438                        crate::data_components::vanilla_components::INSTRUMENT,
439                        crate::data_components::InstrumentComponent::new(
440                            crate::RegistryHolder::reference(instrument),
441                        ),
442                    );
443                }
444            }
445            LootFunction::SetEnchantments { enchantments, add } => {
446                let resolved: Vec<(Identifier, u32)> = enchantments
447                    .iter()
448                    .map(|(key, provider)| (key.clone(), provider.get_int(ctx.rng).max(0) as u32))
449                    .collect();
450                item.set_enchantments(&resolved, *add);
451            }
452            LootFunction::SetItem { item: new_item } => {
453                item.set_item(new_item);
454            }
455            LootFunction::CopyName { source } => {
456                item.copy_name(*source, ctx);
457            }
458            LootFunction::SetLore { lore, mode } => {
459                item.set_lore(lore, *mode);
460            }
461            LootFunction::SetContents {
462                entries,
463                component_type,
464            } => {
465                item.set_contents(entries, component_type, ctx);
466            }
467            LootFunction::ModifyContents {
468                modifier,
469                component_type,
470            } => {
471                item.modify_contents(modifier, component_type, ctx);
472            }
473            LootFunction::SetLootTable { loot_table, seed } => {
474                item.set_loot_table(loot_table, *seed);
475            }
476            LootFunction::SetAttributes { modifiers, replace } => {
477                item.set_attributes(modifiers, *replace, ctx.rng);
478            }
479            LootFunction::FillPlayerHead { entity } => {
480                item.fill_player_head(*entity, ctx);
481            }
482            LootFunction::CopyCustomData { source, operations } => {
483                item.copy_custom_data(*source, operations, ctx);
484            }
485            LootFunction::SetBannerPattern { patterns, append } => {
486                item.set_banner_pattern(patterns, *append);
487            }
488            LootFunction::SetFireworks {
489                explosions,
490                flight_duration,
491            } => {
492                item.set_fireworks(*explosions, *flight_duration);
493            }
494            LootFunction::SetFireworkExplosion { explosion } => {
495                item.set_firework_explosion(explosion);
496            }
497            LootFunction::SetBookCover {
498                title,
499                author,
500                generation,
501            } => {
502                item.set_book_cover(*title, *author, *generation);
503            }
504            LootFunction::SetWrittenBookPages { pages, mode } => {
505                item.set_written_book_pages(pages, *mode);
506            }
507            LootFunction::SetWritableBookPages { pages, mode } => {
508                item.set_writable_book_pages(pages, *mode);
509            }
510            LootFunction::ToggleTooltips { toggles } => {
511                item.toggle_tooltips(toggles);
512            }
513            LootFunction::Discard => {
514                item.count = 0;
515            }
516            LootFunction::Reference(_name) => {
517                // TODO: Implement function registry lookup
518            }
519            LootFunction::Sequence { functions } => {
520                for cond_func in *functions {
521                    if cond_func.conditions.iter().all(|c| c.test(ctx)) {
522                        cond_func.function.apply(item, ctx);
523                    }
524                }
525            }
526            LootFunction::Filtered {
527                item_filter,
528                on_pass,
529                on_fail,
530            } => {
531                let branch = if item_filter.test(item) {
532                    on_pass
533                } else {
534                    on_fail
535                };
536                if let Some(cond_func) = branch
537                    && cond_func.conditions.iter().all(|c| c.test(ctx))
538                {
539                    cond_func.function.apply(item, ctx);
540                }
541            }
542        }
543    }
544}
545
546impl BonusFormula {
547    /// Apply the bonus formula to calculate new count.
548    pub fn apply<R: rand::Rng>(&self, count: i32, level: i32, rng: &mut R) -> i32 {
549        match self {
550            BonusFormula::OreDrops => {
551                if level > 0 {
552                    // Vanilla: count * (max(0, random(0..level+2) - 1) + 1)
553                    let bonus = rng.random_range(0..level + 2) - 1;
554                    let multiplier = bonus.max(0) + 1;
555                    count * multiplier
556                } else {
557                    count
558                }
559            }
560            BonusFormula::UniformBonusCount { bonus_multiplier } => {
561                // Vanilla: count + random(0..bonusMultiplier * level + 1)
562                if level > 0 {
563                    count + rng.random_range(0..bonus_multiplier * level + 1)
564                } else {
565                    count
566                }
567            }
568            BonusFormula::BinomialWithBonusCount { extra, probability } => {
569                // Vanilla: for each of (level + extra) trials, probability p to add 1
570                let trials = level + extra;
571                let mut bonus = 0;
572                for _ in 0..trials {
573                    if rng.random::<f32>() < *probability {
574                        bonus += 1;
575                    }
576                }
577                count + bonus
578            }
579        }
580    }
581}