wowlab_parsers/parsers/transform/

trait.rs

use wowlab_engine_domain::dbc::{
    TraitConditionType, TraitCurrencyFlags, TraitEffectOperation, TraitNodeType, TraitTreeIndex,
};
use wowlab_types::{
    data::{
        PointLimits, TraitEdge, TraitEffectPoints, TraitNode, TraitNodeEntry, TraitSubTree,
        TraitTreeFlat,
    },
    sim::{FastSet, IntMap},
};

#[cfg(feature = "dbc")]
use super::super::dbc::rows::{
    TraitDefinitionRow, TraitEdgeRow, TraitNodeRow, TraitNodeXTraitNodeEntryRow,
    TraitTreeLoadoutRow,
};
#[cfg(feature = "dbc")]
use super::{super::dbc::DbcData, super::errors::TransformError};

const TARGET_HERO_X: i32 = 7500;
const TARGET_HERO_Y: i32 = 1200;

/// Transform a specialization into its complete talent tree structure.
///
/// # Errors
///
/// Returns [`TransformError`] when the specialization or its required talent-tree data is absent.
pub fn transform_trait_tree(dbc: &DbcData, spec_id: i32) -> Result<TraitTreeFlat, TransformError> {
    let spec = dbc
        .chr_specialization
        .get(&spec_id)
        .ok_or_else(|| TransformError::spec_not_found(spec_id))?;

    let class_id = spec
        .ClassID
        .ok_or_else(|| TransformError::spec_not_found(spec_id))?;
    let chr_class = dbc
        .chr_classes
        .get(&class_id)
        .ok_or_else(|| TransformError::class_not_found(class_id))?;

    let all_loadouts = dbc
        .trait_tree_loadout_by_spec
        .get(&spec_id)
        .ok_or_else(|| TransformError::loadout_not_found(spec_id))?;

    let all_tree_ids = select_live_tree_ids(dbc, all_loadouts);

    // IntMap iteration order is nondeterministic.
    let tree_id = all_tree_ids.iter().copied().min().unwrap_or(0);

    let loadouts: Vec<_> = all_loadouts
        .iter()
        .filter(|l| all_tree_ids.contains(&l.TraitTreeID))
        .cloned()
        .collect();

    let order_index_map = build_order_index_map(dbc, &loadouts);

    let all_tree_nodes: Vec<_> = all_tree_ids
        .iter()
        .flat_map(|&tid| {
            dbc.trait_node_by_tree
                .get(&tid)
                .cloned()
                .unwrap_or_default()
        })
        .collect();

    let spec_set_to_specs = build_spec_set_lookup(dbc);

    let tree_nodes: Vec<_> = all_tree_nodes
        .iter()
        .filter(|node| {
            node.TraitSubTreeID > 0
                || TraitNodeType::from_dbc(node.Type) == Some(TraitNodeType::SubtreeSelection)
                || should_include_node_for_spec(dbc, node.ID, spec_id, &spec_set_to_specs)
        })
        .cloned()
        .collect();

    let node_ids: Vec<i32> = tree_nodes.iter().map(|n| n.ID).collect();
    let node_id_set: FastSet<i32> = node_ids.iter().copied().collect();

    let valid_sub_tree_ids =
        get_valid_subtrees_for_spec(dbc, &all_tree_nodes, spec_id, &spec_set_to_specs);

    // #t(block: rust_unnecessary_collect) collect into FastSet intentionally deduplicates sub-tree IDs
    let all_sub_tree_ids: Vec<i32> = tree_nodes
        .iter()
        .map(|n| n.TraitSubTreeID)
        .filter(|&id| id > 0)
        .collect::<FastSet<_>>()
        .into_iter()
        .collect();

    let all_edges = collect_edges(dbc, &node_ids, &node_id_set);

    let indexes = TraitNodeIndexes {
        node_entries: build_node_entries_map(dbc, &tree_nodes),
        hero_tree_offsets: calculate_hero_tree_offsets(&tree_nodes, &all_sub_tree_ids),
        group_to_tree_index: build_group_tree_index_map(dbc, tree_id),
        order_index: order_index_map,
        spec_set_to_specs,
    };

    let mut nodes: Vec<TraitNode> = tree_nodes
        .iter()
        .filter_map(|node| build_talent_node(dbc, node, &indexes, spec_id, None))
        .collect();

    nodes.sort_by_key(|n| n.id);

    let mut edges: Vec<TraitEdge> = all_edges
        .iter()
        .map(|e| TraitEdge {
            id: e.ID,
            from_node_id: e.LeftTraitNodeID,
            to_node_id: e.RightTraitNodeID,
            visual_style: e.VisualStyle,
        })
        .collect();

    edges.sort_by_key(|e| e.id);

    let mut sub_trees = build_subtrees(dbc, &valid_sub_tree_ids, &nodes);

    sub_trees.sort_by_key(|s| s.id);

    let point_limits = calculate_point_limits(dbc, &all_tree_ids);

    let mut all_node_ids: Vec<i32> = all_tree_nodes
        .iter()
        .filter(|n| n.TraitTreeID == tree_id)
        .map(|n| n.ID)
        .collect();

    all_node_ids.sort_unstable();

    Ok(TraitTreeFlat {
        spec_id,
        spec_name: spec.Name_lang.clone().unwrap_or_default(),
        class_name: chr_class.Name_lang.clone().unwrap_or_default(),
        tree_id,
        all_node_ids,
        nodes,
        edges,
        sub_trees,
        point_limits,
    })
}

