rustytorch/crates/core/rtx-graph/src/operations.rs

// Operation types for unified data+compute graph
// Phase 6: Self-Optimizing Platform

use crate::{GraphError, NodeId, Result};
use serde::{Deserialize, Serialize};

/// Unified operation type that can represent both data processing and compute operations
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum OperationType {
    /// Data processing operation (ETL-style)
    Data(DataOp),
    /// Compute operation (ML-style)
    Compute(ComputeOp),
}

/// Data processing operations for ETL-style workloads
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum DataOp {
    /// Read data from a source
    Read {
        /// Data source path or URI
        source: String,
        /// Data format (csv, parquet, json, etc.)
        format: String,
        /// Optional schema specification
        schema: Option<String>,
    },

    /// Transform data using specified operation
    Transform {
        /// Type of transformation (filter, map, aggregate, etc.)
        operation: String,
        /// Transformation predicate or expression
        predicate: String,
    },

    /// Write data to a destination
    Write {
        /// Destination path or URI
        destination: String,
        /// Output format
        format: String,
        /// Write options
        options: Option<String>,
    },

    /// Join two data sources
    Join {
        /// Join type (inner, left, right, outer)
        join_type: String,
        /// Join condition
        condition: String,
    },

    /// Aggregate data
    Aggregate {
        /// Grouping columns
        group_by: Vec<String>,
        /// Aggregation functions
        aggregates: Vec<String>,
    },

    /// Sort data
    Sort {
        /// Columns to sort by
        columns: Vec<String>,
        /// Sort directions (asc/desc for each column)
        directions: Vec<String>,
    },
}

/// Compute operations for ML-style workloads
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ComputeOp {
    /// Load data into tensor format
    TensorLoad {
        /// Source data node
        data_node: NodeId,
        /// Tensor shape
        shape: Vec<usize>,
        /// Data type (f32, f16, i32, etc.)
        dtype: String,
    },

    /// Element-wise addition
    Add {
        /// Left-hand side tensor
        lhs: NodeId,
        /// Right-hand side tensor
        rhs: NodeId,
    },

    /// Element-wise multiplication
    Mul {
        /// Left-hand side tensor
        lhs: NodeId,
        /// Right-hand side tensor
        rhs: NodeId,
    },

    /// Matrix multiplication
    MatMul {
        /// Left-hand side matrix
        lhs: NodeId,
        /// Right-hand side matrix
        rhs: NodeId,
    },

    /// Convolution operation
    Conv2d {
        /// Input tensor
        input: NodeId,
        /// Weight tensor
        weight: NodeId,
        /// Optional bias tensor
        bias: Option<NodeId>,
        /// Stride configuration
        stride: Vec<usize>,
        /// Padding configuration
        padding: Vec<usize>,
        /// Dilation configuration
        dilation: Vec<usize>,
        /// Number of groups for grouped convolution
        groups: usize,
    },

    /// Activation function
    Activation {
        /// Input tensor
        input: NodeId,
        /// Activation type (relu, tanh, sigmoid, etc.)
        activation_type: String,
    },

    /// Pooling operation
    Pool2d {
        /// Input tensor
        input: NodeId,
        /// Pooling type (max, avg)
        pool_type: String,
        /// Kernel size
        kernel_size: Vec<usize>,
        /// Stride
        stride: Vec<usize>,
    },

    /// Batch normalization
    BatchNorm {
        /// Input tensor
        input: NodeId,
        /// Running mean
        running_mean: NodeId,
        /// Running variance
        running_var: NodeId,
        /// Optional weight
        weight: Option<NodeId>,
        /// Optional bias
        bias: Option<NodeId>,
        /// Training mode
        training: bool,
    },

    /// Reshape tensor
    Reshape {
        /// Input tensor
        input: NodeId,
        /// New shape
        shape: Vec<i64>, // -1 for inferred dimensions
    },

    /// Transpose tensor
    Transpose {
        /// Input tensor
        input: NodeId,
        /// Dimension permutation
        dims: Vec<usize>,
    },
}

