joule_ast/

energy.rs

//! Energy budget AST representation for Joule
//!
//! This module defines the AST structures for energy budget attributes.
//! Energy budgets allow developers to specify constraints on function energy consumption.

use joule_common::Span;

/// Energy budget constraints for a function or block.
///
/// These constraints can be attached to functions via the `#[energy_budget(...)]` attribute
/// to specify limits on energy consumption, power usage, and temperature impact.
///
/// # Example
/// ```joule
/// #[energy_budget(max_joules = 0.5)]
/// fn process(data: &[f32]) -> f32 {
///     // Compiler can warn if estimated energy > 0.5J
/// }
///
/// #[energy_budget(max_joules = 1.0, max_temp_delta = 5.0)]
/// fn intensive_compute() {
///     // Also track temperature increase
/// }
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct EnergyBudget {
    /// Maximum energy consumption in joules
    pub max_joules: Option<f64>,
    /// Maximum power consumption in watts
    pub max_watts: Option<f64>,
    /// Maximum temperature delta in Celsius
    pub max_temp_delta: Option<f64>,
    /// Source span for error reporting
    pub span: Span,
}

impl EnergyBudget {
    /// Create a new empty energy budget
    pub const fn new(span: Span) -> Self {
        Self {
            max_joules: None,
            max_watts: None,
            max_temp_delta: None,
            span,
        }
    }

    /// Create an energy budget with `max_joules` constraint
    pub const fn with_max_joules(joules: f64, span: Span) -> Self {
        Self {
            max_joules: Some(joules),
            max_watts: None,
            max_temp_delta: None,
            span,
        }
    }

    /// Check if this energy budget has any constraints
    pub const fn has_constraints(&self) -> bool {
        self.max_joules.is_some() || self.max_watts.is_some() || self.max_temp_delta.is_some()
    }

    /// Set `max_joules` constraint
    pub const fn set_max_joules(&mut self, joules: f64) {
        self.max_joules = Some(joules);
    }

    /// Set `max_watts` constraint
    pub const fn set_max_watts(&mut self, watts: f64) {
        self.max_watts = Some(watts);
    }

    /// Set `max_temp_delta` constraint
    pub const fn set_max_temp_delta(&mut self, delta: f64) {
        self.max_temp_delta = Some(delta);
    }
}

impl Default for EnergyBudget {
    fn default() -> Self {
        Self {
            max_joules: None,
            max_watts: None,
            max_temp_delta: None,
            span: Span::dummy(),
        }
    }
}

/// Represents an attribute attached to an item.
///
/// Joule supports Rust-style attributes like `#[energy_budget(max_joules = 0.5)]`.
#[derive(Debug, Clone, PartialEq)]
pub struct Attribute {
    /// The name/path of the attribute (e.g., "`energy_budget`")
    pub name: crate::Path,
    /// The arguments to the attribute
    pub args: AttributeArgs,
    /// Source span
    pub span: Span,
}

/// Arguments to an attribute
#[derive(Debug, Clone, PartialEq, Default)]
pub enum AttributeArgs {
    /// No arguments: `#[attr]`
    #[default]
    Empty,
    /// Parenthesized arguments: `#[attr(key = value, ...)]`
    Parenthesized(Vec<AttributeArg>),
}

/// A single argument in an attribute
#[derive(Debug, Clone, PartialEq)]
pub struct AttributeArg {
    /// The key name (e.g., "`max_joules`")
    pub key: crate::Ident,
    /// The value (e.g., 0.5)
    pub value: AttributeValue,
    /// Source span
    pub span: Span,
}

/// Value types allowed in attributes
#[derive(Debug, Clone, PartialEq)]
pub enum AttributeValue {
    /// A literal float: `0.5`
    Float(f64),
    /// A literal integer: `100`
    Int(u64),
    /// A string literal: `"hello"`
    String(String),
    /// A boolean: `true` or `false`
    Bool(bool),
    /// An identifier: `SomeVariant`
    Ident(crate::Ident),
}

impl AttributeValue {
    /// Try to get the value as f64
    pub const fn as_f64(&self) -> Option<f64> {
        match self {
            Self::Float(f) => Some(*f),
            Self::Int(i) => Some(*i as f64),
            _ => None,
        }
    }

    /// Try to get the value as u64
    pub const fn as_u64(&self) -> Option<u64> {
        match self {
            Self::Int(i) => Some(*i),
            _ => None,
        }
    }

    /// Try to get the value as bool
    pub const fn as_bool(&self) -> Option<bool> {
        match self {
            Self::Bool(b) => Some(*b),
            _ => None,
        }
    }

    /// Try to get the value as string
    pub fn as_string(&self) -> Option<&str> {
        match self {
            Self::String(s) => Some(s),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use joule_common::Symbol;

    fn make_ident(name: &str) -> crate::Ident {
        crate::Ident::new(Symbol::intern(name), Span::dummy())
    }

    #[test]
    fn test_energy_budget_default() {
        let budget = EnergyBudget::default();
        assert!(budget.max_joules.is_none());
        assert!(budget.max_watts.is_none());
        assert!(budget.max_temp_delta.is_none());
        assert!(!budget.has_constraints());
    }

    #[test]
    fn test_energy_budget_with_constraints() {
        let budget = EnergyBudget::with_max_joules(0.5, Span::dummy());
        assert_eq!(budget.max_joules, Some(0.5));
        assert!(budget.has_constraints());
    }

    #[test]
    fn test_energy_budget_setters() {
        let mut budget = EnergyBudget::new(Span::dummy());
        budget.set_max_joules(1.0);
        budget.set_max_watts(50.0);
        budget.set_max_temp_delta(10.0);

        assert_eq!(budget.max_joules, Some(1.0));
        assert_eq!(budget.max_watts, Some(50.0));
        assert_eq!(budget.max_temp_delta, Some(10.0));
    }

    #[test]
    fn test_attribute_value_conversions() {
        let float_val = AttributeValue::Float(0.5);
        assert_eq!(float_val.as_f64(), Some(0.5));

        let int_val = AttributeValue::Int(100);
        assert_eq!(int_val.as_u64(), Some(100));
        assert_eq!(int_val.as_f64(), Some(100.0)); // ints can be floats

        let bool_val = AttributeValue::Bool(true);
        assert_eq!(bool_val.as_bool(), Some(true));

        let string_val = AttributeValue::String("test".to_string());
        assert_eq!(string_val.as_string(), Some("test"));

        let ident_val = AttributeValue::Ident(make_ident("SomeValue"));
        assert!(ident_val.as_f64().is_none());
    }

    #[test]
    fn test_attribute_args_default() {
        let args = AttributeArgs::default();
        assert!(matches!(args, AttributeArgs::Empty));
    }
}