fn select_live_tree_ids(dbc: &DbcData, loadouts: &[TraitTreeLoadoutRow]) -> Vec<i32> {
    // #t(block: rust_unnecessary_collect) collect into FastSet intentionally deduplicates tree IDs
    let referenced: Vec<i32> = loadouts
        .iter()
        .map(|l| l.TraitTreeID)
        .collect::<FastSet<_>>()
        .into_iter()
        .collect();

    let live: Vec<i32> = referenced
        .iter()
        .copied()
        .filter(|tid| dbc.skill_line_x_trait_tree.contains_key(tid))
        .collect();

    if live.is_empty() { referenced } else { live }
}

fn build_order_index_map(dbc: &DbcData, loadouts: &[TraitTreeLoadoutRow]) -> IntMap<i32, i32> {
    let mut entries: Vec<_> = loadouts
        .iter()
        .flat_map(|loadout| {
            dbc.trait_tree_loadout_entry
                .get(&loadout.ID)
                .cloned()
                .unwrap_or_default()
        })
        .collect();

    entries.sort_by_key(|e| e.OrderIndex);

    entries
        .iter()
        .map(|e| (e.SelectedTraitNodeID, e.OrderIndex))
        .collect()
}

pub(super) fn collect_edges(
    dbc: &DbcData,
    node_ids: &[i32],
    node_id_set: &FastSet<i32>,
) -> Vec<TraitEdgeRow> {
    node_ids
        .iter()
        .flat_map(|&node_id| dbc.trait_edge.get(&node_id).cloned().unwrap_or_default())
        .filter(|e| {
            node_id_set.contains(&e.LeftTraitNodeID) && node_id_set.contains(&e.RightTraitNodeID)
        })
        .collect()
}

fn build_group_tree_index_map(dbc: &DbcData, primary_tree_id: i32) -> IntMap<i32, i32> {
    let mut map = IntMap::default();

    let mut currency_to_tree_index: IntMap<i32, i32> = IntMap::default();

    if let Some(tree_currencies) = dbc.trait_tree_x_trait_currency.get(&primary_tree_id) {
        for tc in tree_currencies {
            if let Some(currency) = dbc.trait_currency.get(&tc.TraitCurrencyID) {
                let flags = TraitCurrencyFlags::from_dbc(currency.Flags);

                if flags.contains(TraitCurrencyFlags::CLASS) {
                    currency_to_tree_index.insert(tc.TraitCurrencyID, TraitTreeIndex::Class as i32);
                } else if flags.contains(TraitCurrencyFlags::SPECIALIZATION) {
                    currency_to_tree_index
                        .insert(tc.TraitCurrencyID, TraitTreeIndex::Specialization as i32);
                }
            }
        }
    }

    for (&group_id, cost_links) in &dbc.trait_node_group_x_trait_cost {
        for cost_link in cost_links {
            if let Some(cost) = dbc.trait_cost.get(&cost_link.TraitCostID) {
                if let Some(&tree_index) = currency_to_tree_index.get(&cost.TraitCurrencyID) {
                    map.insert(group_id, tree_index);
                }
            }
        }
    }

    map
}

