wowlab_loadout/

enrich.rs

// #t(file: rust_alloc_in_loop) enrichment builds per-node output strings during tree walk

use serde::Serialize;
use wowlab_types::{
    data::{
        TraitNode, TraitNodeEntry, TraitNodeType, TraitSelection, TraitTreeIndex,
        TraitTreeWithSelections,
    },
    sim::{FastBuildHasher, FastMap, FastSet},
};

/// Enriched talent loadout with resolved spell info and point summary.
#[derive(Clone, Debug, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct EnrichedLoadout {
    pub spec_id: i32,
    pub spec_name: String,
    pub class_name: String,
    pub talents: Vec<EnrichedTalent>,
    pub summary: TalentPointSummary,
}

/// A single resolved talent with spell info and tree classification.
#[derive(Clone, Debug, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct EnrichedTalent {
    pub node_id: i32,
    pub trait_node_entry_id: i32,
    pub spell_id: i32,
    /// `TraitDefinition.VisibleSpellID`: the spell `spell_id` is displayed and activated as in-game.
    pub override_spell_id: i32,
    /// `TraitDefinition.OverridesSpellID`: the learned talent replaces this spell on the action bar.
    pub replaces_spell_id: i32,
    pub name: String,
    pub ranks: i32,
    pub max_ranks: i32,
    pub tree_index: i32,
    pub sub_tree_id: i32,
    pub tree_section: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub choice_index: Option<u8>,
    pub effect_overrides: Vec<EnrichedTalentEffectOverride>,
}

/// Rank-specific adjustment for one effect on a selected talent spell.
#[derive(Clone, Debug, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
// #t(rust_similar_structs) enriched serialized output is separate from the engine-port selection contract
pub struct EnrichedTalentEffectOverride {
    pub effect_index: i32,
    pub operation: wowlab_types::data::TraitEffectOperation,
    pub value: f64,
}

/// Point totals broken down by tree section.
#[derive(Clone, Debug, Serialize)]
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
pub struct TalentPointSummary {
    pub total_points: i32,
    pub class_points: i32,
    pub spec_points: i32,
    pub hero_points: i32,
}

/// Enriches a decoded loadout with spell metadata and point totals.
#[must_use]
pub fn enrich_loadout(data: &TraitTreeWithSelections) -> EnrichedLoadout {
    let mut node_by_id =
        FastMap::with_capacity_and_hasher(data.tree.nodes.len(), FastBuildHasher::default());

    for node in &data.tree.nodes {
        node_by_id.entry(node.id).or_insert(node);
    }

    let sub_tree_names: FastMap<i32, &str> = data
        .tree
        .sub_trees
        .iter()
        .map(|st| (st.id, st.name.as_str()))
        .collect();

    let active_sub_trees = active_sub_tree_ids(data, &node_by_id);

    let mut talents = Vec::with_capacity(data.selections.len());
    let mut class_points = 0i32;
    let mut spec_points = 0i32;
    let mut hero_points = 0i32;

    for sel in &data.selections {
        if !sel.selected || sel.ranks_purchased <= 0 {
            continue;
        }

        let Some(node) = node_by_id.get(&sel.node_id) else {
            continue;
        };

        if node.sub_tree_id > 0
            && active_sub_trees
                .as_ref()
                .is_some_and(|ids| !ids.contains(&node.sub_tree_id))
        {
            continue;
        }

        let allocations = allocate_entry_ranks(node, sel);

        if allocations.is_empty() {
            continue;
        }

        let tree_section = match TraitTreeIndex::from_dbc(node.tree_index) {
            Some(TraitTreeIndex::Class) => {
                class_points += sel.ranks_purchased;

                "class".to_string()
            }
            Some(TraitTreeIndex::Specialization) => {
                spec_points += sel.ranks_purchased;

                "spec".to_string()
            }
            Some(TraitTreeIndex::Selection) => "selection".to_string(),
            _ => {
                hero_points += sel.ranks_purchased;

                sub_tree_names
                    .get(&node.sub_tree_id)
                    .map_or_else(|| "hero".to_string(), |s| (*s).to_string())
            }
        };

        for allocation in allocations {
            let EntryAllocation {
                entry,
                ranks,
                max_ranks,
            } = allocation;

            talents.push(EnrichedTalent {
                node_id: sel.node_id,
                trait_node_entry_id: entry.id,
                spell_id: entry.spell_id,
                override_spell_id: entry.override_spell_id,
                replaces_spell_id: entry.replaces_spell_id,
                // #t(block: rust_clone_in_loop) building owned structs from borrowed tree entries requires cloning
                name: entry.name.clone(),
                ranks,
                max_ranks,
                tree_index: node.tree_index,
                sub_tree_id: if TraitTreeIndex::from_dbc(node.tree_index)
                    == Some(TraitTreeIndex::Selection)
                {
                    entry.sub_tree_id
                } else {
                    node.sub_tree_id
                },
                tree_section: tree_section.clone(),
                choice_index: sel.choice_index,
                effect_overrides: effect_overrides(entry, ranks),
            });
        }
    }

    EnrichedLoadout {
        spec_id: data.tree.spec_id,
        spec_name: data.tree.spec_name.clone(),
        class_name: data.tree.class_name.clone(),
        talents,
        summary: TalentPointSummary {
            total_points: class_points + spec_points + hero_points,
            class_points,
            spec_points,
            hero_points,
        },
    }
}

