hakorune/src/mir/join_ir/lowering/condition_lowerer.rs

//! Condition Expression Lowerer
//!
//! This module provides the core logic for lowering AST condition expressions
//! to JoinIR instructions. It handles comparisons, logical operators, and
//! arithmetic expressions.
//!
//! ## Design Philosophy
//!
//! **Single Responsibility**: This module ONLY performs AST → JoinIR lowering.
//! It does NOT:
//! - Manage variable environments (that's condition_env.rs)
//! - Extract variables from AST (that's condition_var_extractor.rs)
//! - Manage HOST ↔ JoinIR bindings (that's inline_boundary.rs)

use crate::ast::{ASTNode, BinaryOperator, LiteralValue, UnaryOperator};
use crate::mir::join_ir::{BinOpKind, CompareOp, ConstValue, JoinInst, MirLikeInst, UnaryOp};
use crate::mir::ValueId;

use super::condition_env::ConditionEnv;

/// Lower an AST condition to JoinIR instructions
///
/// # Arguments
///
/// * `cond_ast` - AST node representing the boolean condition
/// * `alloc_value` - ValueId allocator function
/// * `env` - ConditionEnv for variable resolution (JoinIR-local ValueIds)
///
/// # Returns
///
/// * `Ok((ValueId, Vec<JoinInst>))` - Condition result ValueId and evaluation instructions
/// * `Err(String)` - Lowering error message
///
/// # Supported Patterns
///
/// - Comparisons: `i < n`, `x == y`, `a != b`, `x <= y`, `x >= y`, `x > y`
/// - Logical: `a && b`, `a || b`, `!cond`
/// - Variables and literals
///
/// # Example
///
/// ```ignore
/// let mut env = ConditionEnv::new();
/// env.insert("i".to_string(), ValueId(0));
/// env.insert("end".to_string(), ValueId(1));
///
/// let mut value_counter = 2u32;
/// let mut alloc_value = || {
///     let id = ValueId(value_counter);
///     value_counter += 1;
///     id
/// };
///
/// // Lower condition: i < end
/// let (cond_value, cond_insts) = lower_condition_to_joinir(
///     condition_ast,
///     &mut alloc_value,
///     &env,
/// )?;
/// ```
pub fn lower_condition_to_joinir<F>(
    cond_ast: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
) -> Result<(ValueId, Vec<JoinInst>), String>
where
    F: FnMut() -> ValueId,
{
    let mut instructions = Vec::new();
    let result_value = lower_condition_recursive(cond_ast, alloc_value, env, &mut instructions)?;
    Ok((result_value, instructions))
}

/// Recursive helper for condition lowering
///
/// Handles all supported AST node types and emits appropriate JoinIR instructions.
fn lower_condition_recursive<F>(
    cond_ast: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    match cond_ast {
        // Comparison operations: <, ==, !=, <=, >=, >
        ASTNode::BinaryOp {
            operator,
            left,
            right,
            ..
        } => match operator {
            BinaryOperator::Less
            | BinaryOperator::Equal
            | BinaryOperator::NotEqual
            | BinaryOperator::LessEqual
            | BinaryOperator::GreaterEqual
            | BinaryOperator::Greater => {
                lower_comparison(operator, left, right, alloc_value, env, instructions)
            }
            BinaryOperator::And => {
                lower_logical_and(left, right, alloc_value, env, instructions)
            }
            BinaryOperator::Or => {
                lower_logical_or(left, right, alloc_value, env, instructions)
            }
            _ => Err(format!(
                "Unsupported binary operator in condition: {:?}",
                operator
            )),
        },

        // Unary NOT operator
        ASTNode::UnaryOp {
            operator: UnaryOperator::Not,
            operand,
            ..
        } => lower_not_operator(operand, alloc_value, env, instructions),

        // Variables - resolve from ConditionEnv
        ASTNode::Variable { name, .. } => env
            .get(name)
            .ok_or_else(|| format!("Variable '{}' not bound in ConditionEnv", name)),

        // Literals - emit as constants
        ASTNode::Literal { value, .. } => lower_literal(value, alloc_value, instructions),

        _ => Err(format!("Unsupported AST node in condition: {:?}", cond_ast)),
    }
}

