codegen/

gen_spells.rs

use proc_macro2::TokenStream;
use quote::{format_ident, quote};
use wowlab_manifest_schema::{Manifest, ManifestSpellDef};

use crate::{
    gen_aoe::apply_aoe,
    gen_damage::apply_damage,
    gen_emit::{effect_lookup, validated_milliseconds_call, validated_nonnegative_call},
    gen_spell_channel::apply_channel,
    gen_spell_properties::{apply_core_spell_properties, apply_override},
    helpers::to_snake,
    rust_source::{
        ExpressionRequirements, expression_requirements, lit_f64, lit_int, result_expression,
        rust_ident, rust_path,
    },
};

/// Which builder-code shape the spell lowerer emits.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub(crate) enum SpellChainKind {
    Spec,
    Item,
}

/// Options that let spell lowering reproduce either the spec or item spell function.
pub(crate) struct SpellChainOpts<'a> {
    pub kind: SpellChainKind,
    pub inferred_aura: Option<(&'a str, u32)>,
    pub hook_module: &'a str,
    pub hook_suffix: &'a str,
}

pub(crate) fn spell_functions(
    manifest: &Manifest,
    module_name: &str,
) -> anyhow::Result<TokenStream> {
    let functions = manifest
        .spells
        .iter()
        .map(|(name, def)| {
            let function_name = format_ident!("spell_{}", to_snake(name));
            let options = SpellChainOpts {
                kind: SpellChainKind::Spec,
                inferred_aura: infer_self_applied_aura(name, def, manifest),
                hook_module: module_name,
                hook_suffix: "_hook",
            };
            let generated = lower_manifest_spell(def, &options)?;
            let data_parameter = if generated.uses_data() {
                quote!(data)
            } else {
                quote!(_data)
            };
            let expression = generated.into_result_expression()?;

            Ok::<_, anyhow::Error>(quote! {
                fn #function_name(
                    s: SpellDefinitionDraft,
                    #data_parameter: &ResolvedGameData,
                ) -> Result<SpellDefinitionDraft, wowlab_engine_combat::BuilderError> {
                    #expression
                }
            })
        })
        .collect::<anyhow::Result<Vec<_>>>()?;

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

fn infer_self_applied_aura<'a>(
    spell_name: &'a str,
    def: &ManifestSpellDef,
    manifest: &'a Manifest,
) -> Option<(&'a str, u32)> {
    if !def.infer_applies_aura || def.applies_aura.is_some() || !def.applies_auras.is_empty() {
        return None;
    }

    if let Some(aura) = manifest.auras.get(spell_name) {
        return Some((spell_name, aura.id));
    }

    let matches = manifest
        .auras
        .iter()
        .filter(|(_, aura)| aura.id == def.id)
        .map(|(name, _)| name.as_str())
        .collect::<Vec<_>>();

    match matches.as_slice() {
        [only] => Some((only, def.id)),
        _ => None,
    }
}

/// Generated Rust code at the spell-builder stage.
pub(crate) struct GeneratedSpellBuilderCode {
    expression: TokenStream,
    requirements: ExpressionRequirements,
}

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

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

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

