llvm-native-core 0.1.4

//! X86 Expand Reduction — expands vector reduction intrinsics into
//! shuffle+arithmetic sequences for the X86/X86-64 target.
//! Phase 10 — LLVM.TARGET.X86.EXPANDREDUX.1 Court.
//!
//! Clean-room behavioral reconstruction from:
//! - Intel® 64 and IA-32 Architectures Software Developer's Manual Vol 2
//!   (horizontal add/sub instructions: haddps, hsubps, phaddd, etc.)
//! - IEEE 754-2019 Floating-Point Arithmetic (reassociation rules)
//! - LLVM ExpandReductions pass behavioral patterns
//! - Agner Fog's instruction tables (latency/throughput for reduction sequences)
//!
//! Zero LLVM source code consultation.
//!
//! ## Reduction Expansion
//!
//! LLVM IR reduction intrinsics (like `@llvm.vector.reduce.add`) represent
//! a horizontal reduction across all elements of a vector. The X86 backend
//! expands these into efficient shuffle+arithmetic sequences using:
//!
//! - **Tree reduction**: log₂(N) steps of pairwise operations,
//!   e.g. for a 4-element vector:
//!   ```
//!   step 0: shuffle(v, [2,3,0,1]) → v1
//!   step 1: op(v, v1) → v2    (now 2 independent results per 2 elements)
//!   step 2: shuffle(v2, [1,0,?,?]) → v3
//!   step 3: op(v2, v3) → result
//!   ```
//!
//! - **AVX horizontal instructions**: where available (SSE3/AVX), use
//!   `haddps`, `haddpd`, `vhaddps`, `vhaddpd` for float add reductions.
//!
//! - **Fast-math reassociation**: with `reassoc` and `nsz` (no signed zeros)
//!   flags, the reduction order can be rearranged for better ILP.
//!
//! ## Supported reduction kinds
//!
//! | Intrinsic                           | Tree Reduction | AVX HADD | Fast-Math |
//! |-------------------------------------|----------------|----------|-----------|
//! | vector.reduce.add                   | ✅             | ✅       | ✅        |
//! | vector.reduce.mul                   | ✅             | —        | ✅        |
//! | vector.reduce.and                   | ✅             | —        | ✅        |
//! | vector.reduce.or                    | ✅             | —        | ✅        |
//! | vector.reduce.xor                   | ✅             | —        | ✅        |
//! | vector.reduce.smax/smin             | ✅             | —        | —         |
//! | vector.reduce.umax/umin             | ✅             | —        | —         |
//! | vector.reduce.fadd                  | ✅             | ✅       | ✅        |
//! | vector.reduce.fmul                  | ✅             | —        | ✅        |
//! | vector.reduce.fmax/fmin             | ✅             | —        | —         |

use crate::x86::x86_instr_info::{X86InstrInfo, X86Opcode};
use crate::x86::x86_subtarget::X86Subtarget;
use std::collections::{HashMap, HashSet};

// ═══════════════════════════════════════════════════════════════════════════════
// Reduction kinds
// ═══════════════════════════════════════════════════════════════════════════════

/// All supported reduction intrinsics that can be expanded.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum ReductionKind {
    /// Integer addition across vector lanes.
    IntAdd,
    /// Integer multiplication across vector lanes.
    IntMul,
    /// Bitwise AND across vector lanes.
    IntAnd,
    /// Bitwise OR across vector lanes.
    IntOr,
    /// Bitwise XOR across vector lanes.
    IntXor,
    /// Signed integer maximum across vector lanes.
    IntSMax,
    /// Signed integer minimum across vector lanes.
    IntSMin,
    /// Unsigned integer maximum across vector lanes.
    IntUMax,
    /// Unsigned integer minimum across vector lanes.
    IntUMin,
    /// Floating-point addition across vector lanes.
    FAdd,
    /// Floating-point multiplication across vector lanes.
    FMul,
    /// Floating-point maximum across vector lanes.
    FMax,
    /// Floating-point minimum across vector lanes.
    FMin,
}

impl ReductionKind {
    /// Parse a reduction intrinsic name to a ReductionKind.
    pub fn from_intrinsic_name(name: &str) -> Option<Self> {
        match name {
            "llvm.vector.reduce.add" => Some(ReductionKind::IntAdd),
            "llvm.vector.reduce.mul" => Some(ReductionKind::IntMul),
            "llvm.vector.reduce.and" => Some(ReductionKind::IntAnd),
            "llvm.vector.reduce.or" => Some(ReductionKind::IntOr),
            "llvm.vector.reduce.xor" => Some(ReductionKind::IntXor),
            "llvm.vector.reduce.smax" => Some(ReductionKind::IntSMax),
            "llvm.vector.reduce.smin" => Some(ReductionKind::IntSMin),
            "llvm.vector.reduce.umax" => Some(ReductionKind::IntUMax),
            "llvm.vector.reduce.umin" => Some(ReductionKind::IntUMin),
            "llvm.vector.reduce.fadd" => Some(ReductionKind::FAdd),
            "llvm.vector.reduce.fmul" => Some(ReductionKind::FMul),
            "llvm.vector.reduce.fmax" => Some(ReductionKind::FMax),
            "llvm.vector.reduce.fmin" => Some(ReductionKind::FMin),
            _ => None,
        }
    }

    /// Whether this is a floating-point reduction.
    pub fn is_float(self) -> bool {
        matches!(
            self,
            ReductionKind::FAdd | ReductionKind::FMul | ReductionKind::FMax | ReductionKind::FMin
        )
    }

    /// Whether this is an integer reduction.
    pub fn is_integer(self) -> bool {
        !self.is_float()
    }

