codegen/

gen_aura.rs

//! Emission of `aura_<name>` builder functions from manifest aura definitions.

use anyhow::{Context, bail};
use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use wowlab_manifest_schema::{Manifest, ManifestAuraDef, SpellGroupRule};
use wowlab_types::sim::FastMap;

use crate::{
    gen_aura_behavior::{apply_aura_consume_effects, apply_refresh_behavior, apply_tick_behavior},
    gen_aura_modifiers::{
        apply_per_stack_scalars, apply_rating_effects, apply_spell_scoped_effects,
    },
    gen_aura_periodic::{
        apply_aura_duration, apply_periodic_damage, apply_periodic_hook_source,
        apply_periodic_hooks, periodic_behavior_spell_id,
    },
    gen_emit::{
        require_get, scalar_ref_call, spell_idx, validated_max_stacks_call,
        validated_milliseconds_call,
    },
    helpers::to_snake,
    rust_source::{ExpressionRequirements, expression_requirements, lit_int, result_expression},
};

pub(crate) fn aura_functions(
    manifest: &Manifest,
    module_name: &str,
) -> anyhow::Result<TokenStream> {
    let aura_count = manifest
        .spell_groups
        .iter()
        .map(|group| group.auras.len())
        .sum();
    let mut spell_group_map = FastMap::default();

    spell_group_map.reserve(aura_count);

    for (group_index, group) in manifest.spell_groups.iter().enumerate() {
        let group_index =
            u8::try_from(group_index).context("manifest has more than 256 spell groups")?;

        for aura_name in &group.auras {
            spell_group_map.insert(aura_name.as_str(), (group_index, group.rule));
        }
    }

    let mut aura_idx_map = FastMap::default();

    aura_idx_map.reserve(manifest.auras.len());

    for (index, name) in manifest.auras.keys().enumerate() {
        let index = u8::try_from(index).context("manifest has more than 256 auras")?;

        aura_idx_map.insert(name.as_str(), index);
    }

    let spell_id_map = manifest_spell_id_map(manifest)?;
    let functions = manifest
        .auras
        .iter()
        .map(|(name, def)| {
            let options = AuraEmitOptions {
                spell_group: spell_group_map.get(name.as_str()).copied(),
                aura_idx_map: &aura_idx_map,
                spell_id_map: &spell_id_map,
                emit_default_max_stacks: def.max_stacks.is_none(),
                has_item_budget: false,
                hook_module: module_name,
                hook_suffix: "_hook",
            };

            emit_aura_fn(name, def, &options)
        })
        .collect::<anyhow::Result<Vec<_>>>()?;

    Ok(quote!(#(#functions)*))
}

fn manifest_spell_id_map(manifest: &Manifest) -> anyhow::Result<FastMap<&str, u32>> {
    let mut ids: FastMap<&str, u32> = manifest
        .spells
        .iter()
        .map(|(name, spell)| (name.as_str(), spell.id))
        .collect();

    for (name, id) in &manifest.reported_spells {
        ids.entry(name).or_insert(*id);
    }

    let hero_spell_count = manifest
        .hero_talents
        .values()
        .map(|tree| tree.spells.len())
        .sum();
    let mut hero_ids = FastMap::default();

    hero_ids.reserve(hero_spell_count);

    for tree in manifest.hero_talents.values() {
        for (name, id) in &tree.spells {
            if let Some(previous) = hero_ids.insert(name.as_str(), *id)
                && previous != *id
                && !ids.contains_key(name.as_str())
            {
                bail!("hero spell binding {name:?} is ambiguous between ids {previous} and {id}");
            }

            ids.entry(name).or_insert(*id);
        }
    }

    Ok(ids)
}

/// Controls item- and spec-specific aura emission.
pub(crate) struct AuraEmitOptions<'a> {
    pub spell_group: Option<(u8, SpellGroupRule)>,
    pub aura_idx_map: &'a FastMap<&'a str, u8>,
    pub spell_id_map: &'a FastMap<&'a str, u32>,
    pub emit_default_max_stacks: bool,
    pub has_item_budget: bool,
    pub hook_module: &'a str,
    pub hook_suffix: &'a str,
}

struct GeneratedAuraBuilderCode {
    expression: TokenStream,
    requirements: ExpressionRequirements,
    _stage: AuraBuilderStage,
}

struct AuraBuilderStage;

impl GeneratedAuraBuilderCode {
    fn uses_data(&self) -> bool {
        self.requirements.data
    }

    fn uses_budget(&self) -> bool {
        self.requirements.budget
    }

    fn into_result_expression(self) -> anyhow::Result<TokenStream> {
        Ok(result_expression(self.expression)?)
    }

    #[cfg(test)]
    fn body_string(&self) -> String {
        self.expression.to_string()
    }
}

fn lower_manifest_aura(
    def: &ManifestAuraDef,
    options: &AuraEmitOptions<'_>,
) -> anyhow::Result<GeneratedAuraBuilderCode> {
    let id = def.id;
    let spell_id = spell_idx(id);
    let mut expression = match def.on.as_str() {
        "target" => quote!(a.on_target()),
        "player" => quote!(a.on_player()),
        "pet" => quote!(a.on_pet()),
        other => bail!("unsupported aura target {other:?} on aura {id}"),
    };

    expression = apply_core_aura_options(expression, def, options, &spell_id)?;
    expression = apply_tick_behavior(expression, def)?;

    if let Some(remove_stack) = &def.periodic_remove_stack {
        let tick_ms = validated_milliseconds_call(
            &remove_stack.tick_ms,
            id,
            "periodic_remove_stack.tick_ms",
        )?;

        expression = quote!(#expression.periodic_remove_stack(#tick_ms));
    }

    if let Some(value) = &def.primary_stat_percent_per_stack {
        let value = scalar_ref_call(value)?;

        expression = quote!(#expression.primary_pct_per_stack(#value)?);
    }

    if let Some(value) = &def.attack_power_percent_per_stack {
        let value = scalar_ref_call(value)?;

        expression = quote!(#expression.attack_power_pct_per_stack(#value)?);
    }

    expression = apply_spell_scoped_effects(expression, def, options)?;
    expression = apply_periodic_damage(expression, def, id, options.has_item_budget)?;

    if let Some(resource) = &def.periodic_resource {
        let tick_ms =
            validated_milliseconds_call(&resource.tick_ms, id, "periodic_resource.tick_ms")?;
        let amount = scalar_ref_call(&resource.amount)?;

        expression = quote!(#expression.periodic_resource(#tick_ms, #amount));
    }

    if let Some(drain) = &def.periodic_resource_drain {
        let tick_ms =
            validated_milliseconds_call(&drain.tick_ms, id, "periodic_resource_drain.tick_ms")?;
        let amount = scalar_ref_call(&drain.amount)?;

        expression = quote!(#expression.periodic_resource_drain(#tick_ms, #amount));
    }

    if let Some(apply) = &def.periodic_apply_aura {
        let index = options
            .aura_idx_map
            .get(apply.aura.as_str())
            .copied()
            .ok_or_else(|| {
                anyhow::anyhow!(
                    "periodic_apply_aura.aura {:?} not found in [auras.*]",
                    apply.aura
                )
            })?;
        let tick_ms =
            validated_milliseconds_call(&apply.tick_ms, id, "periodic_apply_aura.tick_ms")?;
        let index = lit_int(index);

        expression = quote!(#expression.periodic_apply_aura(#tick_ms, #index));
    }

    expression = apply_periodic_hook_source(expression, def)?;

    if let Some(source_spell_id) = periodic_behavior_spell_id(def) {
        let source_spell_id = lit_int(source_spell_id);

        expression = quote!(#expression.apply_periodic_behavior_from_data(data, #source_spell_id)?);
    }

    expression = apply_periodic_hooks(expression, def, options)?;
    expression = apply_aura_consume_effects(expression, def, options)?;

    if let Some((group_index, rule)) = options.spell_group {
        let group_index = lit_int(group_index);
        let rule = match rule {
            SpellGroupRule::Exclusive => {
                quote!(wowlab_types::combat::SpellGroupRule::Exclusive)
            }
        };

        expression = quote!(#expression.spell_group(#group_index, #rule));
    }

    let requirements = expression_requirements(expression.clone())?;

    Ok(GeneratedAuraBuilderCode {
        expression,
        requirements,
        _stage: AuraBuilderStage,
    })
}

fn apply_core_aura_options(
    mut expression: TokenStream,
    def: &ManifestAuraDef,
    options: &AuraEmitOptions<'_>,
    spell_id: &TokenStream,
) -> anyhow::Result<TokenStream> {
    let id = def.id;

    expression = apply_aura_duration(expression, def)?;

    if options.emit_default_max_stacks {
        let max_stacks = require_get("aura_max_stacks", spell_id);

        expression = quote!(#expression.max_stacks(#max_stacks));
    }

    expression = apply_refresh_behavior(expression, def)?;

    if def.snapshot {
        expression = quote!(#expression.snapshot());
    }

    if let Some(value) = &def.forwarded_effect_value {
        let value = scalar_ref_call(value)?;

        expression = quote!(#expression.forwarded_effect_value(data, #value)?);
    }

    if let Some(multiplier) = &def.damage_mult {
        let multiplier = scalar_ref_call(multiplier)?;

        expression = quote!(#expression.damage_multiplier(#multiplier)?);
    }

    if let Some(multiplier) = &def.damage_mult_stacking {
        let initial = scalar_ref_call(&multiplier.initial)?;
        let per_stack = scalar_ref_call(&multiplier.per_stack)?;

        expression = quote!(#expression.damage_multiplier_stacking(#initial, #per_stack)?);
    }

    if let Some(multiplier) = &def.pet_damage_mult {
        let multiplier = scalar_ref_call(multiplier)?;

        expression = quote!(#expression.pet_damage_multiplier(#multiplier)?);
    }

    if let Some(stacks) = &def.max_stacks {
        let stacks = validated_max_stacks_call(stacks, id, "aura.max_stacks")?;

        expression = quote!(#expression.max_stacks(#stacks));
    }

    expression = apply_per_stack_scalars(expression, def)?;
    expression = apply_rating_effects(expression, def, options)?;

    if def.apply_at_max_stacks {
        expression = quote!(#expression.apply_at_max_stacks());
    }

    if def.reverse {
        expression = quote!(#expression.reverse());
    }

    if def.freeze_stacks {
        expression = quote!(#expression.freeze_stacks());
    }

    if def.async_stacks {
        expression = quote!(#expression.async_stacks());
    }

    if def.periodic_damage_scales_with_stacks {
        expression = quote!(#expression.periodic_damage_scales_with_stacks());
    }

    if def.reapply_on_expire {
        expression = quote!(#expression.reapply_on_expire());
    }

    Ok(expression)
}

/// Emit a single `aura_<name>` builder function.
pub(crate) fn emit_aura_fn(
    name: &str,
    def: &ManifestAuraDef,
    options: &AuraEmitOptions<'_>,
) -> anyhow::Result<TokenStream> {
    let function_name = format_ident!("aura_{}", to_snake(name));
    let generated = lower_manifest_aura(def, options)?;
    let data_parameter = if generated.uses_data() {
        quote!(data)
    } else {
        quote!(_data)
    };
    let uses_budget = generated.uses_budget();
    let expression = generated.into_result_expression()?;

    Ok(if options.has_item_budget {
        let budget_parameter = if uses_budget {
            quote!(budget)
        } else {
            quote!(_budget)
        };

        quote! {
            fn #function_name(
                a: AuraDefinitionDraft,
                #data_parameter: &ResolvedGameData,
                #budget_parameter: ItemBudget,
            ) -> Result<AuraDefinitionDraft, wowlab_engine_combat::BuilderError> {
                #expression
            }
        }
    } else {
        quote! {
            fn #function_name(
                a: AuraDefinitionDraft,
                #data_parameter: &ResolvedGameData,
            ) -> Result<AuraDefinitionDraft, wowlab_engine_combat::BuilderError> {
                #expression
            }
        }
    })
}

#[cfg(test)]
#[path = "gen_aura/periodic_tests.rs"]
mod periodic_tests;

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