frequenz-microgrid-component-graph 0.5.0

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

use crate::Node;

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum Expr {
    /// An empty expression, which as a formula would evaluate to None.
    None,

    /// A negation of an expression.
    Neg { param: Box<Expr> },

    /// A numeric constant.
    Number { value: f64 },

    /// A reference to a component.
    Component { component_id: u64 },

    /// An addition of multiple expressions.
    Add { params: Vec<Expr> },

    /// A subtraction of multiple expressions.
    Sub { params: Vec<Expr> },

    /// A coalesce function.
    Coalesce { params: Vec<Expr> },

    /// A min function.
    Min { params: Vec<Expr> },

    /// A max function.
    Max { params: Vec<Expr> },
}

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

    fn add(self, rhs: Self) -> Self {
        match (self, rhs) {
            (Self::None, other) | (other, Self::None) => other,
            // -a + -b = -(a + b)
            (Self::Neg { param: lhs }, Self::Neg { param: rhs }) => -(*lhs + *rhs),
            // -a + b = b - a
            // a + -b = a - b
            (other, Self::Neg { param }) | (Self::Neg { param }, other) => other - *param,
            // (a + b) + (c + d) = a + b + c + d
            (Self::Add { params: mut lhs }, Self::Add { params: mut rhs }) => {
                lhs.append(&mut rhs);
                Self::Add { params: lhs }
            }
            // (a + b) + c = a + b + c
            (Self::Add { mut params }, rhs) => {
                params.push(rhs);
                Self::Add { params }
            }
            // a + (b + c) = a + b + c
            (lhs, Self::Add { mut params }) => {
                params.insert(0, lhs);
                Self::Add { params }
            }
            // Catch all other cases
            (lhs, rhs) => Self::Add {
                params: vec![lhs, rhs],
            },
        }
    }
}

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

    fn sub(self, rhs: Self) -> Self {
        match (self, rhs) {
            (Self::None, other) => -other,
            (other, Self::None) => other,
            // (a - b) - -c = a - b + c
            (sub @ Self::Sub { .. }, Self::Neg { param }) => sub + *param,
            // -a - (b - c) = c - b - a
            (Self::Neg { param }, sub @ Self::Sub { .. }) => -sub - *param,
            // (a - b) - c = a - b - c
            (Self::Sub { mut params }, rhs) => {
                params.push(rhs);
                Self::Sub { params }
            }
            // -a - -b = b - a
            (Self::Neg { param: lhs }, Self::Neg { param: rhs }) => Self::Sub {
                params: vec![*rhs, *lhs],
            },
            // -a - b = -(a + b)
            (Self::Neg { param }, value) => -(*param + value),
            // a - -b = a + b
            (lhs, Self::Neg { param }) => lhs + *param,
            // Catch all other cases
            (lhs, rhs) => Self::Sub {
                params: vec![lhs, rhs],
            },
        }
    }
}

impl std::ops::Neg for Expr {
    type Output = Self;

    fn neg(self) -> Self {
        match self {
            Self::None => Self::None,
            // -(-a) = a
            Expr::Neg { param: inner } => *inner,
            // -(a - b) = b - a
            // -(a - b - c) = b + c - a
            Expr::Sub { mut params } => {
                let first = params.remove(0);
                Expr::Add { params } - first
            }
            // Catch all other cases
            _ => Expr::Neg {
                param: Box::new(self),
            },
        }
    }
}

impl<N: Node> From<&N> for Expr {
    fn from(node: &N) -> Self {
        Self::Component {
            component_id: node.component_id(),
        }
    }
}

/// Constructors for `FormulaExpression`.
impl Expr {
    #[must_use]
    pub(crate) fn number(value: f64) -> Self {
        Self::Number { value }
    }

    #[must_use]
    pub(crate) fn component(component_id: u64) -> Self {
        Self::Component { component_id }
    }

    #[must_use]
    pub(crate) fn coalesce(self, other: Expr) -> Self {
        match (self, other) {
            (Expr::None, other) | (other, Expr::None) => other,
            (
                Expr::Coalesce { mut params },
                Expr::Coalesce {
                    params: other_params,
                },
            ) => {
                // If both parameters are coalesce expressions, merge them.
                params.extend(other_params);
                Self::Coalesce { params }
            }
            (Expr::Coalesce { mut params }, other) => {
                // If the first parameter is a coalesce expression, add the second
                // parameter to it.
                params.push(other);
                Self::Coalesce { params }
            }
            (
                param,
                Expr::Coalesce {
                    params: other_params,
                },
            ) => {
                // If the second parameter is a coalesce expression, add the first
                // parameter to it.
                let mut params = vec![param];
                params.extend(other_params);
                Self::Coalesce { params }
            }
            (first, second) => {
                // If neither parameter is a coalesce expression, create a new one.
                Self::Coalesce {
                    params: vec![first, second],
                }
            }
        }
    }

