frequenz-microgrid-component-graph 0.5.0

// License: MIT
// Copyright © 2025 Frequenz Energy-as-a-Service GmbH

//! This module defines the `Formula` struct and its operations.

use super::expr::Expr;

/// A trait with methods for combining formulas.
pub trait Formula {
    fn coalesce(self, other: Self) -> Self;
    fn min(self, other: Self) -> Self;
    fn max(self, other: Self) -> Self;
}

/// Blanket implementation of the `Formula` trait for any type that can be
/// converted to and from `Expr`.
impl<T> Formula for T
where
    T: From<Expr>,
    Expr: From<T>,
{
    fn coalesce(self, other: Self) -> Self {
        Expr::coalesce(self.into(), other.into()).into()
    }

    fn min(self, other: Self) -> Self {
        Expr::min(self.into(), other.into()).into()
    }

    fn max(self, other: Self) -> Self {
        Expr::max(self.into(), other.into()).into()
    }
}

// This struct represents the microgrid metric formulas that are generated by
// traversing the component graph.
//
// They are used to represent quantities that are aggregated over multiple
// components in the microgrid, such as power or current.
//
// `AggregationFormula` objects can be added or subtracted from each other, and
// they can converted to a string representation, before they are passed to an
// evaluator.
#[derive(Debug, Clone, PartialEq)]
pub struct AggregationFormula {
    pub(crate) expr: Expr,
}

impl From<Expr> for AggregationFormula {
    fn from(expr: Expr) -> Self {
        AggregationFormula { expr }
    }
}

impl From<AggregationFormula> for Expr {
    fn from(formula: AggregationFormula) -> Self {
        formula.expr
    }
}

impl AggregationFormula {
    pub(crate) fn new(expr: Expr) -> Self {
        AggregationFormula { expr }
    }
}

impl std::ops::Add for AggregationFormula {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        AggregationFormula {
            expr: self.expr + rhs.expr,
        }
    }
}

impl std::ops::Sub for AggregationFormula {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        AggregationFormula {
            expr: self.expr - rhs.expr,
        }
    }
}

impl std::fmt::Display for AggregationFormula {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.expr.fmt(f)
    }
}

impl From<AggregationFormula> for String {
    fn from(formula: AggregationFormula) -> Self {
        formula.expr.to_string()
    }
}

/// Represents a formula that coalesces metrics from multiple components.
///
/// This is typically used for non-aggregating metrics like AC voltage or
/// frequency.
#[derive(Debug, Clone, PartialEq)]
pub struct CoalesceFormula {
    expr: Expr,
}

impl From<Expr> for CoalesceFormula {
    fn from(expr: Expr) -> Self {
        CoalesceFormula { expr }
    }
}

impl From<CoalesceFormula> for Expr {
    fn from(formula: CoalesceFormula) -> Self {
        formula.expr
    }
}

impl CoalesceFormula {
    pub(crate) fn new(expr: Expr) -> Self {
        CoalesceFormula { expr }
    }
}

impl std::fmt::Display for CoalesceFormula {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.expr.fmt(f)
    }
}

impl From<CoalesceFormula> for String {
    fn from(formula: CoalesceFormula) -> Self {
        formula.expr.to_string()
    }
}

#[test]
fn test_aggregation_formula_arith() {
    let formula1 = AggregationFormula::new(Expr::component(1));
    let formula2 = AggregationFormula::new(Expr::component(2));

    let result = formula1.clone() + formula2.clone();
    assert_eq!(result.to_string(), "#1 + #2");

    let result = formula1 - formula2;
    assert_eq!(result.to_string(), "#1 - #2");
}

#[test]
fn test_formula_trait() {
    let formula1 = AggregationFormula::new(Expr::component(1));
    let formula2 = AggregationFormula::new(Expr::component(2));

    let coalesce_result = formula1.clone().coalesce(formula2.clone());
    assert_eq!(coalesce_result.to_string(), "COALESCE(#1, #2)");

    let min_result = formula1.clone().min(formula2.clone());
    assert_eq!(min_result.to_string(), "MIN(#1, #2)");

    let max_result = formula1.max(formula2);
    assert_eq!(max_result.to_string(), "MAX(#1, #2)");

    let formula3 = CoalesceFormula::new(Expr::component(3));
    let formula4 = CoalesceFormula::new(Expr::component(4));

    let coalesce_result = formula3.clone().coalesce(formula4.clone());
    assert_eq!(coalesce_result.to_string(), "COALESCE(#3, #4)");

    let min_result = formula3.clone().min(formula4.clone());
    assert_eq!(min_result.to_string(), "MIN(#3, #4)");

    let max_result = formula3.max(formula4);
    assert_eq!(max_result.to_string(), "MAX(#3, #4)");
}