pub(super) fn build_node_entries_map(
    dbc: &DbcData,
    tree_nodes: &[TraitNodeRow],
) -> IntMap<i32, Vec<TraitNodeXTraitNodeEntryRow>> {
    tree_nodes
        .iter()
        .map(|n| {
            let entries = dbc
                .trait_node_x_trait_node_entry
                .get(&n.ID)
                .cloned()
                .unwrap_or_default();

            let mut sorted = entries;
            // DBC may assign the same selection index to multiple entries.

            sorted.sort_by_key(|entry| {
                let selection_index = if entry.Index >= 0 {
                    entry.Index
                } else {
                    -entry.TraitNodeEntryID
                };

                (selection_index, entry.TraitNodeEntryID)
            });

            (n.ID, sorted)
        })
        .collect()
}

pub(super) struct TraitNodeIndexes {
    pub(super) node_entries: IntMap<i32, Vec<TraitNodeXTraitNodeEntryRow>>,
    pub(super) hero_tree_offsets: IntMap<i32, (i32, i32)>,
    pub(super) group_to_tree_index: IntMap<i32, i32>,
    pub(super) order_index: IntMap<i32, i32>,
    pub(super) spec_set_to_specs: IntMap<i32, Vec<i32>>,
}

pub(super) fn build_talent_node(
    dbc: &DbcData,
    node: &TraitNodeRow,
    indexes: &TraitNodeIndexes,
    spec_id: i32,
    tree_index_override: Option<i32>,
) -> Option<TraitNode> {
    let x_entries = indexes.node_entries.get(&node.ID)?;

    let mut entries: Vec<TraitNodeEntry> = Vec::new();

    for x_entry in x_entries {
        let entry = dbc.trait_node_entry.get(&x_entry.TraitNodeEntryID)?;

        // #t(block: rust_clone_in_loop) need owned String from borrowed DBC data
        if entry.TraitDefinitionID <= 0 {
            let sub_tree_name = dbc
                .trait_sub_tree
                .get(&entry.TraitSubTreeID)
                .and_then(|st| st.Name_lang.clone())
                .unwrap_or_default();

            entries.push(TraitNodeEntry {
                id: entry.ID,
                definition_id: 0,
                spell_id: 0,
                override_spell_id: 0,
                replaces_spell_id: 0,
                sub_tree_id: entry.TraitSubTreeID,
                max_ranks: entry.MaxRanks,
                name: sub_tree_name,
                description: String::new(),
                icon_file_name: String::new(),
                effect_points: Vec::new(),
            });
            continue;
        }

        let definition = dbc.trait_definition.get(&entry.TraitDefinitionID)?;

        let icon_file_name = get_icon_file_name(dbc, definition);
        let name = get_talent_name(dbc, definition);
        let description = get_talent_description(dbc, definition);

        entries.push(TraitNodeEntry {
            id: entry.ID,
            definition_id: definition.ID,
            spell_id: definition.SpellID,
            override_spell_id: definition.VisibleSpellID,
            replaces_spell_id: definition.OverridesSpellID,
            sub_tree_id: entry.TraitSubTreeID,
            max_ranks: entry.MaxRanks,
            name,
            description,
            icon_file_name,
            effect_points: trait_effect_points(dbc, definition.ID),
        });
    }

    if entries.is_empty() {
        return None;
    }

    let max_ranks = if TraitNodeType::from_dbc(node.Type) == Some(TraitNodeType::Tiered) {
        entries.iter().map(|e| e.max_ranks).sum()
    } else {
        entries[0].max_ranks
    };

    let (pos_x, pos_y) = apply_hero_offset(node, &indexes.hero_tree_offsets);
    let tree_index = tree_index_override
        .unwrap_or_else(|| determine_tree_index(dbc, node, &indexes.group_to_tree_index));

    Some(TraitNode {
        id: node.ID,
        pos_x,
        pos_y,
        max_ranks,
        granted_ranks: granted_ranks_for_spec(dbc, node.ID, spec_id, &indexes.spec_set_to_specs),
        node_type: node.Type,
        tree_index,
        order_index: indexes.order_index.get(&node.ID).copied().unwrap_or(-1),
        sub_tree_id: node.TraitSubTreeID,
        entries,
    })
}