impl DataOp {
    /// Validate data operation configuration
    pub fn validate(&self) -> Result<()> {
        match self {
            Self::Read { source, format, .. } => {
                if source.is_empty() {
                    return Err(GraphError::invalid_operation("Read source cannot be empty"));
                }
                if format.is_empty() {
                    return Err(GraphError::invalid_operation("Read format cannot be empty"));
                }
                if !Self::is_valid_format(format) {
                    return Err(GraphError::data_format_error(format!(
                        "Unsupported format: {format}"
                    )));
                }
            }
            Self::Transform {
                operation,
                predicate,
            } => {
                if operation.is_empty() {
                    return Err(GraphError::invalid_operation(
                        "Transform operation cannot be empty",
                    ));
                }
                if predicate.is_empty() {
                    return Err(GraphError::invalid_operation(
                        "Transform predicate cannot be empty",
                    ));
                }
            }
            Self::Write {
                destination,
                format,
                ..
            } => {
                if destination.is_empty() {
                    return Err(GraphError::invalid_operation(
                        "Write destination cannot be empty",
                    ));
                }
                if format.is_empty() {
                    return Err(GraphError::invalid_operation(
                        "Write format cannot be empty",
                    ));
                }
            }
            Self::Join {
                join_type,
                condition,
            } => {
                if !Self::is_valid_join_type(join_type) {
                    return Err(GraphError::invalid_operation(format!(
                        "Invalid join type: {join_type}"
                    )));
                }
                if condition.is_empty() {
                    return Err(GraphError::invalid_operation(
                        "Join condition cannot be empty",
                    ));
                }
            }
            Self::Aggregate {
                group_by: _,
                aggregates,
            } => {
                if aggregates.is_empty() {
                    return Err(GraphError::invalid_operation("Aggregates cannot be empty"));
                }
            }
            Self::Sort {
                columns,
                directions,
            } => {
                if columns.is_empty() {
                    return Err(GraphError::invalid_operation(
                        "Sort columns cannot be empty",
                    ));
                }
                if columns.len() != directions.len() {
                    return Err(GraphError::invalid_operation(
                        "Sort columns and directions must have same length",
                    ));
                }
            }
        }
        Ok(())
    }

    /// Check if format is supported
    fn is_valid_format(format: &str) -> bool {
        matches!(format, "csv" | "parquet" | "json" | "avro" | "orc")
    }

    /// Check if join type is valid
    fn is_valid_join_type(join_type: &str) -> bool {
        matches!(join_type, "inner" | "left" | "right" | "outer" | "cross")
    }
}

impl ComputeOp {
    /// Validate compute operation configuration
    pub fn validate(&self) -> Result<()> {
        match self {
            Self::TensorLoad { shape, dtype, .. } => {
                if shape.is_empty() {
                    return Err(GraphError::invalid_tensor_shape("Shape cannot be empty"));
                }
                if shape.contains(&0) {
                    return Err(GraphError::invalid_tensor_shape(
                        "Shape dimensions cannot be zero",
                    ));
                }
                if !Self::is_valid_dtype(dtype) {
                    return Err(GraphError::invalid_operation(format!(
                        "Unsupported dtype: {dtype}"
                    )));
                }
            }
            Self::Conv2d {
                stride,
                padding,
                dilation: _,
                groups: _,
                ..
            } => {
                if stride.len() != 2 {
                    return Err(GraphError::invalid_operation("Conv2d stride must be 2D"));
                }
                if padding.len() != 2 {
                    return Err(GraphError::invalid_operation("Conv2d padding must be 2D"));
                }
            }
            Self::Activation {
                activation_type, ..
            } => {
                if !Self::is_valid_activation(activation_type) {
                    return Err(GraphError::invalid_operation(format!(
                        "Unsupported activation: {activation_type}"
                    )));
                }
            }
            Self::Pool2d {
                pool_type,
                kernel_size,
                stride,
                ..
            } => {
                if !Self::is_valid_pool_type(pool_type) {
                    return Err(GraphError::invalid_operation(format!(
                        "Unsupported pool type: {pool_type}"
                    )));
                }
                if kernel_size.len() != 2 {
                    return Err(GraphError::invalid_operation(
                        "Pool2d kernel_size must be 2D",
                    ));
                }
                if stride.len() != 2 {
                    return Err(GraphError::invalid_operation("Pool2d stride must be 2D"));
                }
            }
            Self::Reshape { shape, .. } => {
                let inferred_count = shape.iter().filter(|&&dim| dim == -1).count();
                if inferred_count > 1 {
                    return Err(GraphError::invalid_tensor_shape(
                        "Only one dimension can be inferred (-1)",
                    ));
                }
            }
            Self::Transpose { dims, .. } => {
                // Check for duplicate dimensions
                let mut sorted_dims = dims.clone();
                sorted_dims.sort_unstable();
                for i in 1..sorted_dims.len() {
                    if sorted_dims[i] == sorted_dims[i - 1] {
                        return Err(GraphError::invalid_operation(
                            "Transpose dimensions cannot contain duplicates",
                        ));
                    }
                }
            }
            _ => {} // Other operations have no special validation
        }
        Ok(())
    }

