Skip to main content

wowlab_loadout/
enrich.rs

1// #t(file: rust_alloc_in_loop) enrichment builds per-node output strings during tree walk
2
3use serde::Serialize;
4use wowlab_types::{
5    data::{
6        TraitNode, TraitNodeEntry, TraitNodeType, TraitSelection, TraitTreeIndex,
7        TraitTreeWithSelections,
8    },
9    sim::{FastBuildHasher, FastMap, FastSet},
10};
11
12/// Enriched talent loadout with resolved spell info and point summary.
13#[derive(Clone, Debug, Serialize)]
14pub struct EnrichedLoadout {
15    pub spec_id: i32,
16    pub spec_name: String,
17    pub class_name: String,
18    pub talents: Vec<EnrichedTalent>,
19    pub summary: TalentPointSummary,
20}
21
22/// A single resolved talent with spell info and tree classification.
23#[derive(Clone, Debug, Serialize)]
24pub struct EnrichedTalent {
25    pub node_id: i32,
26    pub trait_node_entry_id: i32,
27    pub spell_id: i32,
28    /// `TraitDefinition.VisibleSpellID`: the spell `spell_id` is displayed and activated as in-game.
29    pub override_spell_id: i32,
30    /// `TraitDefinition.OverridesSpellID`: the learned talent replaces this spell on the action bar.
31    pub replaces_spell_id: i32,
32    pub name: String,
33    pub ranks: i32,
34    pub max_ranks: i32,
35    pub tree_index: i32,
36    pub sub_tree_id: i32,
37    pub tree_section: String,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub choice_index: Option<u8>,
40    pub effect_overrides: Vec<EnrichedTalentEffectOverride>,
41}
42
43/// Rank-specific adjustment for one effect on a selected talent spell.
44#[derive(Clone, Debug, Serialize)]
45// #t(rust_similar_structs) enriched serialized output is separate from the engine-port selection contract
46pub struct EnrichedTalentEffectOverride {
47    pub effect_index: i32,
48    pub operation: wowlab_types::data::TraitEffectOperation,
49    pub value: f64,
50}
51
52/// Point totals broken down by tree section.
53#[derive(Clone, Debug, Serialize)]
54pub struct TalentPointSummary {
55    pub total_points: i32,
56    pub class_points: i32,
57    pub spec_points: i32,
58    pub hero_points: i32,
59}
60
61/// Enriches a decoded loadout with spell metadata and point totals.
62#[must_use]
63pub fn enrich_loadout(data: &TraitTreeWithSelections) -> EnrichedLoadout {
64    let mut node_by_id =
65        FastMap::with_capacity_and_hasher(data.tree.nodes.len(), FastBuildHasher::default());
66
67    for node in &data.tree.nodes {
68        node_by_id.entry(node.id).or_insert(node);
69    }
70
71    let sub_tree_names: FastMap<i32, &str> = data
72        .tree
73        .sub_trees
74        .iter()
75        .map(|st| (st.id, st.name.as_str()))
76        .collect();
77
78    let active_sub_trees = active_sub_tree_ids(data, &node_by_id);
79
80    let mut talents = Vec::with_capacity(data.selections.len());
81    let mut class_points = 0i32;
82    let mut spec_points = 0i32;
83    let mut hero_points = 0i32;
84
85    for sel in &data.selections {
86        if !sel.selected || sel.ranks_purchased <= 0 {
87            continue;
88        }
89
90        let Some(node) = node_by_id.get(&sel.node_id) else {
91            continue;
92        };
93
94        if node.sub_tree_id > 0
95            && active_sub_trees
96                .as_ref()
97                .is_some_and(|ids| !ids.contains(&node.sub_tree_id))
98        {
99            continue;
100        }
101
102        let allocations = allocate_entry_ranks(node, sel);
103
104        if allocations.is_empty() {
105            continue;
106        }
107
108        let tree_section = match TraitTreeIndex::from_dbc(node.tree_index) {
109            Some(TraitTreeIndex::Class) => {
110                class_points += sel.ranks_purchased;
111
112                "class".to_string()
113            }
114            Some(TraitTreeIndex::Specialization) => {
115                spec_points += sel.ranks_purchased;
116
117                "spec".to_string()
118            }
119            Some(TraitTreeIndex::Selection) => "selection".to_string(),
120            _ => {
121                hero_points += sel.ranks_purchased;
122
123                sub_tree_names
124                    .get(&node.sub_tree_id)
125                    .map_or_else(|| "hero".to_string(), |s| (*s).to_string())
126            }
127        };
128
129        for allocation in allocations {
130            let EntryAllocation {
131                entry,
132                ranks,
133                max_ranks,
134            } = allocation;
135
136            talents.push(EnrichedTalent {
137                node_id: sel.node_id,
138                trait_node_entry_id: entry.id,
139                spell_id: entry.spell_id,
140                override_spell_id: entry.override_spell_id,
141                replaces_spell_id: entry.replaces_spell_id,
142                // #t(block: rust_clone_in_loop) building owned structs from borrowed tree entries requires cloning
143                name: entry.name.clone(),
144                ranks,
145                max_ranks,
146                tree_index: node.tree_index,
147                sub_tree_id: if TraitTreeIndex::from_dbc(node.tree_index)
148                    == Some(TraitTreeIndex::Selection)
149                {
150                    entry.sub_tree_id
151                } else {
152                    node.sub_tree_id
153                },
154                tree_section: tree_section.clone(),
155                choice_index: sel.choice_index,
156                effect_overrides: effect_overrides(entry, ranks),
157            });
158        }
159    }
160
161    EnrichedLoadout {
162        spec_id: data.tree.spec_id,
163        spec_name: data.tree.spec_name.clone(),
164        class_name: data.tree.class_name.clone(),
165        talents,
166        summary: TalentPointSummary {
167            total_points: class_points + spec_points + hero_points,
168            class_points,
169            spec_points,
170            hero_points,
171        },
172    }
173}
174
175fn effect_overrides(entry: &TraitNodeEntry, ranks: i32) -> Vec<EnrichedTalentEffectOverride> {
176    let rank = f64::from(ranks);
177
178    entry
179        .effect_points
180        .iter()
181        .filter_map(|effect| {
182            effect
183                .rank_values
184                .iter()
185                .find(|(point_rank, _)| (*point_rank - rank).abs() < f64::EPSILON)
186                .map(|&(_, value)| EnrichedTalentEffectOverride {
187                    effect_index: effect.effect_index,
188                    operation: effect.operation,
189                    value,
190                })
191        })
192        .collect()
193}
194
195fn active_sub_tree_ids(
196    data: &TraitTreeWithSelections,
197    node_by_id: &FastMap<i32, &TraitNode>,
198) -> Option<FastSet<i32>> {
199    let mut active =
200        FastSet::with_capacity_and_hasher(data.tree.sub_trees.len(), FastBuildHasher::default());
201    let mut has_selected_picker = false;
202
203    for sel in &data.selections {
204        if !sel.selected {
205            continue;
206        }
207
208        let Some(node) = node_by_id.get(&sel.node_id) else {
209            continue;
210        };
211
212        if TraitTreeIndex::from_dbc(node.tree_index) != Some(TraitTreeIndex::Selection) {
213            continue;
214        }
215
216        has_selected_picker = true;
217
218        let entry_idx = sel.choice_index.unwrap_or(0) as usize;
219        let Some(entry) = node.entries.get(entry_idx) else {
220            continue;
221        };
222
223        if entry.sub_tree_id > 0
224            && data
225                .tree
226                .sub_trees
227                .iter()
228                .any(|sub_tree| sub_tree.id == entry.sub_tree_id)
229        {
230            active.insert(entry.sub_tree_id);
231        }
232    }
233
234    has_selected_picker.then_some(active)
235}
236
237struct EntryAllocation<'a> {
238    entry: &'a TraitNodeEntry,
239    ranks: i32,
240    max_ranks: i32,
241}
242
243fn allocate_entry_ranks<'a>(node: &'a TraitNode, sel: &TraitSelection) -> Vec<EntryAllocation<'a>> {
244    if TraitNodeType::from_dbc(node.node_type) == Some(TraitNodeType::Tiered)
245        && node.entries.iter().any(|e| e.max_ranks > 0)
246    {
247        let mut allocations = Vec::with_capacity(node.entries.len());
248        let mut remaining = sel.ranks_purchased;
249
250        for entry in &node.entries {
251            if remaining <= 0 {
252                break;
253            }
254
255            let allocated = remaining.min(entry.max_ranks);
256
257            if allocated > 0 {
258                allocations.push(EntryAllocation {
259                    entry,
260                    ranks: allocated,
261                    max_ranks: entry.max_ranks,
262                });
263                remaining -= allocated;
264            }
265        }
266
267        return allocations;
268    }
269
270    let entry_idx = sel.choice_index.unwrap_or(0) as usize;
271    let entry = node.entries.get(entry_idx).or_else(|| node.entries.first());
272
273    entry
274        .map(|entry| EntryAllocation {
275            entry,
276            ranks: sel.ranks_purchased,
277            max_ranks: node.max_ranks,
278        })
279        .into_iter()
280        .collect()
281}
282
283impl EnrichedLoadout {
284    /// Returns selected spell IDs and ranks.
285    // docref:start talent-trees-get-spell-ids
286    #[must_use]
287    pub fn talent_selections(&self) -> Vec<(u32, u8)> {
288        self.talents
289            .iter()
290            .filter(|t| t.spell_id > 0)
291            .filter_map(|t| Some((u32::try_from(t.spell_id).ok()?, u8::try_from(t.ranks).ok()?)))
292            .collect()
293    }
294    // docref:end talent-trees-get-spell-ids
295}
296
297#[cfg(test)]
298#[path = "enrich_tests.rs"]
299mod tests;