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

//! Phase 171-fix: Condition Expression Environment
//!
//! This module provides the environment for lowering condition expressions to JoinIR.
//! It maps variable names to JoinIR-local ValueIds, ensuring proper separation between
//! HOST and JoinIR value spaces.
//!
//! ## Design Philosophy
//!
//! **Single Responsibility**: This module ONLY handles variable name → ValueId mapping
//! for condition expressions. It does NOT:
//! - Perform AST lowering (that's condition_lowerer.rs)
//! - Extract variables from AST (that's condition_var_extractor.rs)
//! - Manage HOST ↔ JoinIR bindings (that's inline_boundary.rs)

use crate::mir::ValueId;
use std::collections::HashMap;

/// Environment for condition expression lowering
///
/// Maps variable names to JoinIR-local ValueIds. Used when lowering
/// condition AST nodes to JoinIR instructions.
///
/// # Phase 200-B Extension
///
/// Added `captured` field to track function-scoped captured variables
/// separately from loop parameters. Captured variables have ParamRole::Condition
/// and do NOT participate in header PHI or ExitLine.
///
/// # Example
///
/// ```ignore
/// let mut env = ConditionEnv::new();
/// env.insert("i".to_string(), ValueId(0));   // Loop parameter
/// env.insert("end".to_string(), ValueId(1)); // Condition-only var
///
/// // Phase 200-B: Add captured variable
/// env.captured.insert("digits".to_string(), ValueId(2));
///
/// // Later during lowering:
/// if let Some(value_id) = env.get("i") {
///     // Use value_id in JoinIR instruction
/// }
/// ```
#[derive(Debug, Clone, Default)]
pub struct ConditionEnv {
    /// Loop parameters and condition-only variables (legacy)
    name_to_join: HashMap<String, ValueId>,

    /// Phase 200-B: Captured function-scoped variables (ParamRole::Condition)
    ///
    /// These variables are:
    /// - Declared in function scope before the loop
    /// - Never reassigned (effectively immutable)
    /// - Used in loop condition or body
    /// - NOT included in header PHI or ExitLine (condition-only)
    pub captured: HashMap<String, ValueId>,
}

impl ConditionEnv {
    /// Create a new empty environment
    pub fn new() -> Self {
        Self {
            name_to_join: HashMap::new(),
            captured: HashMap::new(),
        }
    }

    /// Insert a variable binding
    ///
    /// # Arguments
    ///
    /// * `name` - Variable name (e.g., "i", "end")
    /// * `join_id` - JoinIR-local ValueId for this variable
    pub fn insert(&mut self, name: String, join_id: ValueId) {
        self.name_to_join.insert(name, join_id);
    }

    /// Look up a variable by name
    ///
    /// Phase 200-B: Searches both name_to_join (loop params) and captured fields.
    ///
    /// Returns `Some(ValueId)` if the variable exists in the environment,
    /// `None` otherwise.
    pub fn get(&self, name: &str) -> Option<ValueId> {
        self.name_to_join.get(name).copied()
            .or_else(|| self.captured.get(name).copied())
    }

    /// Check if a variable exists in the environment
    ///
    /// Phase 200-B: Checks both name_to_join and captured fields.
    pub fn contains(&self, name: &str) -> bool {
        self.name_to_join.contains_key(name) || self.captured.contains_key(name)
    }

    /// Check if a variable is a captured (Condition role) variable
    ///
    /// Phase 200-B: New method to distinguish captured vars from loop params.
    pub fn is_captured(&self, name: &str) -> bool {
        self.captured.contains_key(name)
    }

    /// Get the number of variables in the environment
    ///
    /// Phase 200-B: Counts both name_to_join and captured fields.
    pub fn len(&self) -> usize {
        self.name_to_join.len() + self.captured.len()
    }

    /// Check if the environment is empty
    ///
    /// Phase 200-B: Checks both name_to_join and captured fields.
    pub fn is_empty(&self) -> bool {
        self.name_to_join.is_empty() && self.captured.is_empty()
    }

