use crate::opt::lang::MathLang;
use egg::{CostFunction, Id, Language};

#[derive(Default)]
pub struct FusionCost;

impl CostFunction<MathLang> for FusionCost {
    type Cost = usize;

    fn cost<C>(&mut self, enode: &MathLang, mut costs: C) -> Self::Cost
    where
        C: FnMut(Id) -> Self::Cost,
    {
        let base = match enode {
            MathLang::FusedExpSigmoid(_) | MathLang::FusedSigmoidLog(_) => 1,
            MathLang::Bool(_) | MathLang::Int(_) | MathLang::Num(_) | MathLang::TypeToken(_) => 1,
            _ => 2,
        };
        enode.fold(base, |acc, id| acc + costs(id))
    }
}