    #[must_use]
    pub(crate) fn min(self, other: Expr) -> Self {
        match (self, other) {
            (Expr::None, expr) | (expr, Expr::None) => expr,
            (
                Expr::Min { mut params },
                Expr::Min {
                    params: other_params,
                },
            ) => {
                // If both parameters are min expressions, merge them.
                params.extend(other_params);
                Self::Min { params }
            }
            (Expr::Min { mut params }, other) | (other, Expr::Min { mut params }) => {
                // If one parameter is a min expression, add the other parameter
                // to it.
                params.push(other);
                Self::Min { params }
            }
            (first, second) => {
                // If neither parameter is a min expression, create a new one.
                Self::Min {
                    params: vec![first, second],
                }
            }
        }
    }

    #[must_use]
    pub(crate) fn max(self, other: Expr) -> Self {
        match (self, other) {
            (Expr::None, expr) | (expr, Expr::None) => expr,
            (
                Expr::Max { mut params },
                Expr::Max {
                    params: other_params,
                },
            ) => {
                // If both parameters are max expressions, merge them.
                params.extend(other_params);
                Self::Max { params }
            }
            (Expr::Max { mut params }, other) | (other, Expr::Max { mut params }) => {
                // If one parameter is a max expression, add the other parameter
                // to it.
                params.push(other);
                Self::Max { params }
            }
            (first, second) => {
                // If neither parameter is a max expression, create a new one.
                Self::Max {
                    params: vec![first, second],
                }
            }
        }
    }
}

impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.generate_string(false))
    }
}

/// Specifies how to add brackets when generating a string representation of an
/// expression.
///
/// - `First`: Add brackets around the first component.
/// - `Rest`: Add brackets around all components except the first.
/// - `All`: Add brackets around all components.
/// - `None`: Do not add brackets.
#[derive(PartialEq)]
enum BracketComponents {
    First,
    Rest,
    All,
    None,
}

/// Display helpers for `FormulaExpression`.
impl Expr {
    /// Generates a string representation of the expression.
    ///
    /// If `bracket_whole` is `true` and the expression has more than one
    /// component, the whole expression is enclosed in brackets.
    fn generate_string(&self, bracket_whole: bool) -> String {
        match self {
            Self::None => String::from("None"),
            Self::Neg { param } => format!("-{}", param.generate_string(true)),
            Self::Number { value } => {
                if value.fract() == 0.0 {
                    // For whole numbers, format with one decimal place.
                    format!("{value:.1}")
                } else {
                    // else format normally.
                    format!("{value}")
                }
            }
            Self::Component { component_id } => format!("#{component_id}"),
            Self::Add { params } => {
                Self::join_params(params, " + ", None, BracketComponents::None, bracket_whole)
            }
            Self::Sub { params } => {
                Self::join_params(params, " - ", None, BracketComponents::Rest, bracket_whole)
            }
            Self::Coalesce { params } => Self::join_params(
                params,
                ", ",
                Some("COALESCE"),
                BracketComponents::None,
                false,
            ),
            Self::Min { params } => {
                Self::join_params(params, ", ", Some("MIN"), BracketComponents::None, false)
            }
            Self::Max { params } => {
                Self::join_params(params, ", ", Some("MAX"), BracketComponents::None, false)
            }
        }
    }