fn trait_effect_points(dbc: &DbcData, definition_id: i32) -> Vec<TraitEffectPoints> {
    let mut effects: Vec<_> = dbc
        .trait_definition_effect_points
        .get(&definition_id)
        .into_iter()
        .flatten()
        .filter_map(|effect| {
            let operation = TraitEffectOperation::from_dbc(effect.OperationType)?;
            let mut points = dbc.curve_point.get(&effect.CurveID)?.clone();

            points.sort_by_key(|point| point.OrderIndex);

            Some(TraitEffectPoints {
                effect_index: effect.EffectIndex,
                operation,
                rank_values: points
                    .into_iter()
                    .map(|point| (point.Pos_0, point.Pos_1))
                    .collect(),
            })
        })
        .collect();

    effects.sort_by_key(|effect| effect.effect_index);

    effects
}

fn apply_hero_offset(node: &TraitNodeRow, offsets: &IntMap<i32, (i32, i32)>) -> (i32, i32) {
    if node.TraitSubTreeID > 0 {
        if let Some(&(dx, dy)) = offsets.get(&node.TraitSubTreeID) {
            return (node.PosX + dx, node.PosY + dy);
        }
    }

    (node.PosX, node.PosY)
}

fn determine_tree_index(
    dbc: &DbcData,
    node: &TraitNodeRow,
    group_to_tree_index: &IntMap<i32, i32>,
) -> i32 {
    if TraitNodeType::from_dbc(node.Type) == Some(TraitNodeType::SubtreeSelection) {
        return TraitTreeIndex::Selection as i32;
    }

    if node.TraitSubTreeID > 0 {
        return TraitTreeIndex::Hero as i32;
    }

    if let Some(memberships) = dbc.trait_node_group_x_trait_node.get(&node.ID) {
        for m in memberships {
            if let Some(&idx) = group_to_tree_index.get(&m.TraitNodeGroupID) {
                return idx;
            }
        }
    }

    TraitTreeIndex::Class as i32
}

fn get_talent_name(dbc: &DbcData, definition: &TraitDefinitionRow) -> String {
    if let Some(name) = &definition.OverrideName_lang {
        return name.clone();
    }

    if definition.SpellID > 0 {
        if let Some(spell) = dbc.spell_name.get(&definition.SpellID) {
            if let Some(name) = &spell.Name_lang {
                return name.clone();
            }
        }

        return format!("Spell {}", definition.SpellID);
    }

    "Unknown Talent".to_string()
}

fn get_talent_description(dbc: &DbcData, definition: &TraitDefinitionRow) -> String {
    if let Some(desc) = &definition.OverrideDescription_lang {
        return desc.clone();
    }

    if definition.SpellID > 0 {
        if let Some(spell) = dbc.spell.get(&definition.SpellID) {
            if let Some(desc) = &spell.Description_lang {
                return desc.clone();
            }
        }
    }

    String::new()
}

fn build_subtrees(
    dbc: &DbcData,
    valid_sub_tree_ids: &FastSet<i32>,
    nodes: &[TraitNode],
) -> Vec<TraitSubTree> {
    valid_sub_tree_ids
        .iter()
        .filter_map(|&sub_tree_id| {
            let sub_tree = dbc.trait_sub_tree.get(&sub_tree_id)?;

            let icon_file_name = dbc
                .ui_texture_atlas_element
                .get(&sub_tree.UiTextureAtlasElementID)
                .map_or_else(
                    || derive_sub_tree_icon(nodes, sub_tree_id),
                    |e| e.Name.clone(),
                );

            Some(TraitSubTree {
                id: sub_tree_id,
                name: sub_tree.Name_lang.clone().unwrap_or_default(),
                description: sub_tree.Description_lang.clone().unwrap_or_default(),
                icon_file_name,
            })
        })
        .collect()
}