    /// Get an iterator over all (name, ValueId) pairs
    ///
    /// Phase 200-B: Note - this only iterates over name_to_join (loop params).
    /// For captured variables, access the `captured` field directly.
    pub fn iter(&self) -> impl Iterator<Item = (&String, &ValueId)> {
        self.name_to_join.iter()
    }

    /// Get all variable names (sorted)
    ///
    /// Phase 200-B: Includes both name_to_join and captured variables.
    pub fn names(&self) -> Vec<String> {
        let mut names: Vec<_> = self.name_to_join.keys()
            .chain(self.captured.keys())
            .cloned()
            .collect();
        names.sort();
        names.dedup(); // Remove duplicates (shouldn't happen, but be safe)
        names
    }

    /// Phase 201-A: Get the maximum ValueId used in this environment
    ///
    /// Returns the highest ValueId.0 value from both name_to_join and captured,
    /// or None if the environment is empty.
    ///
    /// This is used by JoinIR lowering to determine the starting point for
    /// alloc_value() to avoid ValueId collisions.
    pub fn max_value_id(&self) -> Option<u32> {
        let name_max = self.name_to_join.values().map(|v| v.0).max();
        let captured_max = self.captured.values().map(|v| v.0).max();

        match (name_max, captured_max) {
            (Some(a), Some(b)) => Some(a.max(b)),
            (Some(a), None) => Some(a),
            (None, Some(b)) => Some(b),
            (None, None) => None,
        }
    }
}

/// Binding between HOST and JoinIR ValueIds for condition variables
///
/// This structure explicitly connects a variable name to both its HOST ValueId
/// (from the host function's variable_map) and its JoinIR ValueId (allocated
/// locally within the JoinIR fragment).
///
/// # Example
///
/// For condition variable "start" in `loop(start < end)`:
///
/// ```ignore
/// ConditionBinding {
///     name: "start".to_string(),
///     host_value: ValueId(33),  // HOST function's ValueId for "start"
///     join_value: ValueId(1),   // JoinIR-local ValueId for "start"
/// }
/// ```
#[derive(Debug, Clone)]
pub struct ConditionBinding {
    /// Variable name (e.g., "start", "end")
    pub name: String,

    /// HOST function's ValueId for this variable
    ///
    /// This comes from `builder.variable_map` in the host function.
    pub host_value: ValueId,

    /// JoinIR-local ValueId for this variable
    ///
    /// This is allocated within the JoinIR fragment and must be remapped
    /// when merging into the host function.
    pub join_value: ValueId,
}

impl ConditionBinding {
    /// Create a new condition binding
    pub fn new(name: String, host_value: ValueId, join_value: ValueId) -> Self {
        Self {
            name,
            host_value,
            join_value,
        }
    }
}

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

    #[test]
    fn test_condition_env_basic() {
        let mut env = ConditionEnv::new();
        assert!(env.is_empty());
        assert_eq!(env.len(), 0);

        env.insert("i".to_string(), ValueId(0));
        assert!(!env.is_empty());
        assert_eq!(env.len(), 1);
        assert!(env.contains("i"));
        assert_eq!(env.get("i"), Some(ValueId(0)));
    }

    #[test]
    fn test_condition_env_multiple_vars() {
        let mut env = ConditionEnv::new();
        env.insert("i".to_string(), ValueId(0));
        env.insert("start".to_string(), ValueId(1));
        env.insert("end".to_string(), ValueId(2));

        assert_eq!(env.len(), 3);
        assert_eq!(env.get("i"), Some(ValueId(0)));
        assert_eq!(env.get("start"), Some(ValueId(1)));
        assert_eq!(env.get("end"), Some(ValueId(2)));
        assert_eq!(env.get("nonexistent"), None);
    }

    #[test]
    fn test_condition_binding() {
        let binding = ConditionBinding::new(
            "start".to_string(),
            ValueId(33),  // HOST
            ValueId(1),   // JoinIR
        );

        assert_eq!(binding.name, "start");
        assert_eq!(binding.host_value, ValueId(33));
        assert_eq!(binding.join_value, ValueId(1));
    }
}