/// Lower a comparison operation (e.g., `i < end`)
fn lower_comparison<F>(
    operator: &BinaryOperator,
    left: &ASTNode,
    right: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    // Lower left and right sides
    let lhs = lower_value_expression(left, alloc_value, env, instructions)?;
    let rhs = lower_value_expression(right, alloc_value, env, instructions)?;
    let dst = alloc_value();

    let cmp_op = match operator {
        BinaryOperator::Less => CompareOp::Lt,
        BinaryOperator::Equal => CompareOp::Eq,
        BinaryOperator::NotEqual => CompareOp::Ne,
        BinaryOperator::LessEqual => CompareOp::Le,
        BinaryOperator::GreaterEqual => CompareOp::Ge,
        BinaryOperator::Greater => CompareOp::Gt,
        _ => unreachable!(),
    };

    // Emit Compare instruction
    instructions.push(JoinInst::Compute(MirLikeInst::Compare {
        dst,
        op: cmp_op,
        lhs,
        rhs,
    }));

    Ok(dst)
}

/// Lower logical AND operation (e.g., `a && b`)
fn lower_logical_and<F>(
    left: &ASTNode,
    right: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    // Logical AND: evaluate both sides and combine
    let lhs = lower_condition_recursive(left, alloc_value, env, instructions)?;
    let rhs = lower_condition_recursive(right, alloc_value, env, instructions)?;
    let dst = alloc_value();

    // Emit BinOp And instruction
    instructions.push(JoinInst::Compute(MirLikeInst::BinOp {
        dst,
        op: BinOpKind::And,
        lhs,
        rhs,
    }));

    Ok(dst)
}

/// Lower logical OR operation (e.g., `a || b`)
fn lower_logical_or<F>(
    left: &ASTNode,
    right: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    // Logical OR: evaluate both sides and combine
    let lhs = lower_condition_recursive(left, alloc_value, env, instructions)?;
    let rhs = lower_condition_recursive(right, alloc_value, env, instructions)?;
    let dst = alloc_value();

    // Emit BinOp Or instruction
    instructions.push(JoinInst::Compute(MirLikeInst::BinOp {
        dst,
        op: BinOpKind::Or,
        lhs,
        rhs,
    }));

    Ok(dst)
}

/// Lower NOT operator (e.g., `!cond`)
fn lower_not_operator<F>(
    operand: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    let operand_val = lower_condition_recursive(operand, alloc_value, env, instructions)?;
    let dst = alloc_value();

    // Emit UnaryOp Not instruction
    instructions.push(JoinInst::Compute(MirLikeInst::UnaryOp {
        dst,
        op: UnaryOp::Not,
        operand: operand_val,
    }));

    Ok(dst)
}

/// Lower a literal value (e.g., `10`, `true`, `"text"`)
fn lower_literal<F>(
    value: &LiteralValue,
    alloc_value: &mut F,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    let dst = alloc_value();
    let const_value = match value {
        LiteralValue::Integer(n) => ConstValue::Integer(*n),
        LiteralValue::String(s) => ConstValue::String(s.clone()),
        LiteralValue::Bool(b) => ConstValue::Bool(*b),
        LiteralValue::Float(_) => {
            return Err("Float literals not supported in JoinIR conditions yet".to_string());
        }
        _ => {
            return Err(format!("Unsupported literal type in condition: {:?}", value));
        }
    };

    instructions.push(JoinInst::Compute(MirLikeInst::Const {
        dst,
        value: const_value,
    }));

    Ok(dst)
}

/// Lower a value expression (for comparison operands, etc.)
///
/// This handles the common case where we need to evaluate a simple value
/// (variable or literal) as part of a comparison.
pub fn lower_value_expression<F>(
    expr: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    match expr {
        // Variables - look up in ConditionEnv
        ASTNode::Variable { name, .. } => env
            .get(name)
            .ok_or_else(|| format!("Variable '{}' not bound in ConditionEnv", name)),

        // Literals - emit as constants
        ASTNode::Literal { value, .. } => lower_literal(value, alloc_value, instructions),

        // Binary operations (for arithmetic in conditions like i + 1 < n)
        ASTNode::BinaryOp {
            operator,
            left,
            right,
            ..
        } => lower_arithmetic_binop(operator, left, right, alloc_value, env, instructions),

        _ => Err(format!(
            "Unsupported expression in value context: {:?}",
            expr
        )),
    }
}