fn effect_overrides(entry: &TraitNodeEntry, ranks: i32) -> Vec<EnrichedTalentEffectOverride> {
    let rank = f64::from(ranks);

    entry
        .effect_points
        .iter()
        .filter_map(|effect| {
            effect
                .rank_values
                .iter()
                .find(|(point_rank, _)| (*point_rank - rank).abs() < f64::EPSILON)
                .map(|&(_, value)| EnrichedTalentEffectOverride {
                    effect_index: effect.effect_index,
                    operation: effect.operation,
                    value,
                })
        })
        .collect()
}

fn active_sub_tree_ids(
    data: &TraitTreeWithSelections,
    node_by_id: &FastMap<i32, &TraitNode>,
) -> Option<FastSet<i32>> {
    let mut active =
        FastSet::with_capacity_and_hasher(data.tree.sub_trees.len(), FastBuildHasher::default());
    let mut has_selected_picker = false;

    for sel in &data.selections {
        if !sel.selected {
            continue;
        }

        let Some(node) = node_by_id.get(&sel.node_id) else {
            continue;
        };

        if TraitTreeIndex::from_dbc(node.tree_index) != Some(TraitTreeIndex::Selection) {
            continue;
        }

        has_selected_picker = true;

        let entry_idx = sel.choice_index.unwrap_or(0) as usize;
        let Some(entry) = node.entries.get(entry_idx) else {
            continue;
        };

        if entry.sub_tree_id > 0
            && data
                .tree
                .sub_trees
                .iter()
                .any(|sub_tree| sub_tree.id == entry.sub_tree_id)
        {
            active.insert(entry.sub_tree_id);
        }
    }

    has_selected_picker.then_some(active)
}

struct EntryAllocation<'a> {
    entry: &'a TraitNodeEntry,
    ranks: i32,
    max_ranks: i32,
}

fn allocate_entry_ranks<'a>(node: &'a TraitNode, sel: &TraitSelection) -> Vec<EntryAllocation<'a>> {
    if TraitNodeType::from_dbc(node.node_type) == Some(TraitNodeType::Tiered)
        && node.entries.iter().any(|e| e.max_ranks > 0)
    {
        let mut allocations = Vec::with_capacity(node.entries.len());
        let mut remaining = sel.ranks_purchased;

        for entry in &node.entries {
            if remaining <= 0 {
                break;
            }

            let allocated = remaining.min(entry.max_ranks);

            if allocated > 0 {
                allocations.push(EntryAllocation {
                    entry,
                    ranks: allocated,
                    max_ranks: entry.max_ranks,
                });
                remaining -= allocated;
            }
        }

        return allocations;
    }

    let entry_idx = sel.choice_index.unwrap_or(0) as usize;
    let entry = node.entries.get(entry_idx).or_else(|| node.entries.first());

    entry
        .map(|entry| EntryAllocation {
            entry,
            ranks: sel.ranks_purchased,
            max_ranks: node.max_ranks,
        })
        .into_iter()
        .collect()
}

impl EnrichedLoadout {
    /// Returns selected spell IDs and ranks.
    // docref:start talent-trees-get-spell-ids
    #[must_use]
    pub fn talent_selections(&self) -> Vec<(u32, u8)> {
        self.talents
            .iter()
            .filter(|t| t.spell_id > 0)
            .filter_map(|t| Some((u32::try_from(t.spell_id).ok()?, u8::try_from(t.ranks).ok()?)))
            .collect()
    }
    // docref:end talent-trees-get-spell-ids
}

#[cfg(test)]
#[path = "enrich_tests.rs"]
mod tests;