    /// Check if data type is supported
    fn is_valid_dtype(dtype: &str) -> bool {
        matches!(
            dtype,
            "f32" | "f16" | "bf16" | "i32" | "i64" | "i8" | "u8" | "bool"
        )
    }

    /// Check if activation function is supported
    fn is_valid_activation(activation_type: &str) -> bool {
        matches!(
            activation_type,
            "relu" | "tanh" | "sigmoid" | "gelu" | "swish" | "leaky_relu"
        )
    }

    /// Check if pooling type is supported
    fn is_valid_pool_type(pool_type: &str) -> bool {
        matches!(pool_type, "max" | "avg" | "adaptive_max" | "adaptive_avg")
    }
}

impl OperationType {
    /// Validate operation configuration
    pub fn validate(&self) -> Result<()> {
        match self {
            Self::Data(data_op) => data_op.validate(),
            Self::Compute(compute_op) => compute_op.validate(),
        }
    }

    /// Check if this is a data operation
    pub fn is_data_op(&self) -> bool {
        matches!(self, Self::Data(_))
    }

    /// Check if this is a compute operation
    pub fn is_compute_op(&self) -> bool {
        matches!(self, Self::Compute(_))
    }

    /// Get operation name for display/debugging
    pub fn name(&self) -> &'static str {
        match self {
            Self::Data(DataOp::Read { .. }) => "Data::Read",
            Self::Data(DataOp::Transform { .. }) => "Data::Transform",
            Self::Data(DataOp::Write { .. }) => "Data::Write",
            Self::Data(DataOp::Join { .. }) => "Data::Join",
            Self::Data(DataOp::Aggregate { .. }) => "Data::Aggregate",
            Self::Data(DataOp::Sort { .. }) => "Data::Sort",
            Self::Compute(ComputeOp::TensorLoad { .. }) => "Compute::TensorLoad",
            Self::Compute(ComputeOp::Add { .. }) => "Compute::Add",
            Self::Compute(ComputeOp::Mul { .. }) => "Compute::Mul",
            Self::Compute(ComputeOp::MatMul { .. }) => "Compute::MatMul",
            Self::Compute(ComputeOp::Conv2d { .. }) => "Compute::Conv2d",
            Self::Compute(ComputeOp::Activation { .. }) => "Compute::Activation",
            Self::Compute(ComputeOp::Pool2d { .. }) => "Compute::Pool2d",
            Self::Compute(ComputeOp::BatchNorm { .. }) => "Compute::BatchNorm",
            Self::Compute(ComputeOp::Reshape { .. }) => "Compute::Reshape",
            Self::Compute(ComputeOp::Transpose { .. }) => "Compute::Transpose",
        }
    }
}