Skip to main content

steel_registry/loot_table/
functions.rs

1use super::{
2    DyeColor, EquipmentSlotGroup, Identifier, InstrumentRef, ItemStack, LootCondition, LootContext,
3    LootContextEntity, LootEntry, NumberProvider, REGISTRY, RngExt, TaggedRegistryExt,
4    ToolPredicate,
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    /// Conditionally apply function to specific item predicate matches.
197    Filtered {
198        item_filter: ToolPredicate,
199        modifier: &'static ConditionalLootFunction,
200    },
201}
202
203/// Operation mode for list modifications (lore, book pages).
204#[derive(Debug, Clone, Copy)]
205pub enum ListOperation {
206    /// Replace all existing entries.
207    ReplaceAll,
208    /// Replace a section of entries.
209    ReplaceSection { offset: i32, size: Option<i32> },
210    /// Insert before existing entries.
211    InsertBefore { offset: i32 },
212    /// Insert after existing entries.
213    InsertAfter { offset: i32 },
214    /// Append to the end.
215    Append,
216}
217
218/// An attribute modifier for `SetAttributes` function.
219#[derive(Debug, Clone)]
220pub struct AttributeModifier {
221    pub attribute: Identifier,
222    pub operation: AttributeOperation,
223    pub amount: NumberProvider,
224    pub id: Identifier,
225    pub slot: EquipmentSlotGroup,
226}
227
228/// Attribute modifier operation type.
229#[expect(clippy::enum_variant_names, reason = "matches Vanilla naming")]
230#[derive(Debug, Clone, Copy)]
231pub enum AttributeOperation {
232    AddValue,
233    AddMultipliedBase,
234    AddMultipliedTotal,
235}
236
237/// Copy data operation for `CopyCustomData`.
238#[derive(Debug, Clone)]
239pub struct CopyDataOperation {
240    pub source_path: &'static str,
241    pub target_path: &'static str,
242    pub op: CopyDataOp,
243}
244
245/// Operation type for data copying.
246#[derive(Debug, Clone, Copy)]
247pub enum CopyDataOp {
248    Replace,
249    Append,
250    Merge,
251}
252
253/// A banner pattern layer.
254#[derive(Debug, Clone)]
255pub struct BannerPattern {
256    pub pattern: Identifier,
257    pub color: DyeColor,
258}
259
260/// A firework explosion definition.
261#[derive(Debug, Clone)]
262pub struct FireworkExplosion {
263    pub shape: FireworkShape,
264    pub colors: &'static [i32],
265    pub fade_colors: &'static [i32],
266    pub has_trail: bool,
267    pub has_twinkle: bool,
268}
269
270/// Firework explosion shape.
271#[derive(Debug, Clone, Copy)]
272pub enum FireworkShape {
273    SmallBall,
274    LargeBall,
275    Star,
276    Creeper,
277    Burst,
278}
279
280/// Formula types for `apply_bonus` function.
281#[derive(Debug, Clone, Copy)]
282pub enum BonusFormula {
283    /// Ore drops formula: count * (max(0, random(0..level+2) - 1) + 1)
284    OreDrops,
285    /// Uniform bonus: count + random(0..bonusMultiplier * level + 1)
286    UniformBonusCount { bonus_multiplier: i32 },
287    /// Binomial with bonus count: for each of (level + extra) trials, probability p to add 1
288    BinomialWithBonusCount { extra: i32, probability: f32 },
289}
290
291/// Source for copying components.
292#[derive(Debug, Clone, Copy)]
293pub enum CopySource {
294    BlockEntity,
295    This,
296    Attacker,
297    DirectAttacker,
298}
299
300/// Target for `set_name` function.
301#[derive(Debug, Clone, Copy)]
302pub enum NameTarget {
303    CustomName,
304    ItemName,
305}
306
307/// A stew effect for suspicious stew.
308#[derive(Debug, Clone)]
309pub struct StewEffect {
310    pub effect_type: Identifier,
311    pub duration: NumberProvider,
312}
313
314impl LootFunction {
315    /// Apply this function to modify the item stack.
316    ///
317    /// This modifies the item in place. Functions can change:
318    /// - Count (`SetCount`, `ExplosionDecay`, `ApplyBonus`, etc.)
319    /// - Damage/durability (`SetDamage`)
320    /// - Enchantments (`EnchantRandomly`, `EnchantWithLevels`, `SetEnchantments`)
321    /// - Components/NBT (`CopyComponents`, `SetComponents`, `CopyState`)
322    /// - Item type (`FurnaceSmelt`)
323    /// - And more...
324    pub fn apply<R: rand::Rng>(&self, item: &mut ItemStack, ctx: &mut LootContext<'_, R>) {
325        match self {
326            LootFunction::SetCount {
327                count: provider,
328                add,
329            } => {
330                let value = provider.get_int(ctx.rng);
331                if *add {
332                    item.count += value;
333                } else {
334                    item.count = value;
335                }
336            }
337            LootFunction::ExplosionDecay => {
338                if let Some(radius) = ctx.explosion_radius {
339                    // Each item has 1/radius chance to survive
340                    let probability = 1.0 / radius;
341                    let mut result_count = 0;
342                    for _ in 0..item.count {
343                        if ctx.rng.random::<f32>() <= probability {
344                            result_count += 1;
345                        }
346                    }
347                    item.count = result_count;
348                }
349            }
350            LootFunction::ApplyBonus {
351                enchantment,
352                formula,
353            } => {
354                let level = ctx.get_enchantment_level_by_id(enchantment);
355                item.count = formula.apply(item.count, level, ctx.rng);
356            }
357            LootFunction::EnchantedCountIncrease {
358                enchantment,
359                count: provider,
360                limit,
361            } => {
362                let level = ctx.get_enchantment_level_by_id(enchantment);
363                if level > 0 {
364                    let bonus = (provider.get_simple(ctx.rng) * level as f32).round() as i32;
365                    let bonus = if *limit > 0 { bonus.min(*limit) } else { bonus };
366                    item.count += bonus;
367                }
368            }
369            LootFunction::LimitCount { min, max } => {
370                if let Some(min_val) = min {
371                    item.count = item.count.max(*min_val);
372                }
373                if let Some(max_val) = max {
374                    item.count = item.count.min(*max_val);
375                }
376            }
377            LootFunction::SetDamage { damage, add } => {
378                item.set_damage_fraction(damage.get_simple(ctx.rng), *add);
379            }
380            LootFunction::EnchantRandomly { options } => {
381                // TODO: Implement when enchantment system is ready
382                item.enchant_randomly(options, ctx.rng);
383            }
384            LootFunction::EnchantWithLevels { levels, options } => {
385                // TODO: Implement when enchantment system is ready
386                let level = levels.get_int(ctx.rng);
387                item.enchant_with_levels(level, options, ctx.rng);
388            }
389            LootFunction::CopyComponents { source, include } => {
390                // TODO: Implement when block entity system is ready
391                item.copy_components(*source, include, ctx);
392            }
393            LootFunction::CopyState { block, properties } => {
394                // TODO: Implement block state copying
395                item.copy_block_state(block, properties, ctx);
396            }
397            LootFunction::SetComponents { components } => {
398                // TODO: Implement component setting from JSON
399                item.set_components_from_json(components);
400            }
401            LootFunction::SetCustomData { tag } => {
402                item.set_custom_data(&tag());
403            }
404            LootFunction::FurnaceSmelt { use_input_count } => {
405                item.apply_furnace_smelt(*use_input_count);
406            }
407            LootFunction::ExplorationMap {
408                destination,
409                decoration,
410                zoom,
411                skip_existing_chunks,
412            } => {
413                // TODO: Implement exploration map creation
414                item.create_exploration_map(destination, decoration, *zoom, *skip_existing_chunks);
415            }
416            LootFunction::SetName { name, target } => {
417                // TODO: Implement name setting
418                item.set_name(name, *target);
419            }
420            LootFunction::SetOminousBottleAmplifier { amplifier } => {
421                let amp = amplifier.get_int(ctx.rng).clamp(
422                    crate::data_components::OminousBottleAmplifier::MIN_AMPLIFIER,
423                    crate::data_components::OminousBottleAmplifier::MAX_AMPLIFIER,
424                );
425                item.set_ominous_bottle_amplifier(amp);
426            }
427            LootFunction::SetPotion { id } => {
428                item.set_potion(id);
429            }
430            LootFunction::SetStewEffect { effects } => {
431                item.set_stew_effects(effects, ctx.rng);
432            }
433            LootFunction::SetInstrument { options } => {
434                if let Some(instrument) = options.get_random(ctx.rng) {
435                    item.set(
436                        crate::data_components::vanilla_components::INSTRUMENT,
437                        crate::data_components::InstrumentComponent::new(
438                            crate::RegistryHolder::reference(instrument),
439                        ),
440                    );
441                }
442            }
443            LootFunction::SetEnchantments { enchantments, add } => {
444                let resolved: Vec<(Identifier, u32)> = enchantments
445                    .iter()
446                    .map(|(key, provider)| (key.clone(), provider.get_int(ctx.rng).max(0) as u32))
447                    .collect();
448                item.set_enchantments(&resolved, *add);
449            }
450            LootFunction::SetItem { item: new_item } => {
451                item.set_item(new_item);
452            }
453            LootFunction::CopyName { source } => {
454                item.copy_name(*source, ctx);
455            }
456            LootFunction::SetLore { lore, mode } => {
457                item.set_lore(lore, *mode);
458            }
459            LootFunction::SetContents {
460                entries,
461                component_type,
462            } => {
463                item.set_contents(entries, component_type, ctx);
464            }
465            LootFunction::ModifyContents {
466                modifier,
467                component_type,
468            } => {
469                item.modify_contents(modifier, component_type, ctx);
470            }
471            LootFunction::SetLootTable { loot_table, seed } => {
472                item.set_loot_table(loot_table, *seed);
473            }
474            LootFunction::SetAttributes { modifiers, replace } => {
475                item.set_attributes(modifiers, *replace, ctx.rng);
476            }
477            LootFunction::FillPlayerHead { entity } => {
478                item.fill_player_head(*entity, ctx);
479            }
480            LootFunction::CopyCustomData { source, operations } => {
481                item.copy_custom_data(*source, operations, ctx);
482            }
483            LootFunction::SetBannerPattern { patterns, append } => {
484                item.set_banner_pattern(patterns, *append);
485            }
486            LootFunction::SetFireworks {
487                explosions,
488                flight_duration,
489            } => {
490                item.set_fireworks(*explosions, *flight_duration);
491            }
492            LootFunction::SetFireworkExplosion { explosion } => {
493                item.set_firework_explosion(explosion);
494            }
495            LootFunction::SetBookCover {
496                title,
497                author,
498                generation,
499            } => {
500                item.set_book_cover(*title, *author, *generation);
501            }
502            LootFunction::SetWrittenBookPages { pages, mode } => {
503                item.set_written_book_pages(pages, *mode);
504            }
505            LootFunction::SetWritableBookPages { pages, mode } => {
506                item.set_writable_book_pages(pages, *mode);
507            }
508            LootFunction::ToggleTooltips { toggles } => {
509                item.toggle_tooltips(toggles);
510            }
511            LootFunction::Discard => {
512                item.count = 0;
513            }
514            LootFunction::Reference(_name) => {
515                // TODO: Implement function registry lookup
516            }
517            LootFunction::Sequence { functions } => {
518                for cond_func in *functions {
519                    if cond_func.conditions.iter().all(|c| c.test(ctx)) {
520                        cond_func.function.apply(item, ctx);
521                    }
522                }
523            }
524            LootFunction::Filtered {
525                item_filter,
526                modifier,
527            } => {
528                if item_filter.test(item, ctx) && modifier.conditions.iter().all(|c| c.test(ctx)) {
529                    modifier.function.apply(item, ctx);
530                }
531            }
532        }
533    }
534}
535
536impl BonusFormula {
537    /// Apply the bonus formula to calculate new count.
538    pub fn apply<R: rand::Rng>(&self, count: i32, level: i32, rng: &mut R) -> i32 {
539        match self {
540            BonusFormula::OreDrops => {
541                if level > 0 {
542                    // Vanilla: count * (max(0, random(0..level+2) - 1) + 1)
543                    let bonus = rng.random_range(0..level + 2) - 1;
544                    let multiplier = bonus.max(0) + 1;
545                    count * multiplier
546                } else {
547                    count
548                }
549            }
550            BonusFormula::UniformBonusCount { bonus_multiplier } => {
551                // Vanilla: count + random(0..bonusMultiplier * level + 1)
552                if level > 0 {
553                    count + rng.random_range(0..bonus_multiplier * level + 1)
554                } else {
555                    count
556                }
557            }
558            BonusFormula::BinomialWithBonusCount { extra, probability } => {
559                // Vanilla: for each of (level + extra) trials, probability p to add 1
560                let trials = level + extra;
561                let mut bonus = 0;
562                for _ in 0..trials {
563                    if rng.random::<f32>() < *probability {
564                        bonus += 1;
565                    }
566                }
567                count + bonus
568            }
569        }
570    }
571}