fn calculate_point_limits(dbc: &DbcData, all_tree_ids: &[i32]) -> PointLimits {
    let tree_currencies: Vec<_> = all_tree_ids
        .iter()
        .flat_map(|&tid| {
            dbc.trait_tree_x_trait_currency
                .get(&tid)
                .cloned()
                .unwrap_or_default()
        })
        .collect();

    let currency_ids: FastSet<i32> = tree_currencies
        .iter()
        .map(|tc| tc.TraitCurrencyID)
        .collect();

    let currency_totals: IntMap<i32, i32> = currency_ids
        .iter()
        .map(|&currency_id| {
            let total = dbc
                .trait_currency_source
                .get(&currency_id)
                .map_or(0, |sources| sources.iter().map(|s| s.Amount).sum());

            (currency_id, total)
        })
        .collect();

    let mut class_limit = 0;
    let mut spec_limit = 0;
    let mut hero_limit = 0;

    for tc in tree_currencies {
        let Some(currency) = dbc.trait_currency.get(&tc.TraitCurrencyID) else {
            continue;
        };
        let total = currency_totals
            .get(&tc.TraitCurrencyID)
            .copied()
            .unwrap_or(0);

        let flags = TraitCurrencyFlags::from_dbc(currency.Flags);

        if flags.contains(TraitCurrencyFlags::CLASS) {
            class_limit += total;
        } else if flags.contains(TraitCurrencyFlags::SPECIALIZATION) {
            spec_limit += total;
        } else if flags.is_empty() {
            hero_limit = hero_limit.max(total);
        }
    }

    PointLimits {
        class: class_limit,
        spec: spec_limit,
        hero: hero_limit,
    }
}

fn calculate_hero_tree_offsets(
    nodes: &[TraitNodeRow],
    sub_tree_ids: &[i32],
) -> IntMap<i32, (i32, i32)> {
    let mut offsets = IntMap::default();

    for &sub_tree_id in sub_tree_ids {
        let hero_nodes: Vec<_> = nodes
            .iter()
            .filter(|n| n.TraitSubTreeID == sub_tree_id)
            .collect();

        if hero_nodes.is_empty() {
            continue;
        }

        let min_x = hero_nodes.iter().map(|n| n.PosX).min().unwrap_or(0);
        let min_y = hero_nodes.iter().map(|n| n.PosY).min().unwrap_or(0);

        offsets.insert(sub_tree_id, (TARGET_HERO_X - min_x, TARGET_HERO_Y - min_y));
    }

    offsets
}

fn get_icon_file_name(dbc: &DbcData, definition: &TraitDefinitionRow) -> String {
    let icon_file_data_id = if definition.OverrideIcon > 0 {
        definition.OverrideIcon
    } else if definition.SpellID > 0 {
        dbc.spell_misc
            .get(&definition.SpellID)
            .map_or(0, |m| m.SpellIconFileDataID)
    } else {
        0
    };

    super::resolve_icon_file_name(dbc, icon_file_data_id)
}

fn derive_sub_tree_icon(nodes: &[TraitNode], sub_tree_id: i32) -> String {
    let mut entries = nodes
        .iter()
        .filter(|n| n.sub_tree_id == sub_tree_id)
        .flat_map(|n| &n.entries);
    let icon_entry = entries.find(|e| !e.icon_file_name.is_empty());

    icon_entry.map_or_else(
        || wowlab_types::constants::DEFAULT_ICON_FILENAME.to_string(),
        |e| e.icon_file_name.clone(),
    )
}

fn get_valid_subtrees_for_spec(
    dbc: &DbcData,
    tree_nodes: &[TraitNodeRow],
    spec_id: i32,
    spec_set_to_specs: &IntMap<i32, Vec<i32>>,
) -> FastSet<i32> {
    let mut valid_subtrees = FastSet::default();

    let selection_nodes = tree_nodes
        .iter()
        .filter(|n| TraitNodeType::from_dbc(n.Type) == Some(TraitNodeType::SubtreeSelection));

    for node in selection_nodes {
        if !should_include_node_for_spec(dbc, node.ID, spec_id, spec_set_to_specs) {
            continue;
        }

        let entries = dbc
            .trait_node_x_trait_node_entry
            .get(&node.ID)
            .cloned()
            .unwrap_or_default();

        for x_entry in entries {
            if let Some(entry) = dbc.trait_node_entry.get(&x_entry.TraitNodeEntryID) {
                if entry.TraitSubTreeID > 0 {
                    valid_subtrees.insert(entry.TraitSubTreeID);
                }
            }
        }
    }

    valid_subtrees
}

