solverforge_core/constraints/

stream.rs

use crate::constraints::{Collector, Joiner, NamedExpression, WasmFunction};
use serde::{Deserialize, Serialize};

/// A component in a constraint stream pipeline.
///
/// Constraint streams transform entity streams through filtering, joining,
/// grouping, and finally scoring. Components chain together to form constraints.
///
/// # Pipeline Structure
///
/// A typical constraint pipeline:
/// 1. **Source**: `for_each("Lesson")` - iterate over entities
/// 2. **Filter**: `filter(predicate)` - exclude non-matching entities
/// 3. **Join**: `join("Timeslot")` - combine with another entity type
/// 4. **Score**: `penalize("1hard")` or `reward("1soft")` - apply scoring
///
/// # Example
///
/// ```
/// use solverforge_core::constraints::StreamComponent;
///
/// // Simple penalty: penalize each lesson by 1 hard point
/// let stream = vec![
///     StreamComponent::for_each("Lesson"),
///     StreamComponent::penalize("1hard"),
/// ];
///
/// // Unique pair constraint: penalize conflicting lessons
/// let conflict = vec![
///     StreamComponent::for_each_unique_pair("Lesson"),
///     StreamComponent::penalize("1hard"),
/// ];
/// ```
///
/// # Score Weights
///
/// Weights specify the penalty/reward magnitude:
/// - `"1hard"` - 1 hard constraint violation
/// - `"10soft"` - 10 soft constraint points
/// - `"1hard/5soft"` - both hard and soft impact
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind")]
pub enum StreamComponent {
    #[serde(rename = "forEach")]
    ForEach {
        #[serde(rename = "className")]
        class_name: String,
    },
    #[serde(rename = "forEachIncludingUnassigned")]
    ForEachIncludingUnassigned {
        #[serde(rename = "className")]
        class_name: String,
    },
    #[serde(rename = "forEachUniquePair")]
    ForEachUniquePair {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "filter")]
    Filter { predicate: WasmFunction },
    #[serde(rename = "join")]
    Join {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "ifExists")]
    IfExists {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "ifNotExists")]
    IfNotExists {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "ifExistsOther")]
    IfExistsOther {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "ifNotExistsOther")]
    IfNotExistsOther {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "ifExistsIncludingUnassigned")]
    IfExistsIncludingUnassigned {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "ifNotExistsIncludingUnassigned")]
    IfNotExistsIncludingUnassigned {
        #[serde(rename = "className")]
        class_name: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        joiners: Vec<Joiner>,
    },
    #[serde(rename = "groupBy")]
    GroupBy {
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        keys: Vec<WasmFunction>,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        aggregators: Vec<Collector>,
    },
    #[serde(rename = "map")]
    Map {
        #[serde(rename = "mapper")]
        mappers: Vec<WasmFunction>,
    },
    #[serde(rename = "flattenLast")]
    FlattenLast {
        #[serde(skip_serializing_if = "Option::is_none")]
        map: Option<WasmFunction>,
    },
    #[serde(rename = "expand")]
    Expand {
        #[serde(rename = "mapper")]
        mappers: Vec<WasmFunction>,
    },
    #[serde(rename = "complement")]
    Complement {
        #[serde(rename = "className")]
        class_name: String,
    },
    #[serde(rename = "penalize")]
    Penalize {
        weight: String,
        #[serde(rename = "scaleBy", skip_serializing_if = "Option::is_none")]
        scale_by: Option<WasmFunction>,
    },
    #[serde(rename = "reward")]
    Reward {
        weight: String,
        #[serde(rename = "scaleBy", skip_serializing_if = "Option::is_none")]
        scale_by: Option<WasmFunction>,
    },
    #[serde(rename = "concat")]
    Concat {
        #[serde(rename = "otherComponents")]
        other_components: Vec<StreamComponent>,
    },
    #[serde(rename = "distinct")]
    Distinct,
    #[serde(rename = "impact")]
    Impact {
        weight: String,
        #[serde(rename = "scaleBy", skip_serializing_if = "Option::is_none")]
        scale_by: Option<WasmFunction>,
    },
    #[serde(rename = "indictWith")]
    IndictWith {
        #[serde(rename = "indictedObjectProvider")]
        indicted_object_provider: WasmFunction,
    },
    #[serde(rename = "justifyWith")]
    JustifyWith {
        #[serde(rename = "justificationSupplier")]
        justification_supplier: WasmFunction,
    },
}