/// Lower an arithmetic binary operation (e.g., `i + 1`)
fn lower_arithmetic_binop<F>(
    operator: &BinaryOperator,
    left: &ASTNode,
    right: &ASTNode,
    alloc_value: &mut F,
    env: &ConditionEnv,
    instructions: &mut Vec<JoinInst>,
) -> Result<ValueId, String>
where
    F: FnMut() -> ValueId,
{
    let lhs = lower_value_expression(left, alloc_value, env, instructions)?;
    let rhs = lower_value_expression(right, alloc_value, env, instructions)?;
    let dst = alloc_value();

    let bin_op = match operator {
        BinaryOperator::Add => BinOpKind::Add,
        BinaryOperator::Subtract => BinOpKind::Sub,
        BinaryOperator::Multiply => BinOpKind::Mul,
        BinaryOperator::Divide => BinOpKind::Div,
        BinaryOperator::Modulo => BinOpKind::Mod,
        _ => {
            return Err(format!(
                "Unsupported binary operator in expression: {:?}",
                operator
            ));
        }
    };

    instructions.push(JoinInst::Compute(MirLikeInst::BinOp {
        dst,
        op: bin_op,
        lhs,
        rhs,
    }));

    Ok(dst)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{ASTNode, BinaryOperator, LiteralValue, Span};

    /// Helper to create a test ConditionEnv with variables
    fn create_test_env() -> ConditionEnv {
        let mut env = ConditionEnv::new();
        // Register test variables (using JoinIR-local ValueIds)
        env.insert("i".to_string(), ValueId(0));
        env.insert("end".to_string(), ValueId(1));
        env
    }

    #[test]
    fn test_simple_comparison() {
        let env = create_test_env();
        let mut value_counter = 2u32; // Start after i=0, end=1
        let mut alloc_value = || {
            let id = ValueId(value_counter);
            value_counter += 1;
            id
        };

        // AST: i < end
        let ast = ASTNode::BinaryOp {
            operator: BinaryOperator::Less,
            left: Box::new(ASTNode::Variable {
                name: "i".to_string(),
                span: Span::unknown(),
            }),
            right: Box::new(ASTNode::Variable {
                name: "end".to_string(),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        };

        let result = lower_condition_to_joinir(&ast, &mut alloc_value, &env);
        assert!(result.is_ok(), "Simple comparison should succeed");

        let (_cond_value, instructions) = result.unwrap();
        assert_eq!(
            instructions.len(),
            1,
            "Should generate 1 Compare instruction"
        );
    }

    #[test]
    fn test_comparison_with_literal() {
        let env = create_test_env();
        let mut value_counter = 2u32;
        let mut alloc_value = || {
            let id = ValueId(value_counter);
            value_counter += 1;
            id
        };

        // AST: i < 10
        let ast = ASTNode::BinaryOp {
            operator: BinaryOperator::Less,
            left: Box::new(ASTNode::Variable {
                name: "i".to_string(),
                span: Span::unknown(),
            }),
            right: Box::new(ASTNode::Literal {
                value: LiteralValue::Integer(10),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        };

        let result = lower_condition_to_joinir(&ast, &mut alloc_value, &env);
        assert!(result.is_ok(), "Comparison with literal should succeed");

        let (_cond_value, instructions) = result.unwrap();
        // Should have: Const(10), Compare
        assert_eq!(instructions.len(), 2, "Should generate Const + Compare");
    }

    #[test]
    fn test_logical_or() {
        let mut env = ConditionEnv::new();
        env.insert("a".to_string(), ValueId(2));
        env.insert("b".to_string(), ValueId(3));

        let mut value_counter = 4u32;
        let mut alloc_value = || {
            let id = ValueId(value_counter);
            value_counter += 1;
            id
        };

        // AST: a < 5 || b < 5
        let ast = ASTNode::BinaryOp {
            operator: BinaryOperator::Or,
            left: Box::new(ASTNode::BinaryOp {
                operator: BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "a".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(5),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            right: Box::new(ASTNode::BinaryOp {
                operator: BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "b".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(5),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        };

        let result = lower_condition_to_joinir(&ast, &mut alloc_value, &env);
        assert!(result.is_ok(), "OR expression should succeed");

        let (_cond_value, instructions) = result.unwrap();
        // Should have: Const(5), Compare(a<5), Const(5), Compare(b<5), BinOp(Or)
        assert_eq!(instructions.len(), 5, "Should generate proper OR chain");
    }

    #[test]
    fn test_not_operator() {
        let env = create_test_env();
        let mut value_counter = 2u32;
        let mut alloc_value = || {
            let id = ValueId(value_counter);
            value_counter += 1;
            id
        };

        // AST: !(i < end)
        let ast = ASTNode::UnaryOp {
            operator: UnaryOperator::Not,
            operand: Box::new(ASTNode::BinaryOp {
                operator: BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "i".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Variable {
                    name: "end".to_string(),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        };

        let result = lower_condition_to_joinir(&ast, &mut alloc_value, &env);
        assert!(result.is_ok(), "NOT operator should succeed");

        let (_cond_value, instructions) = result.unwrap();
        // Should have: Compare, UnaryOp(Not)
        assert_eq!(instructions.len(), 2, "Should generate Compare + Not");
    }
}