/// Lower one manifest spell into generated spell-builder code.
pub(crate) fn lower_manifest_spell(
    def: &ManifestSpellDef,
    options: &SpellChainOpts<'_>,
) -> anyhow::Result<GeneratedSpellBuilderCode> {
    let id = def.id;
    let raw_id = lit_int(id);
    let mut expression = match options.kind {
        SpellChainKind::Spec => {
            let mut expression = quote!(s.apply_base_from_data(data, #raw_id)?);

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

            if def.infer_damage && !def.infer_primary_damage {
                expression = quote!(#expression.without_primary_data_damage());
            }

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

            expression
        }
        SpellChainKind::Item => {
            let mut expression = quote!(s.apply_cooldown_pools_from_data(data, #raw_id)?);

            if let Some(cooldown) = &def.cooldown {
                let cooldown = validated_nonnegative_call(cooldown, id, "spell.cooldown")?;

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

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

            expression
        }
    };

    match options.kind {
        SpellChainKind::Spec => {
            expression = spec_expression(expression, def, options)?;
        }
        SpellChainKind::Item => {
            if let Some(aura) = def.applies_aura.as_deref() {
                let aura = rust_ident(aura)?;

                expression = quote!(#expression.applies_aura_id(AURA::#aura));
            }

            for aura in &def.applies_auras {
                let aura = rust_ident(aura)?;

                expression = quote!(#expression.applies_aura_id(AURA::#aura));
            }

            if let Some(hook) = &def.hook {
                let module = rust_path(options.hook_module)?;
                let hook = format_ident!("{hook}{}", options.hook_suffix);

                expression = quote!(#expression.on_cast(crate::hooks::#module::#hook));
            }
        }
    }

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

    Ok(GeneratedSpellBuilderCode {
        expression,
        requirements,
    })
}

fn spec_expression(
    mut expression: TokenStream,
    def: &ManifestSpellDef,
    options: &SpellChainOpts<'_>,
) -> anyhow::Result<TokenStream> {
    let id = def.id;

    expression = apply_core_spell_properties(expression, id, def, options.hook_module)?;
    expression = apply_damage_payloads(expression, def, options.hook_module)?;

    if let Some(aura) = def.applies_aura.as_deref() {
        let aura = format_ident!("AURA_{aura}");

        expression = quote!(#expression.applies_aura(#aura));
    } else if let Some((aura, aura_id)) = options.inferred_aura {
        let aura = format_ident!("AURA_{aura}");
        let aura_id = lit_int(aura_id);

        expression = quote!(#expression.applies_inferred_aura(#aura, #aura_id));
    }

    for aura in &def.applies_auras {
        let aura = format_ident!("AURA_{aura}");

        expression = quote!(#expression.applies_aura(#aura));
    }

    for extension in &def.extends_auras {
        let aura = format_ident!("AURA_{}", extension.aura);
        let amount_ms = lit_int(extension.amount_ms);

        expression = quote!(#expression.extends_aura(#aura, #amount_ms));
    }

    if let Some(cooldown_reductions) = &def.reduces_cd {
        for reduction in cooldown_reductions {
            let target = format_ident!("SPELL_{}", reduction.target);
            let amount = validated_milliseconds_call(
                &reduction.amount_ms,
                id,
                "spell.reduces_cd.amount_ms",
            )?;

            expression = quote!(#expression.reduces_cd(#target, #amount));
        }
    }

    if let Some(reduction) = &def.reduces_cd_chance {
        let target = format_ident!("SPELL_{}", reduction.target);
        let amount = validated_milliseconds_call(
            &reduction.amount_ms,
            id,
            "spell.reduces_cd_chance.amount_ms",
        )?;
        let chance = lit_f64(reduction.chance);

        expression = quote!(#expression.reduces_cd_chance(#target, #amount, #chance));
    }

    if let Some(reset) = &def.resets_cd_while {
        let target = format_ident!("SPELL_{}", reset.target);
        let aura = format_ident!("AURA_{}", reset.aura);

        expression = quote!(#expression.resets_cd_while(#target, #aura));
    }

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

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

    expression = apply_aoe(expression, def);

    if let Some(channel) = &def.channel {
        expression = apply_channel(expression, id, channel, options.hook_module)?;
    }

    if let Some(hook) = &def.hook {
        let module = rust_path(options.hook_module)?;
        let hook = format_ident!("{hook}{}", options.hook_suffix);

        expression = quote!(#expression.on_cast(crate::hooks::#module::#hook));
    }

    if let Some(aura_id) = def.requires_aura_id {
        let aura_id = lit_int(aura_id);
        let min_stacks = lit_int(def.requires_aura_min_stacks.unwrap_or(0));

        expression = quote!(#expression.requires_aura_at_min_stacks(#aura_id, #min_stacks));
    }

    if let Some(aura_id) = def.cost_free_when_aura {
        let aura_id = lit_int(aura_id);

        expression = quote!(#expression.cost_free_when_aura(#aura_id));
    }

    if let Some(aura_id) = def.consume_aura_on_cast {
        let aura_id = lit_int(aura_id);

        expression = quote!(#expression.consume_aura_on_cast(#aura_id));
    }

    if let Some(aura) = &def.cooldown_bypass_when_aura {
        let aura = rust_ident(aura)?;

        expression = quote!(#expression.cooldown_bypass_when_aura(AURA::#aura));
    }

    if let Some(threshold) = def.execute_below_pct {
        let threshold = lit_f64(threshold);

        expression = quote!(#expression.execute_below_pct(#threshold));
    }

    if let Some(aura) = &def.instant_when_aura {
        let aura = rust_ident(aura)?;

        expression = quote!(#expression.instant_when_aura(AURA::#aura));
    }

    apply_override(expression, id, def)
}

fn apply_damage_payloads(
    mut expression: TokenStream,
    def: &ManifestSpellDef,
    hook_module: &str,
) -> anyhow::Result<TokenStream> {
    if let Some(damage) = def.damage.as_ref() {
        expression = apply_damage(&expression, damage)?;
    }

    for damage in &def.damage_effects {
        let lookup = effect_lookup(&quote!(data), damage.spell_id, damage.effect);
        let multiplier = lit_f64(damage.multiplier);

        expression = match (damage.kind.as_str(), damage.ap_type.as_deref()) {
            ("ap", None) => {
                quote!(#expression.additional_damage_ap_from_data(#lookup, #multiplier)?)
            }
            ("ap", Some(ap_type)) => {
                let ap_type = match ap_type {
                    "mainhand" => format_ident!("MainHand"),
                    "offhand" => format_ident!("OffHand"),
                    "both" => format_ident!("Both"),
                    "none" => format_ident!("None"),
                    other => anyhow::bail!("unknown damage effect ap_type: {other}"),
                };

                quote! {
                    #expression.additional_damage_ap_from_data_typed(
                        #lookup,
                        #multiplier,
                        wowlab_engine_combat::WeaponApType::#ap_type,
                    )?
                }
            }
            ("sp", _) => {
                quote!(#expression.additional_damage_sp_from_data(#lookup, #multiplier)?)
            }
            (other, _) => anyhow::bail!("unknown damage effect kind: {other}"),
        };
    }

    for trigger_predicate in &def.trigger_program_predicates {
        if trigger_predicate.effects.is_empty() {
            anyhow::bail!("trigger program predicate effects must not be empty");
        }

        for effect_index in &trigger_predicate.effects {
            if *effect_index == 0 {
                anyhow::bail!("trigger program predicate effect indices are one-based");
            }

            let effect_index = lit_int(effect_index);
            let module = rust_path(hook_module)?;
            let predicate = rust_ident(&trigger_predicate.predicate)?;

            expression = quote! {
                #expression.trigger_program_predicate(
                    #effect_index,
                    crate::hooks::#module::#predicate,
                )
            };
        }
    }

    Ok(expression)
}

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