fn build_spec_set_lookup(dbc: &DbcData) -> IntMap<i32, Vec<i32>> {
    let mut map: IntMap<i32, Vec<i32>> = IntMap::default();

    for members in dbc.spec_set_member.values() {
        for member in members {
            map.entry(member.SpecSet)
                .or_default()
                .push(member.ChrSpecializationID);
        }
    }

    map
}

fn granted_ranks_for_spec(
    dbc: &DbcData,
    node_id: i32,
    spec_id: i32,
    spec_set_to_specs: &IntMap<i32, Vec<i32>>,
) -> i32 {
    let node_conditions = dbc
        .trait_node_x_trait_cond
        .get(&node_id)
        .into_iter()
        .flatten()
        .filter_map(|reference| dbc.trait_cond.get(&reference.TraitCondID));
    let group_conditions = dbc
        .trait_node_group_x_trait_node
        .get(&node_id)
        .into_iter()
        .flatten()
        .flat_map(|membership| {
            dbc.trait_node_group_x_trait_cond
                .get(&membership.TraitNodeGroupID)
                .into_iter()
                .flatten()
        })
        .filter_map(|reference| dbc.trait_cond.get(&reference.TraitCondID));

    let starter_conditions = node_conditions
        .chain(group_conditions)
        .filter(|condition| {
            TraitConditionType::try_from(condition.CondType).ok()
                == Some(TraitConditionType::StarterBuild)
        })
        .filter(|condition| {
            spec_set_to_specs
                .get(&condition.SpecSetID)
                .is_some_and(|specs| specs.contains(&spec_id))
        });

    starter_conditions
        .map(|condition| condition.GrantedRanks)
        .max()
        .unwrap_or(0)
}

fn should_include_node_for_spec(
    dbc: &DbcData,
    node_id: i32,
    spec_id: i32,
    spec_set_to_specs: &IntMap<i32, Vec<i32>>,
) -> bool {
    let mut has_spec_condition = false;
    let mut spec_allowed = false;

    let node_conds = dbc
        .trait_node_x_trait_cond
        .get(&node_id)
        .cloned()
        .unwrap_or_default();

    for node_cond in &node_conds {
        let Some(cond) = dbc.trait_cond.get(&node_cond.TraitCondID) else {
            continue;
        };

        if TraitConditionType::try_from(cond.CondType).ok()
            != Some(TraitConditionType::Specialization)
            || cond.SpecSetID <= 0
        {
            continue;
        }

        has_spec_condition = true;

        if let Some(specs) = spec_set_to_specs.get(&cond.SpecSetID) {
            if specs.contains(&spec_id) {
                spec_allowed = true;
            }
        }
    }

    let node_groups = dbc
        .trait_node_group_x_trait_node
        .get(&node_id)
        .cloned()
        .unwrap_or_default();

    for node_group in &node_groups {
        let group_cond_refs = dbc
            .trait_node_group_x_trait_cond
            .get(&node_group.TraitNodeGroupID)
            .cloned()
            .unwrap_or_default();

        for cond_ref in &group_cond_refs {
            let Some(cond) = dbc.trait_cond.get(&cond_ref.TraitCondID) else {
                continue;
            };

            if TraitConditionType::try_from(cond.CondType).ok()
                != Some(TraitConditionType::Specialization)
                || cond.SpecSetID <= 0
            {
                continue;
            }

            has_spec_condition = true;

            if let Some(specs) = spec_set_to_specs.get(&cond.SpecSetID) {
                if specs.contains(&spec_id) {
                    spec_allowed = true;
                }
            }
        }
    }

    !has_spec_condition || spec_allowed
}

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