impl StreamComponent {
    /// Iterates over all entities of the given class.
    ///
    /// This is the most common way to start a constraint stream.
    /// Only considers entities with assigned planning variables.
    ///
    /// ```
    /// use solverforge_core::constraints::StreamComponent;
    ///
    /// let source = StreamComponent::for_each("Lesson");
    /// ```
    pub fn for_each(class_name: impl Into<String>) -> Self {
        StreamComponent::ForEach {
            class_name: class_name.into(),
        }
    }

    pub fn for_each_including_unassigned(class_name: impl Into<String>) -> Self {
        StreamComponent::ForEachIncludingUnassigned {
            class_name: class_name.into(),
        }
    }

    /// Iterates over all unique pairs of entities of the given class.
    ///
    /// Use this for constraints that compare two entities of the same type,
    /// like detecting room conflicts between lessons.
    ///
    /// ```
    /// use solverforge_core::constraints::StreamComponent;
    ///
    /// // Pairs: (A,B), (A,C), (B,C) - no duplicates like (B,A)
    /// let pairs = StreamComponent::for_each_unique_pair("Lesson");
    /// ```
    pub fn for_each_unique_pair(class_name: impl Into<String>) -> Self {
        StreamComponent::ForEachUniquePair {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn for_each_unique_pair_with_joiners(
        class_name: impl Into<String>,
        joiners: Vec<Joiner>,
    ) -> Self {
        StreamComponent::ForEachUniquePair {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn filter(predicate: WasmFunction) -> Self {
        StreamComponent::Filter { predicate }
    }

    pub fn join(class_name: impl Into<String>) -> Self {
        StreamComponent::Join {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn join_with_joiners(class_name: impl Into<String>, joiners: Vec<Joiner>) -> Self {
        StreamComponent::Join {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn if_exists(class_name: impl Into<String>) -> Self {
        StreamComponent::IfExists {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn if_exists_with_joiners(class_name: impl Into<String>, joiners: Vec<Joiner>) -> Self {
        StreamComponent::IfExists {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn if_not_exists(class_name: impl Into<String>) -> Self {
        StreamComponent::IfNotExists {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn if_not_exists_with_joiners(class_name: impl Into<String>, joiners: Vec<Joiner>) -> Self {
        StreamComponent::IfNotExists {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn if_exists_other(class_name: impl Into<String>) -> Self {
        StreamComponent::IfExistsOther {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn if_exists_other_with_joiners(
        class_name: impl Into<String>,
        joiners: Vec<Joiner>,
    ) -> Self {
        StreamComponent::IfExistsOther {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn if_not_exists_other(class_name: impl Into<String>) -> Self {
        StreamComponent::IfNotExistsOther {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn if_not_exists_other_with_joiners(
        class_name: impl Into<String>,
        joiners: Vec<Joiner>,
    ) -> Self {
        StreamComponent::IfNotExistsOther {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn if_exists_including_unassigned(class_name: impl Into<String>) -> Self {
        StreamComponent::IfExistsIncludingUnassigned {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn if_exists_including_unassigned_with_joiners(
        class_name: impl Into<String>,
        joiners: Vec<Joiner>,
    ) -> Self {
        StreamComponent::IfExistsIncludingUnassigned {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn if_not_exists_including_unassigned(class_name: impl Into<String>) -> Self {
        StreamComponent::IfNotExistsIncludingUnassigned {
            class_name: class_name.into(),
            joiners: Vec::new(),
        }
    }

    pub fn if_not_exists_including_unassigned_with_joiners(
        class_name: impl Into<String>,
        joiners: Vec<Joiner>,
    ) -> Self {
        StreamComponent::IfNotExistsIncludingUnassigned {
            class_name: class_name.into(),
            joiners,
        }
    }

    pub fn group_by(keys: Vec<WasmFunction>, aggregators: Vec<Collector>) -> Self {
        StreamComponent::GroupBy { keys, aggregators }
    }

    pub fn group_by_key(key: WasmFunction) -> Self {
        StreamComponent::GroupBy {
            keys: vec![key],
            aggregators: Vec::new(),
        }
    }

    pub fn group_by_collector(aggregator: Collector) -> Self {
        StreamComponent::GroupBy {
            keys: Vec::new(),
            aggregators: vec![aggregator],
        }
    }

    pub fn map(mappers: Vec<WasmFunction>) -> Self {
        StreamComponent::Map { mappers }
    }

    pub fn map_single(mapper: WasmFunction) -> Self {
        StreamComponent::Map {
            mappers: vec![mapper],
        }
    }

    pub fn flatten_last() -> Self {
        StreamComponent::FlattenLast { map: None }
    }

    pub fn flatten_last_with_map(map: WasmFunction) -> Self {
        StreamComponent::FlattenLast { map: Some(map) }
    }

    pub fn expand(mappers: Vec<WasmFunction>) -> Self {
        StreamComponent::Expand { mappers }
    }

    pub fn complement(class_name: impl Into<String>) -> Self {
        StreamComponent::Complement {
            class_name: class_name.into(),
        }
    }

    /// Penalizes matching entities by a fixed weight.
    ///
    /// The weight reduces the solution score. Higher penalties are worse.
    ///
    /// ```
    /// use solverforge_core::constraints::StreamComponent;
    ///
    /// // Hard constraint: 1 point per violation
    /// let hard = StreamComponent::penalize("1hard");
    ///
    /// // Soft constraint: 100 points per violation
    /// let soft = StreamComponent::penalize("100soft");
    /// ```
    pub fn penalize(weight: impl Into<String>) -> Self {
        StreamComponent::Penalize {
            weight: weight.into(),
            scale_by: None,
        }
    }

    pub fn penalize_with_weigher(weight: impl Into<String>, scale_by: WasmFunction) -> Self {
        StreamComponent::Penalize {
            weight: weight.into(),
            scale_by: Some(scale_by),
        }
    }

    /// Rewards matching entities by a fixed weight.
    ///
    /// The weight increases the solution score. Higher rewards are better.
    ///
    /// ```
    /// use solverforge_core::constraints::StreamComponent;
    ///
    /// // Reward preferred assignments
    /// let bonus = StreamComponent::reward("10soft");
    /// ```
    pub fn reward(weight: impl Into<String>) -> Self {
        StreamComponent::Reward {
            weight: weight.into(),
            scale_by: None,
        }
    }

    pub fn reward_with_weigher(weight: impl Into<String>, scale_by: WasmFunction) -> Self {
        StreamComponent::Reward {
            weight: weight.into(),
            scale_by: Some(scale_by),
        }
    }

    pub fn concat(other_components: Vec<StreamComponent>) -> Self {
        StreamComponent::Concat { other_components }
    }

    pub fn distinct() -> Self {
        StreamComponent::Distinct
    }

    pub fn impact(weight: impl Into<String>) -> Self {
        StreamComponent::Impact {
            weight: weight.into(),
            scale_by: None,
        }
    }

    pub fn impact_with_weigher(weight: impl Into<String>, scale_by: WasmFunction) -> Self {
        StreamComponent::Impact {
            weight: weight.into(),
            scale_by: Some(scale_by),
        }
    }

    pub fn indict_with(indicted_object_provider: WasmFunction) -> Self {
        StreamComponent::IndictWith {
            indicted_object_provider,
        }
    }

    pub fn indict_with_expr(indicted_object_provider: NamedExpression) -> Self {
        StreamComponent::IndictWith {
            indicted_object_provider: indicted_object_provider.into(),
        }
    }

    pub fn justify_with(justification_supplier: WasmFunction) -> Self {
        StreamComponent::JustifyWith {
            justification_supplier,
        }
    }

    pub fn justify_with_expr(justification_supplier: NamedExpression) -> Self {
        StreamComponent::JustifyWith {
            justification_supplier: justification_supplier.into(),
        }
    }

    // ===== Expression-based convenience methods =====

    /// Creates a filter component from a named expression.
    ///
    /// # Example
    ///
    /// ```
    /// use solverforge_core::wasm::{Expr, FieldAccessExt};
    /// use solverforge_core::constraints::{StreamComponent, IntoNamedExpression};
    ///
    /// let has_room = Expr::is_not_null(Expr::param(0).get("Lesson", "room"))
    ///     .named_as("has_room");
    /// let filter = StreamComponent::filter_expr(has_room);
    /// assert!(matches!(filter, StreamComponent::Filter { .. }));
    /// ```
    pub fn filter_expr(expr: NamedExpression) -> Self {
        StreamComponent::Filter {
            predicate: expr.into(),
        }
    }

    /// Creates a map component from named expressions.
    pub fn map_expr(mappers: Vec<NamedExpression>) -> Self {
        StreamComponent::Map {
            mappers: mappers.into_iter().map(|e| e.into()).collect(),
        }
    }

    /// Creates a single-mapper map component from a named expression.
    pub fn map_single_expr(mapper: NamedExpression) -> Self {
        StreamComponent::Map {
            mappers: vec![mapper.into()],
        }
    }

    /// Creates a groupBy component with expression-based key extractors.
    pub fn group_by_expr(keys: Vec<NamedExpression>, aggregators: Vec<Collector>) -> Self {
        StreamComponent::GroupBy {
            keys: keys.into_iter().map(|e| e.into()).collect(),
            aggregators,
        }
    }

    /// Creates a groupBy component with a single expression-based key.
    pub fn group_by_key_expr(key: NamedExpression) -> Self {
        StreamComponent::GroupBy {
            keys: vec![key.into()],
            aggregators: Vec::new(),
        }
    }

    /// Creates a penalize component with an expression-based weigher.
    pub fn penalize_with_expr(weight: impl Into<String>, scale_by: NamedExpression) -> Self {
        StreamComponent::Penalize {
            weight: weight.into(),
            scale_by: Some(scale_by.into()),
        }
    }

    /// Creates a reward component with an expression-based weigher.
    pub fn reward_with_expr(weight: impl Into<String>, scale_by: NamedExpression) -> Self {
        StreamComponent::Reward {
            weight: weight.into(),
            scale_by: Some(scale_by.into()),
        }
    }

    /// Creates a flattenLast component with an expression-based mapper.
    pub fn flatten_last_with_expr(map: NamedExpression) -> Self {
        StreamComponent::FlattenLast {
            map: Some(map.into()),
        }
    }

    /// Creates an expand component from named expressions.
    pub fn expand_expr(mappers: Vec<NamedExpression>) -> Self {
        StreamComponent::Expand {
            mappers: mappers.into_iter().map(|e| e.into()).collect(),
        }
    }

    /// Creates an impact component with an expression-based weigher.
    pub fn impact_with_expr(weight: impl Into<String>, scale_by: NamedExpression) -> Self {
        StreamComponent::Impact {
            weight: weight.into(),
            scale_by: Some(scale_by.into()),
        }
    }
}

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

    #[test]
    fn test_for_each() {
        let component = StreamComponent::for_each("Lesson");
        match component {
            StreamComponent::ForEach { class_name } => {
                assert_eq!(class_name, "Lesson");
            }
            _ => panic!("Expected ForEach"),
        }
    }

    #[test]
    fn test_for_each_including_unassigned() {
        let component = StreamComponent::for_each_including_unassigned("Lesson");
        match component {
            StreamComponent::ForEachIncludingUnassigned { class_name } => {
                assert_eq!(class_name, "Lesson");
            }
            _ => panic!("Expected ForEachIncludingUnassigned"),
        }
    }

    #[test]
    fn test_for_each_unique_pair() {
        let component = StreamComponent::for_each_unique_pair("Lesson");
        match component {
            StreamComponent::ForEachUniquePair {
                class_name,
                joiners,
            } => {
                assert_eq!(class_name, "Lesson");
                assert!(joiners.is_empty());
            }
            _ => panic!("Expected ForEachUniquePair"),
        }
    }

    #[test]
    fn test_for_each_unique_pair_with_joiners() {
        let component = StreamComponent::for_each_unique_pair_with_joiners(
            "Lesson",
            vec![Joiner::equal(WasmFunction::new("get_timeslot"))],
        );
        match component {
            StreamComponent::ForEachUniquePair { joiners, .. } => {
                assert_eq!(joiners.len(), 1);
            }
            _ => panic!("Expected ForEachUniquePair"),
        }
    }

    #[test]
    fn test_filter() {
        let component = StreamComponent::filter(WasmFunction::new("is_valid"));
        match component {
            StreamComponent::Filter { predicate } => {
                assert_eq!(predicate.name(), "is_valid");
            }
            _ => panic!("Expected Filter"),
        }
    }

    #[test]
    fn test_join() {
        let component = StreamComponent::join("Room");
        match component {
            StreamComponent::Join {
                class_name,
                joiners,
            } => {
                assert_eq!(class_name, "Room");
                assert!(joiners.is_empty());
            }
            _ => panic!("Expected Join"),
        }
    }

    #[test]
    fn test_join_with_joiners() {
        let component = StreamComponent::join_with_joiners(
            "Room",
            vec![Joiner::equal(WasmFunction::new("get_room"))],
        );
        match component {
            StreamComponent::Join { joiners, .. } => {
                assert_eq!(joiners.len(), 1);
            }
            _ => panic!("Expected Join"),
        }
    }

    #[test]
    fn test_if_exists() {
        let component = StreamComponent::if_exists("Conflict");
        match component {
            StreamComponent::IfExists { class_name, .. } => {
                assert_eq!(class_name, "Conflict");
            }
            _ => panic!("Expected IfExists"),
        }
    }

    #[test]
    fn test_if_not_exists() {
        let component = StreamComponent::if_not_exists("Conflict");
        match component {
            StreamComponent::IfNotExists { class_name, .. } => {
                assert_eq!(class_name, "Conflict");
            }
            _ => panic!("Expected IfNotExists"),
        }
    }

    #[test]
    fn test_group_by() {
        let component = StreamComponent::group_by(
            vec![WasmFunction::new("get_room")],
            vec![Collector::count()],
        );
        match component {
            StreamComponent::GroupBy { keys, aggregators } => {
                assert_eq!(keys.len(), 1);
                assert_eq!(aggregators.len(), 1);
            }
            _ => panic!("Expected GroupBy"),
        }
    }

    #[test]
    fn test_group_by_key() {
        let component = StreamComponent::group_by_key(WasmFunction::new("get_room"));
        match component {
            StreamComponent::GroupBy { keys, aggregators } => {
                assert_eq!(keys.len(), 1);
                assert!(aggregators.is_empty());
            }
            _ => panic!("Expected GroupBy"),
        }
    }

    #[test]
    fn test_group_by_collector() {
        let component = StreamComponent::group_by_collector(Collector::count());
        match component {
            StreamComponent::GroupBy { keys, aggregators } => {
                assert!(keys.is_empty());
                assert_eq!(aggregators.len(), 1);
            }
            _ => panic!("Expected GroupBy"),
        }
    }

    #[test]
    fn test_map() {
        let component =
            StreamComponent::map(vec![WasmFunction::new("get_a"), WasmFunction::new("get_b")]);
        match component {
            StreamComponent::Map { mappers } => {
                assert_eq!(mappers.len(), 2);
            }
            _ => panic!("Expected Map"),
        }
    }

    #[test]
    fn test_map_single() {
        let component = StreamComponent::map_single(WasmFunction::new("get_value"));
        match component {
            StreamComponent::Map { mappers } => {
                assert_eq!(mappers.len(), 1);
            }
            _ => panic!("Expected Map"),
        }
    }

    #[test]
    fn test_flatten_last() {
        let component = StreamComponent::flatten_last();
        match component {
            StreamComponent::FlattenLast { map } => {
                assert!(map.is_none());
            }
            _ => panic!("Expected FlattenLast"),
        }
    }

    #[test]
    fn test_flatten_last_with_map() {
        let component = StreamComponent::flatten_last_with_map(WasmFunction::new("get_items"));
        match component {
            StreamComponent::FlattenLast { map } => {
                assert!(map.is_some());
            }
            _ => panic!("Expected FlattenLast"),
        }
    }

    #[test]
    fn test_expand() {
        let component = StreamComponent::expand(vec![WasmFunction::new("get_extra")]);
        match component {
            StreamComponent::Expand { mappers } => {
                assert_eq!(mappers.len(), 1);
            }
            _ => panic!("Expected Expand"),
        }
    }

    #[test]
    fn test_complement() {
        let component = StreamComponent::complement("Timeslot");
        match component {
            StreamComponent::Complement { class_name } => {
                assert_eq!(class_name, "Timeslot");
            }
            _ => panic!("Expected Complement"),
        }
    }

    #[test]
    fn test_penalize() {
        let component = StreamComponent::penalize("1hard");
        match component {
            StreamComponent::Penalize { weight, scale_by } => {
                assert_eq!(weight, "1hard");
                assert!(scale_by.is_none());
            }
            _ => panic!("Expected Penalize"),
        }
    }

    #[test]
    fn test_penalize_with_weigher() {
        let component =
            StreamComponent::penalize_with_weigher("1hard", WasmFunction::new("get_weight"));
        match component {
            StreamComponent::Penalize { weight, scale_by } => {
                assert_eq!(weight, "1hard");
                assert!(scale_by.is_some());
            }
            _ => panic!("Expected Penalize"),
        }
    }

    #[test]
    fn test_reward() {
        let component = StreamComponent::reward("1soft");
        match component {
            StreamComponent::Reward { weight, scale_by } => {
                assert_eq!(weight, "1soft");
                assert!(scale_by.is_none());
            }
            _ => panic!("Expected Reward"),
        }
    }

    #[test]
    fn test_reward_with_weigher() {
        let component =
            StreamComponent::reward_with_weigher("1soft", WasmFunction::new("get_bonus"));
        match component {
            StreamComponent::Reward { scale_by, .. } => {
                assert!(scale_by.is_some());
            }
            _ => panic!("Expected Reward"),
        }
    }

    #[test]
    fn test_for_each_json_serialization() {
        let component = StreamComponent::for_each("Lesson");
        let json = serde_json::to_string(&component).unwrap();
        assert!(json.contains("\"kind\":\"forEach\""));
        assert!(json.contains("\"className\":\"Lesson\""));

        let parsed: StreamComponent = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, component);
    }

    #[test]
    fn test_filter_json_serialization() {
        let component = StreamComponent::filter(WasmFunction::new("is_valid"));
        let json = serde_json::to_string(&component).unwrap();
        assert!(json.contains("\"kind\":\"filter\""));
        assert!(json.contains("\"predicate\":\"is_valid\""));

        let parsed: StreamComponent = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, component);
    }

    #[test]
    fn test_join_json_serialization() {
        let component = StreamComponent::join_with_joiners(
            "Room",
            vec![Joiner::equal(WasmFunction::new("get_room"))],
        );
        let json = serde_json::to_string(&component).unwrap();
        assert!(json.contains("\"kind\":\"join\""));
        assert!(json.contains("\"className\":\"Room\""));
        assert!(json.contains("\"joiners\""));

        let parsed: StreamComponent = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, component);
    }

    #[test]
    fn test_group_by_json_serialization() {
        let component = StreamComponent::group_by(
            vec![WasmFunction::new("get_room")],
            vec![Collector::count()],
        );
        let json = serde_json::to_string(&component).unwrap();
        assert!(json.contains("\"kind\":\"groupBy\""));
        assert!(json.contains("\"keys\""));
        assert!(json.contains("\"aggregators\""));

        let parsed: StreamComponent = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, component);
    }

    #[test]
    fn test_penalize_json_serialization() {
        let component = StreamComponent::penalize("1hard");
        let json = serde_json::to_string(&component).unwrap();
        assert!(json.contains("\"kind\":\"penalize\""));
        assert!(json.contains("\"weight\":\"1hard\""));

        let parsed: StreamComponent = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed, component);
    }

    #[test]
    fn test_component_clone() {
        let component = StreamComponent::for_each("Lesson");
        let cloned = component.clone();
        assert_eq!(component, cloned);
    }

    #[test]
    fn test_component_debug() {
        let component = StreamComponent::for_each("Lesson");
        let debug = format!("{:?}", component);
        assert!(debug.contains("ForEach"));
    }

    // ===== Expression-based method tests =====

    #[test]
    fn test_filter_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::{Expr, FieldAccessExt};

        let has_room = Expr::is_not_null(Expr::param(0).get("Lesson", "room")).named_as("has_room");
        let component = StreamComponent::filter_expr(has_room);

        match component {
            StreamComponent::Filter { predicate } => {
                assert_eq!(predicate.name(), "has_room");
            }
            _ => panic!("Expected Filter"),
        }
    }

    #[test]
    fn test_map_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::{Expr, FieldAccessExt};

        let get_room = Expr::param(0).get("Lesson", "room").named_as("get_room");
        let get_timeslot = Expr::param(0)
            .get("Lesson", "timeslot")
            .named_as("get_timeslot");
        let component = StreamComponent::map_expr(vec![get_room, get_timeslot]);

        match component {
            StreamComponent::Map { mappers } => {
                assert_eq!(mappers.len(), 2);
                assert_eq!(mappers[0].name(), "get_room");
                assert_eq!(mappers[1].name(), "get_timeslot");
            }
            _ => panic!("Expected Map"),
        }
    }

    #[test]
    fn test_map_single_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::{Expr, FieldAccessExt};

        let get_room = Expr::param(0).get("Lesson", "room").named_as("get_room");
        let component = StreamComponent::map_single_expr(get_room);

        match component {
            StreamComponent::Map { mappers } => {
                assert_eq!(mappers.len(), 1);
                assert_eq!(mappers[0].name(), "get_room");
            }
            _ => panic!("Expected Map"),
        }
    }

    #[test]
    fn test_group_by_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::{Expr, FieldAccessExt};

        let get_room = Expr::param(0).get("Lesson", "room").named_as("get_room");
        let component = StreamComponent::group_by_expr(vec![get_room], vec![Collector::count()]);

        match component {
            StreamComponent::GroupBy { keys, aggregators } => {
                assert_eq!(keys.len(), 1);
                assert_eq!(keys[0].name(), "get_room");
                assert_eq!(aggregators.len(), 1);
            }
            _ => panic!("Expected GroupBy"),
        }
    }

    #[test]
    fn test_group_by_key_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::{Expr, FieldAccessExt};

        let get_room = Expr::param(0).get("Lesson", "room").named_as("get_room");
        let component = StreamComponent::group_by_key_expr(get_room);

        match component {
            StreamComponent::GroupBy { keys, aggregators } => {
                assert_eq!(keys.len(), 1);
                assert!(aggregators.is_empty());
            }
            _ => panic!("Expected GroupBy"),
        }
    }

    #[test]
    fn test_penalize_with_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::Expr;

        let weight_fn = Expr::int(1).named_as("weight");
        let component = StreamComponent::penalize_with_expr("1hard", weight_fn);

        match component {
            StreamComponent::Penalize { weight, scale_by } => {
                assert_eq!(weight, "1hard");
                assert!(scale_by.is_some());
                assert_eq!(scale_by.unwrap().name(), "weight");
            }
            _ => panic!("Expected Penalize"),
        }
    }

    #[test]
    fn test_reward_with_expr() {
        use crate::constraints::IntoNamedExpression;
        use crate::wasm::Expr;

        let bonus = Expr::int(10).named_as("bonus");
        let component = StreamComponent::reward_with_expr("1soft", bonus);

        match component {
            StreamComponent::Reward { weight, scale_by } => {
                assert_eq!(weight, "1soft");
                assert!(scale_by.is_some());
            }
            _ => panic!("Expected Reward"),
        }
    }
}