    /// The scalar operation name (e.g. "add", "fadd", "and").
    pub fn scalar_op(self) -> &'static str {
        match self {
            ReductionKind::IntAdd => "add",
            ReductionKind::IntMul => "mul",
            ReductionKind::IntAnd => "and",
            ReductionKind::IntOr => "or",
            ReductionKind::IntXor => "xor",
            ReductionKind::IntSMax => "icmp+smax",
            ReductionKind::IntSMin => "icmp+smin",
            ReductionKind::IntUMax => "icmp+umax",
            ReductionKind::IntUMin => "icmp+umin",
            ReductionKind::FAdd => "fadd",
            ReductionKind::FMul => "fmul",
            ReductionKind::FMax => "fmax",
            ReductionKind::FMin => "fmin",
        }
    }

    /// Whether this reduction is commutative (operand order doesn't matter).
    pub fn is_commutative(self) -> bool {
        matches!(
            self,
            ReductionKind::IntAdd
                | ReductionKind::IntMul
                | ReductionKind::IntAnd
                | ReductionKind::IntOr
                | ReductionKind::IntXor
                | ReductionKind::FAdd
                | ReductionKind::FMul
                | ReductionKind::FMax
                | ReductionKind::FMin
        )
    }

    /// Whether fast-math reassociation can be applied.
    pub fn allows_reassociation(self) -> bool {
        matches!(
            self,
            ReductionKind::IntAdd
                | ReductionKind::IntMul
                | ReductionKind::IntAnd
                | ReductionKind::IntOr
                | ReductionKind::IntXor
                | ReductionKind::FAdd
                | ReductionKind::FMul
        )
    }

    /// The X86 horizontal instruction for this reduction (if available).
    pub fn horizontal_instruction(self, element_type: &str) -> Option<&'static str> {
        match (self, element_type) {
            (ReductionKind::FAdd, "float") => Some("haddps"),
            (ReductionKind::FAdd, "double") => Some("haddpd"),
            (ReductionKind::IntAdd, "i16") => Some("phaddw"),
            (ReductionKind::IntAdd, "i32") => Some("phaddd"),
            (ReductionKind::IntAdd, "i8") => Some("phaddb"),
            _ => None,
        }
    }

    /// Whether this reduction has native horizontal instruction support on
    /// the given ISA level.
    pub fn has_horizontal_support(self, isa_level: &str) -> bool {
        match isa_level {
            "SSE3" | "SSSE3" | "SSE41" | "SSE42" => {
                matches!(self, ReductionKind::FAdd | ReductionKind::IntAdd)
            }
            "AVX" | "AVX2" | "AVX512F" | "AVX512BW" | "AVX512DQ" | "AVX512VL" => {
                matches!(self, ReductionKind::FAdd | ReductionKind::IntAdd)
            }
            _ => false,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Fast Math Flags
// ═══════════════════════════════════════════════════════════════════════════════

/// Fast Math Flags controlling reduction expansion behavior.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct FastMathFlags {
    /// Allow reassociation (a + b) + c → a + (b + c).
    pub reassoc: bool,
    /// No NaNs (assume operands are never NaN).
    pub no_nans: bool,
    /// No infinities.
    pub no_infs: bool,
    /// No signed zeros (+0 == -0).
    pub no_signed_zeros: bool,
    /// Allow reciprocal approximations.
    pub allow_reciprocal: bool,
    /// Contract operations into FMA.
    pub allow_contract: bool,
    /// Approximate functions.
    pub approx_func: bool,
}

impl FastMathFlags {
    /// Create flags from a string representation (e.g. "reassoc nnan nsz").
    pub fn from_str(s: &str) -> Self {
        let mut flags = FastMathFlags::default();
        for token in s.split_whitespace() {
            match token {
                "reassoc" => flags.reassoc = true,
                "nnan" => flags.no_nans = true,
                "ninf" => flags.no_infs = true,
                "nsz" => flags.no_signed_zeros = true,
                "arcp" => flags.allow_reciprocal = true,
                "contract" => flags.allow_contract = true,
                "afn" => flags.approx_func = true,
                "fast" => {
                    // 'fast' implies all flags.
                    flags.reassoc = true;
                    flags.no_nans = true;
                    flags.no_infs = true;
                    flags.no_signed_zeros = true;
                    flags.allow_reciprocal = true;
                    flags.allow_contract = true;
                    flags.approx_func = true;
                }
                _ => {}
            }
        }
        flags
    }

    /// Whether all fast-math flags are set.
    pub fn is_fast(&self) -> bool {
        self.reassoc
            && self.no_nans
            && self.no_infs
            && self.no_signed_zeros
            && self.allow_reciprocal
            && self.allow_contract
            && self.approx_func
    }

    /// Returns true if reassociation is allowed.
    pub fn can_reassociate(&self) -> bool {
        self.reassoc
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Shuffle patterns for reduction
// ═══════════════════════════════════════════════════════════════════════════════

/// A shuffle step in a reduction sequence.
#[derive(Debug, Clone)]
pub struct ShuffleStep {
    /// The shuffle pattern description.
    pub description: String,
    /// The shuffle mask (element indices).
    pub mask: Vec<u32>,
    /// The operation to perform after shuffling.
    pub operation: String,
    /// Whether this step uses a native horizontal instruction.
    pub is_native_horizontal: bool,
}

/// Compute the shuffle mask for a tree-reduction step.
///
/// For a vector of N elements, step s (0-indexed) pairs elements:
///   dst[i] = src[i + (1 << s)] for even i
///   dst[i] = undef for other i
/// Then operation pairs src[i] with dst[i].
pub fn tree_reduction_shuffle_mask(num_elements: u32, step: u32) -> Vec<u32> {
    let mut mask = Vec::with_capacity(num_elements as usize);
    let offset = 1u32 << step;

    for i in 0..num_elements {
        mask.push(i + offset);
    }

    mask
}

/// Compute the horizontal-add shuffle mask for SSE (haddps).
///
/// haddps on two XMM registers: dst[0]=a[0]+a[1], dst[1]=a[2]+a[3],
/// dst[2]=b[0]+b[1], dst[3]=b[2]+b[3].
/// When reducing a single register, we use the same register for both sources.
pub fn hadd_shuffle_mask(num_elements: u32) -> Vec<u32> {
    match num_elements {
        2 => vec![1, 1],       // self-hadd: [a0+a1, a1+a1]
        4 => vec![1, 3, 5, 7], // SSE haddps pattern
        8 => vec![1, 3, 5, 7, 9, 11, 13, 15],
        _ => (0..num_elements).map(|i| i ^ 1).collect(), // pairwise
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Tree reduction planning
// ═══════════════════════════════════════════════════════════════════════════════

/// A single step in a reduction plan.
#[derive(Debug, Clone)]
pub struct ReductionStep {
    /// The step number (0-indexed).
    pub step_index: u32,
    /// The shuffle mask for this step.
    pub shuffle_mask: Vec<u32>,
    /// The operation to perform.
    pub operation: String,
    /// The result name.
    pub result_name: String,
    /// Whether this step uses a native horizontal instruction.
    pub is_native_horizontal: bool,
    /// Whether this is the final step (produces a scalar).
    pub is_final: bool,
    /// The source register names.
    pub sources: Vec<String>,
}

/// A complete plan for expanding a reduction.
#[derive(Debug, Clone)]
pub struct ReductionPlan {
    /// The reduction kind.
    pub kind: ReductionKind,
    /// The number of elements in the source vector.
    pub num_elements: u32,
    /// The element type (e.g. "float", "i32").
    pub element_type: String,
    /// The vector type (e.g. "<4 x float>").
    pub vector_type: String,
    /// The steps of the reduction.
    pub steps: Vec<ReductionStep>,
    /// Total number of shuffle+op pairs.
    pub total_steps: usize,
    /// Whether the plan uses native horizontal instructions.
    pub uses_horizontal: bool,
    /// Whether fast-math reassociation was applied.
    pub uses_fast_math: bool,
    /// The estimated latency in cycles.
    pub estimated_latency: f64,
    /// The source value name being reduced.
    pub source_name: String,
}

/// Plan a tree reduction for a vector of N elements.
///
/// A tree reduction takes log₂(N) steps, each halving the number of
/// independent results until only one (a scalar) remains.
pub fn plan_tree_reduction(
    kind: ReductionKind,
    num_elements: u32,
    element_type: &str,
    source_name: &str,
    fmf: &FastMathFlags,
) -> ReductionPlan {
    let steps_count = num_elements.ilog2() as usize;
    let mut steps = Vec::with_capacity(steps_count);

    let vector_type = format!("<{} x {}>", num_elements, element_type);
    let op = kind.scalar_op();
    let float = kind.is_float();

    for step in 0..steps_count as u32 {
        let mask = tree_reduction_shuffle_mask(num_elements, step);
        let remaining = num_elements >> (step + 1);
        let is_final = remaining == 0;

        let result_type = if is_final {
            element_type.to_string()
        } else {
            format!("<{} x {}>", remaining.max(1), element_type)
        };

        let result_name = if is_final {
            format!("%reduction.result")
        } else {
            format!("%reduction.step{}", step)
        };

        let source = if step == 0 {
            vec![source_name.to_string()]
        } else {
            vec![format!("%reduction.step{}", step - 1)]
        };

        steps.push(ReductionStep {
            step_index: step,
            shuffle_mask: mask,
            operation: if float {
                format!("f{}", op)
            } else {
                op.to_string()
            },
            result_name,
            is_native_horizontal: false,
            is_final,
            sources: source,
        });
    }

    // Estimate latency: each step is 1 shuffle + 1 arithmetic op.
    // Typical: shuffle ~1 cycle, arithmetic ~3-5 cycles.
    let latency_per_step = if float { 5.0 } else { 4.0 };
    let estimated_latency = steps_count as f64 * latency_per_step;

    ReductionPlan {
        kind,
        num_elements,
        element_type: element_type.to_string(),
        vector_type,
        steps,
        total_steps: steps_count,
        uses_horizontal: false,
        uses_fast_math: fmf.can_reassociate(),
        estimated_latency,
        source_name: source_name.to_string(),
    }
}

/// Plan a reduction using native horizontal instructions (SSE3 haddps, etc.).
pub fn plan_horizontal_reduction(
    kind: ReductionKind,
    num_elements: u32,
    element_type: &str,
    source_name: &str,
    fmf: &FastMathFlags,
) -> ReductionPlan {
    let mut steps = Vec::new();
    let vector_type = format!("<{} x {}>", num_elements, element_type);
    let h_inst = kind.horizontal_instruction(element_type);
    let float = kind.is_float();

    // Horizontal instructions typically pair adjacent elements.
    let mut remaining = num_elements;
    let mut step = 0u32;

    while remaining > 1 {
        let next_remaining = remaining / 2;

        let result_type = if next_remaining == 0 {
            element_type.to_string()
        } else {
            format!("<{} x {}>", next_remaining, element_type)
        };

        let result_name = if next_remaining == 0 {
            "%reduction.result".to_string()
        } else {
            format!("%reduction.hstep{}", step)
        };

        let source = if step == 0 {
            vec![source_name.to_string(), source_name.to_string()]
        } else {
            let prev = format!("%reduction.hstep{}", step - 1);
            vec![prev.clone(), prev]
        };

        steps.push(ReductionStep {
            step_index: step,
            shuffle_mask: hadd_shuffle_mask(remaining),
            operation: h_inst.unwrap_or("hadd").to_string(),
            result_name,
            is_native_horizontal: h_inst.is_some(),
            is_final: next_remaining == 0,
            sources: source,
        });

        remaining = next_remaining;
        step += 1;
    }

    // Horizontal instructions are typically faster.
    let latency_per_step = if float { 3.0 } else { 3.0 };
    let estimated_latency = steps.len() as f64 * latency_per_step;

    ReductionPlan {
        kind,
        num_elements,
        element_type: element_type.to_string(),
        vector_type,
        steps,
        total_steps: step as usize,
        uses_horizontal: h_inst.is_some(),
        uses_fast_math: fmf.can_reassociate(),
        estimated_latency,
        source_name: source_name.to_string(),
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Fast-math reassociation for reductions
// ═══════════════════════════════════════════════════════════════════════════════

/// Reassociate a reduction tree for better ILP when fast-math is enabled.
///
/// With reassociation, instead of sequential tree reduction:
///   (((a+b) + c) + d)
/// We can do a balanced tree:
///   (a+b) + (c+d)
///
/// This reduces the critical path from O(N) to O(log₂ N).
pub fn reassociate_reduction(plan: &ReductionPlan, fmf: &FastMathFlags) -> ReductionPlan {
    if !fmf.can_reassociate() || !plan.kind.allows_reassociation() {
        return plan.clone();
    }

    // The tree reduction is already O(log N) — this is mainly about
    // reordering for better latency hiding. With reassoc+nsz, we can
    // also group by sign or pattern.

    let mut new_plan = plan.clone();
    new_plan.uses_fast_math = true;

    // For float reductions with reassoc, prefer the balanced tree.
    // For integer reductions, the tree is already balanced.
    // The main difference: with reassoc, we can use FMA if contract is enabled.

    if fmf.allow_contract && plan.kind.is_float() {
        // Annotate steps that could be contracted into FMA.
        for step in &mut new_plan.steps {
            if step.step_index >= 2 {
                // Potential FMA contraction opportunity.
                // In practice, this would be lowered to vfmaddXXXps.
            }
        }
    }

    new_plan
}

// ═══════════════════════════════════════════════════════════════════════════════
// Code generation for reduction sequences
// ═══════════════════════════════════════════════════════════════════════════════

/// Generate LLVM IR for a reduction plan.
pub fn generate_reduction_code(plan: &ReductionPlan) -> Vec<String> {
    let mut code = Vec::new();

    for step in &plan.steps {
        // Step 1: Shuffle.
        let mask_str: Vec<String> = step
            .shuffle_mask
            .iter()
            .map(|m| format!("i32 {}", m))
            .collect();
        let mask_name = format!("%mask.step{}", step.step_index);
        let shuffle_name = format!("%shuffle.step{}", step.step_index);

        if step.is_native_horizontal {
            // Native horizontal: use the horizontal intrinsic directly.
            if step.sources.len() >= 2 {
                code.push(format!(
                    "  {} = call {} @llvm.x86.sse3.{}({}, {})",
                    step.result_name,
                    step.operation,
                    step.operation,
                    step.sources[0],
                    step.sources[1],
                ));
            } else {
                code.push(format!(
                    "  {} = call {} @llvm.x86.sse3.{}({}, {})",
                    step.result_name,
                    step.operation,
                    step.operation,
                    step.sources[0],
                    step.sources[0],
                ));
            }
        } else {
            // Tree reduction: shuffle + arithmetic.
            let vector_type = if step.is_final {
                format!("<{} x {}>", step.shuffle_mask.len(), plan.element_type)
            } else {
                plan.vector_type.clone()
            };

            // Generate the shuffle mask constant.
            let mask_constant = format!(
                "  {} = <{} x i32> <{}>",
                mask_name,
                step.shuffle_mask.len(),
                mask_str.join(", ")
            );
            code.push(mask_constant);

            // Shuffle the source.
            let source = &step.sources[0];
            if step.is_final {
                // For the final step, extract the scalar from the shuffle result.
                code.push(format!(
                    "  {} = shufflevector {} {}, {} {}, {}",
                    shuffle_name, plan.vector_type, source, plan.vector_type, source, mask_name
                ));
                code.push(format!(
                    "  {} = extractelement {} {}, i32 0",
                    step.result_name,
                    format!("<{} x {}>", step.shuffle_mask.len(), plan.element_type),
                    shuffle_name
                ));
            } else {
                code.push(format!(
                    "  {} = shufflevector {} {}, {} {}, {}",
                    shuffle_name, plan.vector_type, source, plan.vector_type, source, mask_name
                ));
                // Arithmetic op.
                code.push(format!(
                    "  {} = {} {} {}, {}",
                    step.result_name, step.operation, plan.vector_type, source, shuffle_name
                ));
            }
        }
    }

    code
}

/// Generate X86 assembly-like pseudo-code for a reduction plan.
pub fn generate_x86_pseudo_code(plan: &ReductionPlan) -> Vec<String> {
    let mut code = Vec::new();
    code.push(format!(
        "; Reduction: {:?} on {}",
        plan.kind, plan.vector_type
    ));

    for step in &plan.steps {
        if step.is_native_horizontal {
            code.push(format!(
                "  {} {}, {} ; native horizontal step {}",
                step.operation,
                step.result_name,
                step.sources.join(", "),
                step.step_index
            ));
        } else {
            // Shuffle instruction.
            let shuffle_inst = match plan.num_elements {
                2 => "pshufd",      // or unpcklps
                4 => "pshufd",      // SSE 128-bit
                8 => "vperm2f128",  // AVX 256-bit
                16 => "vshuff32x4", // AVX-512
                _ => "pshufd",
            };

            let mask_str: Vec<String> = step.shuffle_mask.iter().map(|m| m.to_string()).collect();
            code.push(format!(
                "  {} {}, {} ; mask=[{}]",
                shuffle_inst,
                step.result_name.replace("%", ""),
                step.sources.join(", "),
                mask_str.join(",")
            ));

            // Arithmetic instruction.
            code.push(format!(
                "  {} {}, {}",
                step.operation,
                step.result_name.replace("%", ""),
                step.sources.join(", ")
            ));
        }
    }

    code
}

// ═══════════════════════════════════════════════════════════════════════════════
// FMF propagation
// ═══════════════════════════════════════════════════════════════════════════════

/// Propagate fast-math flags from a reduction call to its expansion.
///
/// The expanded sequence should inherit appropriate FMF from the original
/// reduction intrinsic. This ensures that downstream optimizations (like
/// instcombine) can still apply FMF-based transformations.
pub fn propagate_fmf(flags: &FastMathFlags) -> String {
    let mut parts = Vec::new();
    if flags.reassoc {
        parts.push("reassoc");
    }
    if flags.no_nans {
        parts.push("nnan");
    }
    if flags.no_infs {
        parts.push("ninf");
    }
    if flags.no_signed_zeros {
        parts.push("nsz");
    }
    if flags.allow_reciprocal {
        parts.push("arcp");
    }
    if flags.allow_contract {
        parts.push("contract");
    }
    if flags.approx_func {
        parts.push("afn");
    }

    if parts.is_empty() {
        String::new()
    } else {
        parts.join(" ")
    }
}

/// Annotate generated reduction code with FMF.
pub fn annotate_fmf(code: &[String], fmf: &FastMathFlags) -> Vec<String> {
    let fmf_str = propagate_fmf(fmf);
    if fmf_str.is_empty() {
        return code.to_vec();
    }

    code.iter()
        .map(|line| {
            // Annotate arithmetic instructions with FMF.
            if line.contains("fadd")
                || line.contains("fsub")
                || line.contains("fmul")
                || line.contains("fdiv")
            {
                format!("{} {}", line, fmf_str)
            } else {
                line.clone()
            }
        })
        .collect()
}

// ═══════════════════════════════════════════════════════════════════════════════
// Main expander pass
// ═══════════════════════════════════════════════════════════════════════════════

/// Statistics for the reduction expander pass.
#[derive(Debug, Clone, Default)]
pub struct ExpandReductionStats {
    pub reductions_found: usize,
    pub reductions_expanded: usize,
    pub tree_reductions: usize,
    pub horizontal_reductions: usize,
    pub fast_math_applied: usize,
    pub fma_contractions: usize,
    pub int_reductions: usize,
    pub fp_reductions: usize,
    pub sse3_hadd_used: usize,
    pub avx_hadd_used: usize,
    pub total_steps_generated: usize,
    pub reductions_too_small: usize,
}

/// The X86 Expand Reduction pass.
///
/// Expands reduction intrinsics into shuffle+arithmetic sequences,
/// selecting between tree-reduction and horizontal-instruction approaches
/// based on ISA capabilities and fast-math flags.
#[derive(Debug, Clone)]
pub struct X86ExpandReduction {
    /// Whether SSE3 horizontal instructions are available.
    pub has_sse3: bool,
    /// Whether AVX is available (vhaddps/vhaddpd).
    pub has_avx: bool,
    /// Whether AVX2 is available (phaddd/phaddw with 256-bit).
    pub has_avx2: bool,
    /// Whether AVX-512 is available.
    pub has_avx512: bool,
    /// Preferred strategy: tree or horizontal.
    pub prefer_horizontal: bool,
    /// Whether to use fast-math reassociation.
    pub enable_fast_math: bool,
    /// Whether to emit debug comments in generated code.
    pub emit_debug_comments: bool,
    /// Pass statistics.
    pub stats: ExpandReductionStats,
}

impl X86ExpandReduction {
    /// Create a new reduction expander for the given subtarget.
    pub fn new(st: &X86Subtarget) -> Self {
        X86ExpandReduction {
            has_sse3: st.has_sse3(),
            has_avx: st.has_avx(),
            has_avx2: st.has_avx2(),
            has_avx512: st.has_avx512(),
            prefer_horizontal: true,
            enable_fast_math: true,
            emit_debug_comments: false,
            stats: ExpandReductionStats::default(),
        }
    }

    /// Expand a single reduction intrinsic.
    ///
    /// `intrinsic_name` e.g. "llvm.vector.reduce.fadd"
    /// `num_elements` — number of elements in the vector (2, 4, 8, 16, ...)
    /// `element_type` — "float", "double", "i32", etc.
    /// `source_name` — the SSA value name being reduced.
    /// `fmf` — fast-math flags.
    ///
    /// Returns the generated LLVM IR code as strings, and a result name.
    pub fn expand_reduction(
        &mut self,
        intrinsic_name: &str,
        num_elements: u32,
        element_type: &str,
        source_name: &str,
        fmf: &FastMathFlags,
    ) -> Option<(Vec<String>, String)> {
        let kind = ReductionKind::from_intrinsic_name(intrinsic_name)?;
        self.stats.reductions_found += 1;

        if kind.is_float() {
            self.stats.fp_reductions += 1;
        } else {
            self.stats.int_reductions += 1;
        }

        // Reject single-element reductions (they're just the element itself).
        if num_elements < 2 {
            self.stats.reductions_too_small += 1;
            return None;
        }

        // Determine ISA level string.
        let isa_level = if self.has_avx512 {
            "AVX512F"
        } else if self.has_avx2 {
            "AVX2"
        } else if self.has_avx {
            "AVX"
        } else if self.has_sse3 {
            "SSE3"
        } else {
            "SSE2"
        };

        // Choose strategy.
        let use_horizontal = self.prefer_horizontal
            && kind.has_horizontal_support(isa_level)
            && (self.has_sse3 || self.has_avx);

        let plan = if use_horizontal {
            plan_horizontal_reduction(kind, num_elements, element_type, source_name, fmf)
        } else {
            plan_tree_reduction(kind, num_elements, element_type, source_name, fmf)
        };

        // Apply fast-math reassociation if enabled.
        let final_plan = if self.enable_fast_math && fmf.can_reassociate() {
            self.stats.fast_math_applied += 1;
            reassociate_reduction(&plan, fmf)
        } else {
            plan
        };

        let mut code = Vec::new();

        if self.emit_debug_comments {
            code.push(format!(
                "; Expanding {} of {} ({} elements, {} strategy)",
                intrinsic_name,
                source_name,
                num_elements,
                if use_horizontal { "horizontal" } else { "tree" }
            ));
        }

        // Generate the code.
        let generated = generate_reduction_code(&final_plan);
        code.extend(generated);

        // Apply FMF annotations.
        let annotated = annotate_fmf(&code, fmf);

        let result_name = if final_plan.steps.is_empty() {
            source_name.to_string()
        } else {
            final_plan.steps.last().unwrap().result_name.clone()
        };

        self.stats.reductions_expanded += 1;

        if use_horizontal {
            self.stats.horizontal_reductions += 1;
            if self.has_avx {
                self.stats.avx_hadd_used += 1;
            } else {
                self.stats.sse3_hadd_used += 1;
            }
        } else {
            self.stats.tree_reductions += 1;
        }

        self.stats.total_steps_generated += final_plan.total_steps;

        Some((annotated, result_name))
    }

    /// Expand all reduction intrinsics found in a set of instructions.
    pub fn expand_all(
        &mut self,
        instructions: &[(usize, &str, u32, &str, &str, &str)],
    ) -> Vec<(usize, Vec<String>, String)> {
        let mut results = Vec::new();

        for &(idx, intrinsic_name, num_elements, element_type, source_name, fmf_str) in instructions
        {
            let fmf = FastMathFlags::from_str(fmf_str);
            if let Some((code, result)) = self.expand_reduction(
                intrinsic_name,
                num_elements,
                element_type,
                source_name,
                &fmf,
            ) {
                results.push((idx, code, result));
            }
        }

        results
    }

    /// Return the ISA level name for the current subtarget.
    pub fn isa_level_name(&self) -> &'static str {
        if self.has_avx512 {
            "AVX-512"
        } else if self.has_avx2 {
            "AVX2"
        } else if self.has_avx {
            "AVX"
        } else if self.has_sse3 {
            "SSE3"
        } else {
            "SSE2"
        }
    }

    /// Return a summary of pass statistics.
    pub fn stats_summary(&self) -> String {
        format!(
            "Redux Expand ({}): {} found, {} expanded ({} tree, {} horizontal), \
             {} fast-math, {} FMA contracted, {} int, {} fp, \
             {} SSE3 hadd, {} AVX hadd, {} steps generated, {} too small",
            self.isa_level_name(),
            self.stats.reductions_found,
            self.stats.reductions_expanded,
            self.stats.tree_reductions,
            self.stats.horizontal_reductions,
            self.stats.fast_math_applied,
            self.stats.fma_contractions,
            self.stats.int_reductions,
            self.stats.fp_reductions,
            self.stats.sse3_hadd_used,
            self.stats.avx_hadd_used,
            self.stats.total_steps_generated,
            self.stats.reductions_too_small,
        )
    }

    /// Reset all statistics.
    pub fn clear(&mut self) {
        self.stats = ExpandReductionStats::default();
    }
}

impl Default for X86ExpandReduction {
    fn default() -> Self {
        X86ExpandReduction {
            has_sse3: true,
            has_avx: false,
            has_avx2: false,
            has_avx512: false,
            prefer_horizontal: true,
            enable_fast_math: true,
            emit_debug_comments: false,
            stats: ExpandReductionStats::default(),
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Specialized reducers for common X86 patterns
// ═══════════════════════════════════════════════════════════════════════════════

/// Generate an optimal x86 reduction sequence for float add (haddps-based).
pub fn generate_fadd_reduction_sse3(num_elements: u32, source: &str) -> Vec<String> {
    let mut code = Vec::new();
    let mut current = source.to_string();

    match num_elements {
        4 => {
            // SSE: 4xfloat → haddps x, x → haddps x, x → extract
            code.push(format!(
                "  %h0 = call <4 x float> @llvm.x86.sse3.haddps({}, {})",
                current, current
            ));
            code.push(format!(
                "  %h1 = call <4 x float> @llvm.x86.sse3.haddps(%h0, %h0)"
            ));
            code.push("  %result = extractelement <4 x float> %h1, i32 0".to_string());
        }
        8 => {
            // AVX: 8xfloat → vextractf128 + haddps chain
            code.push(format!(
                "  %h0 = call <8 x float> @llvm.x86.avx.haddps({}, {})",
                current, current
            ));
            code.push(format!(
                "  %h1 = call <8 x float> @llvm.x86.avx.haddps(%h0, %h0)"
            ));
            code.push(format!("  %lo = extractelement <8 x float> %h1, i32 0"));
            code.push(format!("  %hi = extractelement <8 x float> %h1, i32 4"));
            code.push("  %result = fadd float %lo, %hi".to_string());
        }
        _ => {
            // Generic tree reduction.
            code.push(format!(
                "; Generic fadd reduction for {} elements",
                num_elements
            ));
            let plan = plan_tree_reduction(
                ReductionKind::FAdd,
                num_elements,
                "float",
                source,
                &FastMathFlags::default(),
            );
            code.extend(generate_reduction_code(&plan));
        }
    }

    code
}

/// Generate an optimal x86 reduction sequence for integer add (phaddd-based).
pub fn generate_iadd_reduction_ssse3(num_elements: u32, source: &str) -> Vec<String> {
    let mut code = Vec::new();

    match num_elements {
        4 => {
            // SSSE3: 4xi32 → phaddd x, x → phaddd x, x → extract
            code.push(format!(
                "  %h0 = call <4 x i32> @llvm.x86.ssse3.phadd.d.128({}, {})",
                source, source
            ));
            code.push(format!(
                "  %h1 = call <4 x i32> @llvm.x86.ssse3.phadd.d.128(%h0, %h0)"
            ));
            code.push("  %result = extractelement <4 x i32> %h1, i32 0".to_string());
        }
        8 => {
            // AVX2: 8xi32 → vphaddd chain
            code.push(format!(
                "  %h0 = call <8 x i32> @llvm.x86.avx2.phadd.d({}, {})",
                source, source
            ));
            code.push(format!(
                "  %h1 = call <8 x i32> @llvm.x86.avx2.phadd.d(%h0, %h0)"
            ));
            code.push("  %result = extractelement <8 x i32> %h1, i32 0".to_string());
        }
        _ => {
            code.push(format!(
                "; Generic iadd reduction for {} elements",
                num_elements
            ));
            let plan = plan_tree_reduction(
                ReductionKind::IntAdd,
                num_elements,
                "i32",
                source,
                &FastMathFlags::default(),
            );
            code.extend(generate_reduction_code(&plan));
        }
    }

    code
}

/// Generate a dot-product reduction (pairwise mul then horizontal add).
pub fn generate_dot_product_reduction(
    num_elements: u32,
    element_type: &str,
    lhs: &str,
    rhs: &str,
) -> Vec<String> {
    let mut code = Vec::new();
    let vec_type = format!("<{} x {}>", num_elements, element_type);

    // Step 1: Pairwise multiply.
    code.push(format!("  %mul = fmul {} {}, {}", vec_type, lhs, rhs));

    // Step 2: Horizontal add.
    let add_code = if element_type == "float" {
        generate_fadd_reduction_sse3(num_elements, "%mul")
    } else if element_type == "double" {
        // Double: use haddpd.
        let mut dc = Vec::new();
        dc.push(format!(
            "  %h0 = call {} @llvm.x86.sse3.haddpd(%mul, %mul)",
            vec_type
        ));
        dc.push(format!(
            "  %result = extractelement {} %h0, i32 0",
            vec_type
        ));
        dc
    } else {
        generate_iadd_reduction_ssse3(num_elements, "%mul")
    };

    // Replace %mul references in add_code — they already match.
    code.extend(add_code);
    code
}

/// Generate a sum-of-absolute-differences (SAD) reduction.
///
/// SAD: Σ |a[i] - b[i]| for i in 0..N
pub fn generate_sad_reduction(a: &str, b: &str) -> Vec<String> {
    let mut code = Vec::new();

    // SAD is natively supported via PSADBW on SSE.
    // For 8-bit elements:
    code.push(format!(
        "  %sub = call <16 x i8> @llvm.x86.sse2.psub.us.b({}, {})",
        a, b
    ));
    code.push("  ; Absolute difference: use pabsb (SSSE3) or manual abs".to_string());
    code.push(format!(
        "  %abs = call <16 x i8> @llvm.x86.ssse3.pabs.b.128(%sub)"
    ));
    code.push(format!(
        "  %sad = call <2 x i64> @llvm.x86.sse2.psadbw(%abs, zeroinitializer)"
    ));
    code.push("  %result = extractelement <2 x i64> %sad, i32 0".to_string());

    code
}

// ═══════════════════════════════════════════════════════════════════════════════
// Factory functions
// ═══════════════════════════════════════════════════════════════════════════════

/// Create a reduction expander for SSE3.
pub fn make_x86_expand_reduction_sse3() -> X86ExpandReduction {
    X86ExpandReduction {
        has_sse3: true,
        has_avx: false,
        has_avx2: false,
        has_avx512: false,
        prefer_horizontal: true,
        enable_fast_math: true,
        emit_debug_comments: false,
        stats: ExpandReductionStats::default(),
    }
}

/// Create a reduction expander for AVX.
pub fn make_x86_expand_reduction_avx() -> X86ExpandReduction {
    X86ExpandReduction {
        has_sse3: true,
        has_avx: true,
        has_avx2: false,
        has_avx512: false,
        prefer_horizontal: true,
        enable_fast_math: true,
        emit_debug_comments: false,
        stats: ExpandReductionStats::default(),
    }
}

/// Create a reduction expander for AVX2.
pub fn make_x86_expand_reduction_avx2() -> X86ExpandReduction {
    X86ExpandReduction {
        has_sse3: true,
        has_avx: true,
        has_avx2: true,
        has_avx512: false,
        prefer_horizontal: true,
        enable_fast_math: true,
        emit_debug_comments: false,
        stats: ExpandReductionStats::default(),
    }
}

/// Create a reduction expander for AVX-512.
pub fn make_x86_expand_reduction_avx512() -> X86ExpandReduction {
    X86ExpandReduction {
        has_sse3: true,
        has_avx: true,
        has_avx2: true,
        has_avx512: true,
        prefer_horizontal: true,
        enable_fast_math: true,
        emit_debug_comments: false,
        stats: ExpandReductionStats::default(),
    }
}

/// Create a conservative reduction expander (tree-only, no fast-math).
pub fn make_x86_expand_reduction_conservative() -> X86ExpandReduction {
    X86ExpandReduction {
        has_sse3: false,
        has_avx: false,
        has_avx2: false,
        has_avx512: false,
        prefer_horizontal: false,
        enable_fast_math: false,
        emit_debug_comments: false,
        stats: ExpandReductionStats::default(),
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════════════════

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

    // ═══════════════════════════════════════════════════════════════════
    // ReductionKind tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_reduction_kind_from_intrinsic() {
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.add"),
            Some(ReductionKind::IntAdd)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.fadd"),
            Some(ReductionKind::FAdd)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.mul"),
            Some(ReductionKind::IntMul)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.fmul"),
            Some(ReductionKind::FMul)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.and"),
            Some(ReductionKind::IntAnd)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.or"),
            Some(ReductionKind::IntOr)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.xor"),
            Some(ReductionKind::IntXor)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.smax"),
            Some(ReductionKind::IntSMax)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.smin"),
            Some(ReductionKind::IntSMin)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.umax"),
            Some(ReductionKind::IntUMax)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.umin"),
            Some(ReductionKind::IntUMin)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.fmax"),
            Some(ReductionKind::FMax)
        );
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.fmin"),
            Some(ReductionKind::FMin)
        );
    }

    #[test]
    fn test_reduction_kind_unknown() {
        assert_eq!(
            ReductionKind::from_intrinsic_name("llvm.vector.reduce.unknown"),
            None
        );
        assert_eq!(ReductionKind::from_intrinsic_name(""), None);
    }

    #[test]
    fn test_reduction_kind_is_float() {
        assert!(ReductionKind::FAdd.is_float());
        assert!(ReductionKind::FMul.is_float());
        assert!(!ReductionKind::IntAdd.is_float());
        assert!(!ReductionKind::IntAnd.is_float());
    }

    #[test]
    fn test_reduction_kind_is_integer() {
        assert!(ReductionKind::IntAdd.is_integer());
        assert!(ReductionKind::IntAnd.is_integer());
        assert!(!ReductionKind::FAdd.is_integer());
    }

    #[test]
    fn test_reduction_kind_commutative() {
        assert!(ReductionKind::IntAdd.is_commutative());
        assert!(ReductionKind::IntMul.is_commutative());
        assert!(ReductionKind::FAdd.is_commutative());
        assert!(ReductionKind::FMul.is_commutative());
        assert!(ReductionKind::IntAnd.is_commutative());
        assert!(ReductionKind::IntOr.is_commutative());
        assert!(ReductionKind::IntXor.is_commutative());
        assert!(!ReductionKind::IntSMax.is_commutative()); // min/max are not strictly commutative for NaN
        assert!(!ReductionKind::IntSMin.is_commutative());
    }

    #[test]
    fn test_reduction_kind_allows_reassociation() {
        assert!(ReductionKind::FAdd.allows_reassociation());
        assert!(ReductionKind::FMul.allows_reassociation());
        assert!(ReductionKind::IntAdd.allows_reassociation());
        assert!(!ReductionKind::FMax.allows_reassociation());
    }

    #[test]
    fn test_reduction_kind_horizontal_instruction() {
        assert_eq!(
            ReductionKind::FAdd.horizontal_instruction("float"),
            Some("haddps")
        );
        assert_eq!(
            ReductionKind::FAdd.horizontal_instruction("double"),
            Some("haddpd")
        );
        assert_eq!(
            ReductionKind::IntAdd.horizontal_instruction("i32"),
            Some("phaddd")
        );
        assert_eq!(ReductionKind::FMul.horizontal_instruction("float"), None);
    }

    #[test]
    fn test_reduction_kind_has_horizontal_support() {
        assert!(ReductionKind::FAdd.has_horizontal_support("SSE3"));
        assert!(ReductionKind::FAdd.has_horizontal_support("AVX"));
        assert!(ReductionKind::FAdd.has_horizontal_support("AVX2"));
        assert!(!ReductionKind::FAdd.has_horizontal_support("SSE2"));
        assert!(!ReductionKind::FMul.has_horizontal_support("AVX2"));
    }

    #[test]
    fn test_reduction_kind_scalar_op() {
        assert_eq!(ReductionKind::IntAdd.scalar_op(), "add");
        assert_eq!(ReductionKind::FAdd.scalar_op(), "fadd");
        assert_eq!(ReductionKind::IntAnd.scalar_op(), "and");
        assert_eq!(ReductionKind::FMul.scalar_op(), "fmul");
    }

    // ═══════════════════════════════════════════════════════════════════
    // FastMathFlags tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_fmf_from_str_empty() {
        let fmf = FastMathFlags::from_str("");
        assert!(!fmf.reassoc);
        assert!(!fmf.no_nans);
    }

    #[test]
    fn test_fmf_from_str_reassoc() {
        let fmf = FastMathFlags::from_str("reassoc");
        assert!(fmf.reassoc);
        assert!(!fmf.no_nans);
    }

    #[test]
    fn test_fmf_from_str_multiple() {
        let fmf = FastMathFlags::from_str("reassoc nnan nsz");
        assert!(fmf.reassoc);
        assert!(fmf.no_nans);
        assert!(fmf.no_signed_zeros);
        assert!(!fmf.no_infs);
    }

    #[test]
    fn test_fmf_from_str_fast() {
        let fmf = FastMathFlags::from_str("fast");
        assert!(fmf.is_fast());
        assert!(fmf.can_reassociate());
    }

    #[test]
    fn test_fmf_from_str_arcp() {
        let fmf = FastMathFlags::from_str("arcp contract");
        assert!(fmf.allow_reciprocal);
        assert!(fmf.allow_contract);
        assert!(!fmf.approx_func);
    }

    #[test]
    fn test_fmf_default() {
        let fmf = FastMathFlags::default();
        assert!(!fmf.reassoc);
        assert!(!fmf.no_nans);
        assert!(!fmf.is_fast());
    }

    #[test]
    fn test_fmf_can_reassociate() {
        let mut fmf = FastMathFlags::default();
        assert!(!fmf.can_reassociate());
        fmf.reassoc = true;
        assert!(fmf.can_reassociate());
    }

    // ═══════════════════════════════════════════════════════════════════
    // Shuffle mask tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_tree_reduction_mask_step0_4elems() {
        let mask = tree_reduction_shuffle_mask(4, 0);
        // Step 0 pairs: src[i] with src[i+1]
        // dst[0] = src[1], dst[1] = src[2], dst[2] = src[3], dst[3] = src[4]
        assert_eq!(mask, vec![1, 2, 3, 4]);
    }

    #[test]
    fn test_tree_reduction_mask_step1_4elems() {
        let mask = tree_reduction_shuffle_mask(4, 1);
        // Step 1 pairs: src[i] with src[i+2]
        assert_eq!(mask, vec![2, 3, 4, 5]);
    }

    #[test]
    fn test_tree_reduction_mask_step0_8elems() {
        let mask = tree_reduction_shuffle_mask(8, 0);
        assert_eq!(mask.len(), 8);
        assert_eq!(mask[0], 1);
        assert_eq!(mask[1], 2);
    }

    #[test]
    fn test_tree_reduction_mask_step2_8elems() {
        let mask = tree_reduction_shuffle_mask(8, 2);
        assert_eq!(mask.len(), 8);
        assert_eq!(mask[0], 4);
    }

    #[test]
    fn test_hadd_shuffle_mask_4() {
        let mask = hadd_shuffle_mask(4);
        assert_eq!(mask, vec![1, 3, 5, 7]);
    }

    #[test]
    fn test_hadd_shuffle_mask_2() {
        let mask = hadd_shuffle_mask(2);
        assert_eq!(mask, vec![1, 1]);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Tree reduction plan tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_plan_tree_reduction_fadd_4() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert_eq!(plan.total_steps, 2); // log2(4) = 2
        assert_eq!(plan.steps.len(), 2);
        assert!(!plan.uses_horizontal);
        assert_eq!(plan.num_elements, 4);
    }

    #[test]
    fn test_plan_tree_reduction_fadd_8() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            8,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert_eq!(plan.total_steps, 3); // log2(8) = 3
        assert_eq!(plan.steps.len(), 3);
    }

    #[test]
    fn test_plan_tree_reduction_fadd_16() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            16,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert_eq!(plan.total_steps, 4); // log2(16) = 4
    }

    #[test]
    fn test_plan_tree_reduction_int_and() {
        let plan = plan_tree_reduction(
            ReductionKind::IntAnd,
            4,
            "i32",
            "%v",
            &FastMathFlags::default(),
        );
        assert_eq!(plan.total_steps, 2);
        assert_eq!(plan.kind, ReductionKind::IntAnd);
    }

    #[test]
    fn test_plan_tree_reduction_last_step_is_final() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(plan.steps.last().unwrap().is_final);
        assert!(!plan.steps[0].is_final);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Horizontal plan tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_plan_horizontal_reduction_fadd_4() {
        let plan = plan_horizontal_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(plan.uses_horizontal);
        assert!(plan.total_steps >= 2); // haddps takes 2 steps for 4-wide
    }

    #[test]
    fn test_plan_horizontal_reduction_fadd_8() {
        let plan = plan_horizontal_reduction(
            ReductionKind::FAdd,
            8,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(plan.uses_horizontal);
        assert!(plan.total_steps >= 3);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Reassociation tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_reassociate_with_reassoc() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let fmf = FastMathFlags::from_str("reassoc");
        let reassoc = reassociate_reduction(&plan, &fmf);
        assert!(reassoc.uses_fast_math);
    }

    #[test]
    fn test_reassociate_without_reassoc() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let fmf = FastMathFlags::default();
        let reassoc = reassociate_reduction(&plan, &fmf);
        assert!(!reassoc.uses_fast_math); // No reassoc flag, plan unchanged.
    }

    #[test]
    fn test_reassociate_non_reassociable_kind() {
        let plan = plan_tree_reduction(
            ReductionKind::FMax,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let fmf = FastMathFlags::from_str("reassoc");
        let reassoc = reassociate_reduction(&plan, &fmf);
        // FMax can't be reassociated, so plan should be unchanged.
        assert_eq!(reassoc.total_steps, plan.total_steps);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Code generation tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_generate_reduction_code_tree() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let code = generate_reduction_code(&plan);
        assert!(code.len() >= 4); // 2 shuffle masks + 2 shuffles + 2 ops
        assert!(code.iter().any(|l| l.contains("shufflevector")));
        assert!(code.iter().any(|l| l.contains("fadd")));
    }

    #[test]
    fn test_generate_reduction_code_int() {
        let plan = plan_tree_reduction(
            ReductionKind::IntAdd,
            4,
            "i32",
            "%v",
            &FastMathFlags::default(),
        );
        let code = generate_reduction_code(&plan);
        assert!(code.iter().any(|l| l.contains("add")));
    }

    #[test]
    fn test_generate_x86_pseudo_code() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let code = generate_x86_pseudo_code(&plan);
        assert!(code.len() >= 4);
        assert!(code.iter().any(|l| l.contains("pshufd")));
    }

    // ═══════════════════════════════════════════════════════════════════
    // FMF propagation tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_propagate_fmf_empty() {
        let fmf = FastMathFlags::default();
        assert_eq!(propagate_fmf(&fmf), "");
    }

    #[test]
    fn test_propagate_fmf_reassoc() {
        let fmf = FastMathFlags::from_str("reassoc");
        assert_eq!(propagate_fmf(&fmf), "reassoc");
    }

    #[test]
    fn test_propagate_fmf_multiple() {
        let fmf = FastMathFlags::from_str("reassoc nsz nnan");
        assert_eq!(propagate_fmf(&fmf), "reassoc nnan nsz");
    }

    #[test]
    fn test_propagate_fmf_fast() {
        let fmf = FastMathFlags::from_str("fast");
        let result = propagate_fmf(&fmf);
        assert!(result.contains("reassoc"));
        assert!(result.contains("nnan"));
        assert!(result.contains("nsz"));
    }

    #[test]
    fn test_annotate_fmf_no_flags() {
        let code = vec!["  %r = fadd float %a, %b".to_string()];
        let fmf = FastMathFlags::default();
        let annotated = annotate_fmf(&code, &fmf);
        assert_eq!(annotated, code); // No change if no flags.
    }

    #[test]
    fn test_annotate_fmf_with_flags() {
        let code = vec![
            "  %r = fadd <4 x float> %a, %b".to_string(),
            "  %s = fmul <4 x float> %r, %c".to_string(),
        ];
        let fmf = FastMathFlags::from_str("reassoc nsz");
        let annotated = annotate_fmf(&code, &fmf);
        assert!(annotated[0].ends_with("reassoc nsz"));
        assert!(annotated[1].ends_with("reassoc nsz"));
    }

    // ═══════════════════════════════════════════════════════════════════
    // Specialized reducer tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_generate_fadd_reduction_sse3_4() {
        let code = generate_fadd_reduction_sse3(4, "%v");
        assert!(code.len() >= 3);
        assert!(code.iter().any(|l| l.contains("haddps")));
        assert!(code.iter().any(|l| l.contains("extractelement")));
    }

    #[test]
    fn test_generate_fadd_reduction_sse3_8() {
        let code = generate_fadd_reduction_sse3(8, "%v");
        assert!(code.len() >= 5);
        assert!(code.iter().any(|l| l.contains("haddps")));
        assert!(code.iter().any(|l| l.contains("extractelement")));
    }

    #[test]
    fn test_generate_iadd_reduction_ssse3_4() {
        let code = generate_iadd_reduction_ssse3(4, "%v");
        assert!(code.iter().any(|l| l.contains("phadd")));
        assert!(code.iter().any(|l| l.contains("extractelement")));
    }

    #[test]
    fn test_generate_dot_product_reduction() {
        let code = generate_dot_product_reduction(4, "float", "%a", "%b");
        assert!(code.iter().any(|l| l.contains("fmul")));
        assert!(code
            .iter()
            .any(|l| l.contains("haddps") || l.contains("fadd")));
    }

    #[test]
    fn test_generate_sad_reduction() {
        let code = generate_sad_reduction("%a", "%b");
        assert!(code.iter().any(|l| l.contains("psadbw")));
    }

    // ═══════════════════════════════════════════════════════════════════
    // Expander integration tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_expander_new_sse3() {
        let mut st = X86Subtarget::default();
        st.set_sse3(true);
        let expander = X86ExpandReduction::new(&st);
        assert!(expander.has_sse3);
        assert!(!expander.has_avx);
        assert!(expander.prefer_horizontal);
    }

    #[test]
    fn test_expander_new_avx() {
        let mut st = X86Subtarget::default();
        st.set_avx(true);
        st.set_sse3(true);
        let expander = X86ExpandReduction::new(&st);
        assert!(expander.has_avx);
        assert!(expander.has_sse3);
    }

    #[test]
    fn test_expander_new_avx512() {
        let mut st = X86Subtarget::default();
        st.set_avx512f(true);
        st.set_avx(true);
        st.set_sse3(true);
        let expander = X86ExpandReduction::new(&st);
        assert!(expander.has_avx512);
    }

    #[test]
    fn test_expand_reduction_fadd_sse3() {
        let mut expander = make_x86_expand_reduction_sse3();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fadd",
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        let (code, _result_name) = result.unwrap();
        assert!(!code.is_empty());
        assert_eq!(expander.stats.reductions_found, 1);
        assert_eq!(expander.stats.reductions_expanded, 1);
        assert_eq!(expander.stats.fp_reductions, 1);
    }

    #[test]
    fn test_expand_reduction_single_element() {
        let mut expander = make_x86_expand_reduction_sse3();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fadd",
            1,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_none());
        assert_eq!(expander.stats.reductions_too_small, 1);
    }

    #[test]
    fn test_expand_reduction_unknown_intrinsic() {
        let mut expander = make_x86_expand_reduction_sse3();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.nonexistent",
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_none());
    }

    #[test]
    fn test_expand_reduction_int_add() {
        let mut expander = make_x86_expand_reduction_sse3();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.add",
            4,
            "i32",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        assert_eq!(expander.stats.int_reductions, 1);
    }

    #[test]
    fn test_expand_reduction_with_fast_math() {
        let mut expander = make_x86_expand_reduction_sse3();
        let fmf = FastMathFlags::from_str("reassoc nnan nsz");
        let result = expander.expand_reduction("llvm.vector.reduce.fadd", 4, "float", "%v", &fmf);
        assert!(result.is_some());
        assert_eq!(expander.stats.fast_math_applied, 1);
    }

    #[test]
    fn test_expand_all() {
        let mut expander = make_x86_expand_reduction_sse3();
        let insts = vec![
            (0usize, "llvm.vector.reduce.fadd", 4u32, "float", "%v0", ""),
            (1usize, "llvm.vector.reduce.add", 4u32, "i32", "%v1", ""),
        ];
        let results = expander.expand_all(&insts);
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn test_expander_isa_level_name() {
        let sse3 = make_x86_expand_reduction_sse3();
        assert_eq!(sse3.isa_level_name(), "SSE3");

        let avx = make_x86_expand_reduction_avx();
        assert_eq!(avx.isa_level_name(), "AVX");

        let avx512 = make_x86_expand_reduction_avx512();
        assert_eq!(avx512.isa_level_name(), "AVX-512");
    }

    #[test]
    fn test_stats_summary() {
        let mut expander = make_x86_expand_reduction_sse3();
        expander.stats.reductions_found = 10;
        expander.stats.reductions_expanded = 8;
        expander.stats.tree_reductions = 3;
        expander.stats.horizontal_reductions = 5;
        let summary = expander.stats_summary();
        assert!(summary.contains("10"));
        assert!(summary.contains("8"));
        assert!(summary.contains("3"));
        assert!(summary.contains("5"));
    }

    #[test]
    fn test_clear() {
        let mut expander = make_x86_expand_reduction_sse3();
        expander.stats.reductions_found = 100;
        expander.clear();
        assert_eq!(expander.stats.reductions_found, 0);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Factory function tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_make_sse3() {
        let e = make_x86_expand_reduction_sse3();
        assert!(e.has_sse3);
        assert!(!e.has_avx);
        assert!(e.prefer_horizontal);
    }

    #[test]
    fn test_make_avx() {
        let e = make_x86_expand_reduction_avx();
        assert!(e.has_avx);
        assert!(!e.has_avx2);
    }

    #[test]
    fn test_make_avx2() {
        let e = make_x86_expand_reduction_avx2();
        assert!(e.has_avx2);
        assert!(!e.has_avx512);
    }

    #[test]
    fn test_make_avx512() {
        let e = make_x86_expand_reduction_avx512();
        assert!(e.has_avx512);
    }

    #[test]
    fn test_make_conservative() {
        let e = make_x86_expand_reduction_conservative();
        assert!(!e.has_sse3);
        assert!(!e.prefer_horizontal);
        assert!(!e.enable_fast_math);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Default impl tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_default_expander() {
        let e = X86ExpandReduction::default();
        assert!(e.has_sse3);
        assert!(!e.has_avx);
        assert!(e.prefer_horizontal);
        assert!(e.enable_fast_math);
    }

    // ═══════════════════════════════════════════════════════════════════
    // Edge case tests
    // ═══════════════════════════════════════════════════════════════════

    #[test]
    fn test_expand_2_element_vector() {
        let mut expander = make_x86_expand_reduction_sse3();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fadd",
            2,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        let (code, _) = result.unwrap();
        // 2-element reduction: 1 step.
        assert!(!code.is_empty());
    }

    #[test]
    fn test_expand_16_element_vector() {
        let mut expander = make_x86_expand_reduction_avx512();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fadd",
            16,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        let (_code, _) = result.unwrap();
        assert_eq!(expander.stats.reductions_expanded, 1);
    }

    #[test]
    fn test_expand_all_kinds() {
        let mut expander = make_x86_expand_reduction_sse3();
        let kinds = [
            "llvm.vector.reduce.add",
            "llvm.vector.reduce.mul",
            "llvm.vector.reduce.and",
            "llvm.vector.reduce.or",
            "llvm.vector.reduce.xor",
            "llvm.vector.reduce.smax",
            "llvm.vector.reduce.smin",
            "llvm.vector.reduce.umax",
            "llvm.vector.reduce.umin",
            "llvm.vector.reduce.fadd",
            "llvm.vector.reduce.fmul",
            "llvm.vector.reduce.fmax",
            "llvm.vector.reduce.fmin",
        ];
        for kind_str in &kinds {
            let mut e = make_x86_expand_reduction_sse3();
            let result = e.expand_reduction(kind_str, 4, "i32", "%v", &FastMathFlags::default());
            // Float kinds with i32 element type would be nonsensical, but the
            // expander should handle it gracefully.
            assert!(result.is_some() || e.stats.reductions_found > 0);
        }
    }

    #[test]
    fn test_fmf_from_str_unknown_token() {
        let fmf = FastMathFlags::from_str("unknown_flag another_one");
        assert!(!fmf.reassoc);
        assert!(!fmf.no_nans);
    }

    #[test]
    fn test_propagate_fmf_all() {
        let fmf = FastMathFlags {
            reassoc: true,
            no_nans: true,
            no_infs: true,
            no_signed_zeros: true,
            allow_reciprocal: true,
            allow_contract: true,
            approx_func: true,
        };
        let s = propagate_fmf(&fmf);
        assert!(s.contains("reassoc"));
        assert!(s.contains("nnan"));
        assert!(s.contains("ninf"));
        assert!(s.contains("nsz"));
        assert!(s.contains("arcp"));
        assert!(s.contains("contract"));
        assert!(s.contains("afn"));
    }

    #[test]
    fn test_annotate_fmf_preserves_non_arithmetic() {
        let code = vec![
            "  %x = load <4 x float>, ptr %p".to_string(),
            "  %r = fadd <4 x float> %x, %y".to_string(),
            "  store <4 x float> %r, ptr %q".to_string(),
        ];
        let fmf = FastMathFlags::from_str("reassoc");
        let annotated = annotate_fmf(&code, &fmf);
        assert_eq!(annotated[0], code[0]); // load unchanged
        assert!(annotated[1].ends_with("reassoc")); // fadd annotated
        assert_eq!(annotated[2], code[2]); // store unchanged
    }

    #[test]
    fn test_reduction_plan_vector_type() {
        let plan = plan_tree_reduction(
            ReductionKind::FMul,
            4,
            "double",
            "%v",
            &FastMathFlags::default(),
        );
        assert_eq!(plan.vector_type, "<4 x double>");
        assert_eq!(plan.element_type, "double");
    }

    #[test]
    fn test_tree_reduction_mask_large() {
        let mask = tree_reduction_shuffle_mask(16, 3);
        assert_eq!(mask.len(), 16);
        assert_eq!(mask[0], 8); // offset = 1<<3 = 8
    }

    #[test]
    fn test_no_panic_on_zero_elements() {
        // num_elements = 0 shouldn't be passed, but if it is...
        let mask = tree_reduction_shuffle_mask(0, 0);
        assert!(mask.is_empty());
    }

    #[test]
    fn test_reduction_kind_ordering() {
        // All reduction kinds should be comparable.
        let kinds = vec![
            ReductionKind::IntAdd,
            ReductionKind::FAdd,
            ReductionKind::IntAnd,
        ];
        let mut sorted = kinds.clone();
        sorted.sort();
        assert_eq!(sorted.len(), 3);
    }

    #[test]
    fn test_expander_debug_comments() {
        let mut expander = make_x86_expand_reduction_sse3();
        expander.emit_debug_comments = true;
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fadd",
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        let (code, _) = result.unwrap();
        // First line should be a debug comment.
        assert!(code[0].starts_with(";"));
    }

    #[test]
    fn test_expand_avx2_strategy() {
        let mut expander = make_x86_expand_reduction_avx2();
        // AVX2 has horizontal add for 8-element integer vectors.
        let result = expander.expand_reduction(
            "llvm.vector.reduce.add",
            8,
            "i32",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        assert_eq!(expander.stats.avx_hadd_used, 1);
    }

    #[test]
    fn test_expand_tree_fallback_when_no_horizontal() {
        let mut expander = make_x86_expand_reduction_conservative();
        // Conservative expander has no SSE3, so no horizontal support.
        expander.prefer_horizontal = true; // But SSH3 isn't available.
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fadd",
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        // Should use tree reduction.
        assert_eq!(expander.stats.tree_reductions, 1);
        assert_eq!(expander.stats.horizontal_reductions, 0);
    }

    #[test]
    fn test_generate_x86_pseudo_code_avx() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            8,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let code = generate_x86_pseudo_code(&plan);
        assert!(code.iter().any(|l| l.contains("vperm2f128")));
    }

    #[test]
    fn test_generate_x86_pseudo_code_avx512() {
        let plan = plan_tree_reduction(
            ReductionKind::FAdd,
            16,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        let code = generate_x86_pseudo_code(&plan);
        assert!(code.iter().any(|l| l.contains("vshuff32x4")));
    }

    #[test]
    fn test_all_intrinsic_names_covered() {
        let names = [
            "llvm.vector.reduce.add",
            "llvm.vector.reduce.mul",
            "llvm.vector.reduce.and",
            "llvm.vector.reduce.or",
            "llvm.vector.reduce.xor",
            "llvm.vector.reduce.smax",
            "llvm.vector.reduce.smin",
            "llvm.vector.reduce.umax",
            "llvm.vector.reduce.umin",
            "llvm.vector.reduce.fadd",
            "llvm.vector.reduce.fmul",
            "llvm.vector.reduce.fmax",
            "llvm.vector.reduce.fmin",
        ];
        for name in &names {
            assert!(
                ReductionKind::from_intrinsic_name(name).is_some(),
                "Missing: {}",
                name
            );
        }
    }

    #[test]
    fn test_expander_stats_default_all_zero() {
        let stats = ExpandReductionStats::default();
        assert_eq!(stats.reductions_found, 0);
        assert_eq!(stats.reductions_expanded, 0);
        assert_eq!(stats.tree_reductions, 0);
        assert_eq!(stats.horizontal_reductions, 0);
        assert_eq!(stats.fast_math_applied, 0);
        assert_eq!(stats.fma_contractions, 0);
        assert_eq!(stats.int_reductions, 0);
        assert_eq!(stats.fp_reductions, 0);
        assert_eq!(stats.sse3_hadd_used, 0);
        assert_eq!(stats.avx_hadd_used, 0);
        assert_eq!(stats.total_steps_generated, 0);
        assert_eq!(stats.reductions_too_small, 0);
    }

    #[test]
    fn test_dot_product_reduction_double() {
        let code = generate_dot_product_reduction(2, "double", "%a", "%b");
        assert!(code.iter().any(|l| l.contains("fmul")));
        assert!(code.iter().any(|l| l.contains("haddpd")));
    }

    #[test]
    fn test_dot_product_reduction_float_8() {
        let code = generate_dot_product_reduction(8, "float", "%a", "%b");
        assert!(code.iter().any(|l| l.contains("fmul")));
    }

    #[test]
    fn test_expand_reduction_fmul_4() {
        let mut expander = make_x86_expand_reduction_sse3();
        let result = expander.expand_reduction(
            "llvm.vector.reduce.fmul",
            4,
            "float",
            "%v",
            &FastMathFlags::default(),
        );
        assert!(result.is_some());
        let (code, _) = result.unwrap();
        assert!(code.iter().any(|l| l.contains("fmul")));
    }
}