    /// Joins a list of expressions into a string, with the specified separator.
    ///
    /// It also takes an optional prefix, and specifics on how to bracket the
    /// components and the whole expression.
    fn join_params(
        params: &[Expr],
        separator: &str,
        prefix: Option<&str>,
        bracket_components: BracketComponents,
        bracket_whole: bool,
    ) -> String {
        let (mut result, suffix) = match prefix {
            Some(prefix) => (format!("{prefix}("), String::from(")")),
            None => (String::new(), String::new()),
        };
        let mut num_components = 0;
        for expression in params.iter() {
            if num_components > 0 {
                result.push_str(separator);
            }
            if (bracket_components == BracketComponents::First && num_components == 0)
                || (bracket_components == BracketComponents::Rest && num_components > 0)
                || (bracket_components == BracketComponents::All)
            {
                result.push_str(&expression.generate_string(true));
            } else {
                result.push_str(&expression.generate_string(false));
            }
            num_components += 1;
        }
        if bracket_whole && num_components > 1 {
            String::from("(") + &result + &suffix + ")"
        } else {
            result + &suffix
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Expr;

    #[track_caller]
    fn assert_expr(exprs: &[Expr], expected: &str) {
        for expr in exprs {
            assert_eq!(expr.to_string(), expected);
        }
    }

    #[test]
    fn test_arithmatic() {
        let comp = Expr::component;

        assert_expr(
            &[
                comp(10) + comp(11) + comp(12) + comp(13),
                comp(10) - -comp(11) + (comp(12) + comp(13)),
                (comp(10) + comp(11)) - -(comp(12) - -comp(13)),
            ],
            "#10 + #11 + #12 + #13",
        );

        assert_expr(
            &[
                -(comp(10) + comp(11) + comp(12)),
                -comp(10) - comp(11) - comp(12),
                -comp(10) - (comp(11) + comp(12)),
                -(comp(10) + comp(11)) - comp(12),
            ],
            "-(#10 + #11 + #12)",
        );

        assert_expr(
            &[
                comp(11) - comp(10),
                comp(11) + -comp(10),
                -comp(10) + comp(11),
                -comp(10) - -comp(11),
            ],
            "#11 - #10",
        );

        assert_expr(
            &[
                (comp(11) + comp(12)) - comp(10),
                (comp(11) + comp(12)) + -comp(10),
                -comp(10) + (comp(11) + comp(12)),
                -comp(10) - -(comp(11) + comp(12)),
            ],
            "#11 + #12 - #10",
        );

        assert_expr(
            &[
                (comp(11) - comp(12)) - comp(10),
                (comp(11) - comp(12)) + -comp(10),
                -comp(10) + (comp(11) - comp(12)),
                -comp(10) - -(comp(11) - comp(12)),
            ],
            "#11 - #12 - #10",
        );

        assert_expr(
            &[
                comp(11) - comp(12) + comp(10),
                (comp(11) - comp(12)) - -comp(10),
                (comp(11) - comp(12)) + comp(10),
                -(comp(12) - comp(11)) + comp(10),
            ],
            "#11 - #12 + #10",
        );

        assert_expr(
            &[
                (comp(11) + comp(12)) - (comp(10) + comp(13)),
                (comp(11) + comp(12)) + -(comp(10) + comp(13)),
                -(comp(10) + comp(13)) + (comp(11) + comp(12)),
                -(comp(10) + comp(13)) - -(comp(11) + comp(12)),
            ],
            "#11 + #12 - (#10 + #13)",
        );

        assert_expr(
            &[
                (comp(11) - comp(12)) - (comp(10) + comp(13)),
                (comp(11) - comp(12)) + -(comp(10) + comp(13)),
                -(comp(10) + comp(13)) + (comp(11) - comp(12)),
                -(comp(10) + comp(13)) - -(comp(11) - comp(12)),
            ],
            "#11 - #12 - (#10 + #13)",
        );

        assert_expr(
            &[(comp(11) + comp(12)) - (comp(10) - comp(13))],
            "#11 + #12 - (#10 - #13)",
        );
        assert_expr(
            &[(comp(11) + comp(12)) + -(comp(10) - comp(13))],
            "#11 + #12 + #13 - #10",
        );
        assert_expr(
            &[
                -(comp(10) - comp(13)) + (comp(11) + comp(12)),
                -(comp(10) - comp(13)) - -(comp(11) + comp(12)),
            ],
            "#13 - #10 + #11 + #12",
        );
    }

    #[test]
    fn test_functions() {
        let comp = Expr::component;
        let coalesce = Expr::coalesce;
        let number = Expr::number;

        assert_expr(
            &[
                comp(1) - (coalesce(comp(5), comp(7) + comp(6)) + coalesce(comp(2), comp(3)))
                    + coalesce(
                        number(0.0).max(comp(5)),
                        number(0.0).max(comp(7)) + number(0.0).max(comp(6)),
                    ),
            ],
            concat!(
                "#1 - (COALESCE(#5, #7 + #6) + COALESCE(#2, #3)) + ",
                "COALESCE(MAX(0.0, #5), MAX(0.0, #7) + MAX(0.0, #6))"
            ),
        );

        assert_expr(
            &[number(0.0).min(comp(5)).min(comp(7) + comp(6))
                - coalesce(comp(5), comp(7) + comp(6))
                    .max(comp(7))
                    .max(number(22.44))],
            "MIN(0.0, #5, #7 + #6) - MAX(COALESCE(#5, #7 + #6), #7, 22.44)",
        )
    }
}