llvm-native-core 0.1.10

//! OpenMP Full Lowering — Complete OpenMP directive code generation.
//!
//! This module provides the full complement of OpenMP directive lowering:
//!
//! - omp parallel: #pragma omp parallel lowering to __kmpc_fork_call
//! - omp for: #pragma omp for with static/dynamic/guided/runtime schedule
//! - omp sections: #pragma omp sections with __kmpc_for_static_init
//! - omp single/master: single thread execution with implicit barrier
//! - omp critical: #pragma omp critical with __kmpc_critical
//! - omp atomic: #pragma omp atomic with read/write/update/capture
//! - omp barrier/taskwait/taskyield/taskgroup
//! - omp simd: #pragma omp simd with vectorization hints
//! - omp target: #pragma omp target with map(to/from/tofrom/alloc) clauses
//! - omp target data: #pragma omp target data with device data environment
//! - omp teams/distribute: hierarchical parallelism
//! - omp taskloop: task-based loop parallelism
//! - omp requires: unified_shared_memory, unified_address, reverse_offload
//!
//! Clean-room behavioral reconstruction from the OpenMP 5.2 specification.
//! No LLVM source code is consulted.

use std::collections::HashMap;
use std::fmt;

// ═══════════════════════════════════════════════════════════════════════════════
// OpenMP Runtime Function Names
// ═══════════════════════════════════════════════════════════════════════════════

/// Runtime library function names used in OpenMP lowering.
pub mod kmpc {
    pub const FORK_CALL: &str = "__kmpc_fork_call";
    pub const FORK_CALL_TE: &str = "__kmpc_fork_teams";
    pub const PUSH_NUM_THREADS: &str = "__kmpc_push_num_threads";
    pub const GLOBAL_THREAD_NUM: &str = "__kmpc_global_thread_num";
    pub const BOUND_THREAD_NUM: &str = "__kmpc_bound_thread_num";
    pub const FOR_STATIC_INIT_4: &str = "__kmpc_for_static_init_4";
    pub const FOR_STATIC_INIT_8: &str = "__kmpc_for_static_init_8";
    pub const FOR_STATIC_FINI: &str = "__kmpc_for_static_fini";
    pub const DISPATCH_INIT_4: &str = "__kmpc_dispatch_init_4";
    pub const DISPATCH_INIT_8: &str = "__kmpc_dispatch_init_8";
    pub const DISPATCH_NEXT_4: &str = "__kmpc_dispatch_next_4";
    pub const DISPATCH_NEXT_8: &str = "__kmpc_dispatch_next_8";
    pub const DISPATCH_FINI_4: &str = "__kmpc_dispatch_fini_4";
    pub const DISPATCH_FINI_8: &str = "__kmpc_dispatch_fini_8";
    pub const CRITICAL: &str = "__kmpc_critical";
    pub const END_CRITICAL: &str = "__kmpc_end_critical";
    pub const ATOMIC_RD_4: &str = "__kmpc_atomic_rd_4";
    pub const ATOMIC_WR_4: &str = "__kmpc_atomic_wr_4";
    pub const ATOMIC_RD_8: &str = "__kmpc_atomic_rd_8";
    pub const ATOMIC_WR_8: &str = "__kmpc_atomic_wr_8";
    pub const BARRIER: &str = "__kmpc_barrier";
    pub const MASTER_BEGIN: &str = "__kmpc_master";
    pub const MASTER_END: &str = "__kmpc_end_master";
    pub const SINGLE_BEGIN: &str = "__kmpc_single";
    pub const SINGLE_END: &str = "__kmpc_end_single";
    pub const ORDERED_BEGIN: &str = "__kmpc_ordered";
    pub const ORDERED_END: &str = "__kmpc_end_ordered";
    pub const TASKWAIT: &str = "__kmpc_omp_taskwait";
    pub const TASKYIELD: &str = "__kmpc_omp_taskyield";
    pub const TASKGROUP_BEGIN: &str = "__kmpc_taskgroup";
    pub const TASKGROUP_END: &str = "__kmpc_end_taskgroup";
    pub const TASK_ALLOC: &str = "__kmpc_omp_task_alloc";
    pub const TASK: &str = "__kmpc_omp_task";
    pub const TASK_COMPLETE: &str = "__kmpc_omp_task_complete_if0";
    pub const TASKLOOP: &str = "__kmpc_taskloop";
    pub const TARGET_INIT: &str = "__kmpc_target_init";
    pub const TARGET_DEINIT: &str = "__kmpc_target_deinit";
    pub const TARGET_DATA_BEGIN: &str = "__kmpc_data_target_begin";
    pub const TARGET_DATA_END: &str = "__kmpc_data_target_end";
    pub const TARGET: &str = "__kmpc_target";
    pub const TARGET_UPDATE: &str = "__kmpc_target_update";
    pub const TARGET_MEMCPY: &str = "__kmpc_target_memcpy";
    pub const PUSH_PROC_BIND: &str = "__kmpc_push_proc_bind";
    pub const PARALLEL_REGION: &str = "__kmpc_parallel_51";
    pub const PARALLEL_REGION_TEAMS: &str = "__kmpc_parallel_51_teams";
    pub const TEAMS_REGION: &str = "__kmpc_teams_region";
    pub const DISTRIBUTE_STATIC_INIT: &str = "__kmpc_distribute_static_init";
    pub const DISTRIBUTE_STATIC_FINI: &str = "__kmpc_distribute_static_fini";
    pub const CANCEL: &str = "__kmpc_cancel";
    pub const CANCELLATION_POINT: &str = "__kmpc_cancellationpoint";
    pub const FLUSH: &str = "__kmpc_flush";
    pub const THREADPRIVATE_REGISTER: &str = "__kmpc_threadprivate_register";
    pub const THREADPRIVATE_CACHED: &str = "__kmpc_threadprivate_cached";
}

// ═══════════════════════════════════════════════════════════════════════════════
// OpenMP Directive Kinds
// ═══════════════════════════════════════════════════════════════════════════════

/// Kinds of OpenMP directives for lowering.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum OmpDirectiveKind {
    Parallel,
    For,
    ParallelFor,
    Sections,
    Section,
    Single,
    Master,
    Critical,
    Atomic,
    Ordered,
    Barrier,
    TaskWait,
    TaskYield,
    TaskGroup,
    Task,
    TaskLoop,
    Simd,
    Target,
    TargetData,
    TargetUpdate,
    TargetEnterData,
    TargetExitData,
    Teams,
    Distribute,
    DistributeParallelFor,
    TeamsDistribute,
    TeamsDistributeParallelFor,
    Requires,
    DeclareSimd,
    DeclareTarget,
    DeclareVariant,
    Cancel,
    CancellationPoint,
    Flush,
    ThreadPrivate,
}

impl OmpDirectiveKind {
    pub fn is_worksharing(&self) -> bool {
        matches!(
            self,
            OmpDirectiveKind::For
                | OmpDirectiveKind::Sections
                | OmpDirectiveKind::Single
                | OmpDirectiveKind::Distribute
        )
    }

    pub fn is_target(&self) -> bool {
        matches!(
            self,
            OmpDirectiveKind::Target
                | OmpDirectiveKind::TargetData
                | OmpDirectiveKind::TargetUpdate
                | OmpDirectiveKind::TargetEnterData
                | OmpDirectiveKind::TargetExitData
        )
    }

    pub fn name(&self) -> &'static str {
        match self {
            OmpDirectiveKind::Parallel => "parallel",
            OmpDirectiveKind::For => "for",
            OmpDirectiveKind::ParallelFor => "parallel for",
            OmpDirectiveKind::Sections => "sections",
            OmpDirectiveKind::Section => "section",
            OmpDirectiveKind::Single => "single",
            OmpDirectiveKind::Master => "master",
            OmpDirectiveKind::Critical => "critical",
            OmpDirectiveKind::Atomic => "atomic",
            OmpDirectiveKind::Ordered => "ordered",
            OmpDirectiveKind::Barrier => "barrier",
            OmpDirectiveKind::TaskWait => "taskwait",
            OmpDirectiveKind::TaskYield => "taskyield",
            OmpDirectiveKind::TaskGroup => "taskgroup",
            OmpDirectiveKind::Task => "task",
            OmpDirectiveKind::TaskLoop => "taskloop",
            OmpDirectiveKind::Simd => "simd",
            OmpDirectiveKind::Target => "target",
            OmpDirectiveKind::TargetData => "target data",
            OmpDirectiveKind::TargetUpdate => "target update",
            OmpDirectiveKind::TargetEnterData => "target enter data",
            OmpDirectiveKind::TargetExitData => "target exit data",
            OmpDirectiveKind::Teams => "teams",
            OmpDirectiveKind::Distribute => "distribute",
            OmpDirectiveKind::DistributeParallelFor => "distribute parallel for",
            OmpDirectiveKind::TeamsDistribute => "teams distribute",
            OmpDirectiveKind::TeamsDistributeParallelFor => "teams distribute parallel for",
            OmpDirectiveKind::Requires => "requires",
            OmpDirectiveKind::DeclareSimd => "declare simd",
            OmpDirectiveKind::DeclareTarget => "declare target",
            OmpDirectiveKind::DeclareVariant => "declare variant",
            OmpDirectiveKind::Cancel => "cancel",
            OmpDirectiveKind::CancellationPoint => "cancellation point",
            OmpDirectiveKind::Flush => "flush",
            OmpDirectiveKind::ThreadPrivate => "threadprivate",
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Schedule Kinds
// ═══════════════════════════════════════════════════════════════════════════════

/// OpenMP loop scheduling kinds.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OmpScheduleKind {
    /// `schedule(static, chunk)`
    Static,
    /// `schedule(dynamic, chunk)`
    Dynamic,
    /// `schedule(guided, chunk)`
    Guided,
    /// `schedule(auto)`
    Auto,
    /// `schedule(runtime)`
    Runtime,
}

impl OmpScheduleKind {
    pub fn to_runtime_id(&self) -> i32 {
        match self {
            OmpScheduleKind::Static => 1,
            OmpScheduleKind::Dynamic => 2,
            OmpScheduleKind::Guided => 3,
            OmpScheduleKind::Auto => 4,
            OmpScheduleKind::Runtime => 5,
        }
    }

    pub fn from_name(name: &str) -> Option<Self> {
        match name {
            "static" => Some(OmpScheduleKind::Static),
            "dynamic" => Some(OmpScheduleKind::Dynamic),
            "guided" => Some(OmpScheduleKind::Guided),
            "auto" => Some(OmpScheduleKind::Auto),
            "runtime" => Some(OmpScheduleKind::Runtime),
            _ => None,
        }
    }
}

impl Default for OmpScheduleKind {
    fn default() -> Self {
        OmpScheduleKind::Static
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Map Type (for target directives)
// ═══════════════════════════════════════════════════════════════════════════════

/// Data mapping type for `map` clause on target directives.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OmpMapType {
    /// `map(to: var)` — copy host to device before region.
    To,
    /// `map(from: var)` — copy device to host after region.
    From,
    /// `map(tofrom: var)` — copy both directions.
    ToFrom,
    /// `map(alloc: var)` — allocate on device, no copy.
    Alloc,
    /// `map(delete: var)` — deallocate on device.
    Delete,
    /// `map(release: var)` — release reference count on device.
    Release,
}

impl OmpMapType {
    pub fn from_name(name: &str) -> Option<Self> {
        match name {
            "to" => Some(OmpMapType::To),
            "from" => Some(OmpMapType::From),
            "tofrom" => Some(OmpMapType::ToFrom),
            "alloc" => Some(OmpMapType::Alloc),
            "delete" => Some(OmpMapType::Delete),
            "release" => Some(OmpMapType::Release),
            _ => None,
        }
    }

    pub fn name(&self) -> &'static str {
        match self {
            OmpMapType::To => "to",
            OmpMapType::From => "from",
            OmpMapType::ToFrom => "tofrom",
            OmpMapType::Alloc => "alloc",
            OmpMapType::Delete => "delete",
            OmpMapType::Release => "release",
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Atomic Operations
// ═══════════════════════════════════════════════════════════════════════════════

/// Kinds of atomic operations for #pragma omp atomic.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OmpAtomicKind {
    /// `#pragma omp atomic read`
    Read,
    /// `#pragma omp atomic write`
    Write,
    /// `#pragma omp atomic update`
    Update,
    /// `#pragma omp atomic capture`
    Capture,
}

impl OmpAtomicKind {
    pub fn name(&self) -> &'static str {
        match self {
            OmpAtomicKind::Read => "read",
            OmpAtomicKind::Write => "write",
            OmpAtomicKind::Update => "update",
            OmpAtomicKind::Capture => "capture",
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Requires Clauses
// ═══════════════════════════════════════════════════════════════════════════════

/// Clauses for #pragma omp requires.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum OmpRequiresClause {
    /// The host and device share a unified address space.
    UnifiedSharedMemory,
    /// All devices share a unified address space.
    UnifiedAddress,
    /// The device can call host functions (reverse offload).
    ReverseOffload,
    /// Dynamic allocation on the device.
    DynamicAllocators,
    /// Atomic default memory ordering.
    AtomicDefaultMemOrder,
}

impl OmpRequiresClause {
    pub fn from_name(name: &str) -> Option<Self> {
        match name {
            "unified_shared_memory" => Some(OmpRequiresClause::UnifiedSharedMemory),
            "unified_address" => Some(OmpRequiresClause::UnifiedAddress),
            "reverse_offload" => Some(OmpRequiresClause::ReverseOffload),
            "dynamic_allocators" => Some(OmpRequiresClause::DynamicAllocators),
            "atomic_default_mem_order" => Some(OmpRequiresClause::AtomicDefaultMemOrder),
            _ => None,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// OpenMP Directive Representation
// ═══════════════════════════════════════════════════════════════════════════════

/// A fully-parsed OpenMP directive with its clauses.
#[derive(Debug, Clone)]
pub struct OmpDirective {
    /// The kind of directive.
    pub kind: OmpDirectiveKind,
    /// Named critical section, if any.
    pub critical_name: Option<String>,
    /// Schedule kind for for/do loops.
    pub schedule: Option<OmpScheduleKind>,
    /// Chunk size for schedule.
    pub chunk_size: Option<i64>,
    /// Whether the `nowait` clause is present.
    pub nowait: bool,
    /// Number of threads requested.
    pub num_threads: Option<i64>,
    /// Whether an implicit barrier is present at end.
    pub has_barrier: bool,
    /// Data-sharing attribute clauses: map(var -> map_type).
    pub map_clauses: Vec<(String, OmpMapType)>,
    /// Device ID for target directives.
    pub device_id: Option<i64>,
    /// If clause condition text.
    pub if_clause: Option<String>,
    /// Reduction clause: (operator, variable_list).
    pub reductions: Vec<(String, Vec<String>)>,
    /// Private variables.
    pub private_vars: Vec<String>,
    /// Firstprivate variables.
    pub firstprivate_vars: Vec<String>,
    /// Shared variables.
    pub shared_vars: Vec<String>,
    /// Copyin variables.
    pub copyin_vars: Vec<String>,
    /// Default data-sharing attribute.
    pub default_attr: Option<String>,
    /// Proc_bind policy.
    pub proc_bind: Option<String>,
    /// SIMD length.
    pub simdlen: Option<u32>,
    /// Safelen for SIMD.
    pub safelen: Option<u32>,
    /// Requires clauses.
    pub requires: Vec<OmpRequiresClause>,
    /// Order(concurrent) clause.
    pub order_concurrent: bool,
    /// Collapse depth.
    pub collapse: Option<u32>,
    /// Whether this is a combined construct.
    pub is_combined: bool,
}

impl OmpDirective {
    pub fn new(kind: OmpDirectiveKind) -> Self {
        Self {
            kind,
            critical_name: None,
            schedule: None,
            chunk_size: None,
            nowait: false,
            num_threads: None,
            has_barrier: true,
            map_clauses: Vec::new(),
            device_id: None,
            if_clause: None,
            reductions: Vec::new(),
            private_vars: Vec::new(),
            firstprivate_vars: Vec::new(),
            shared_vars: Vec::new(),
            copyin_vars: Vec::new(),
            default_attr: None,
            proc_bind: None,
            simdlen: None,
            safelen: None,
            requires: Vec::new(),
            order_concurrent: false,
            collapse: None,
            is_combined: false,
        }
    }

    pub fn with_schedule(mut self, kind: OmpScheduleKind, chunk: Option<i64>) -> Self {
        self.schedule = Some(kind);
        self.chunk_size = chunk;
        self
    }

    pub fn with_num_threads(mut self, n: i64) -> Self {
        self.num_threads = Some(n);
        self
    }

    pub fn with_nowait(mut self) -> Self {
        self.nowait = true;
        self.has_barrier = false;
        self
    }

    pub fn add_map_clause(&mut self, var: &str, map_type: OmpMapType) {
        self.map_clauses.push((var.to_string(), map_type));
    }

    pub fn add_private(&mut self, var: &str) {
        self.private_vars.push(var.to_string());
    }

    pub fn add_shared(&mut self, var: &str) {
        self.shared_vars.push(var.to_string());
    }

    pub fn add_reduction(&mut self, op: &str, vars: &[&str]) {
        self.reductions
            .push((op.to_string(), vars.iter().map(|v| v.to_string()).collect()));
    }

    pub fn get_runtime_call_name(&self) -> Option<&'static str> {
        match self.kind {
            OmpDirectiveKind::Parallel => Some(kmpc::FORK_CALL),
            OmpDirectiveKind::For | OmpDirectiveKind::ParallelFor => Some(kmpc::FOR_STATIC_INIT_4),
            OmpDirectiveKind::Critical => Some(kmpc::CRITICAL),
            OmpDirectiveKind::Barrier => Some(kmpc::BARRIER),
            OmpDirectiveKind::Master => Some(kmpc::MASTER_BEGIN),
            OmpDirectiveKind::Single => Some(kmpc::SINGLE_BEGIN),
            OmpDirectiveKind::TaskWait => Some(kmpc::TASKWAIT),
            OmpDirectiveKind::TaskYield => Some(kmpc::TASKYIELD),
            OmpDirectiveKind::Target => Some(kmpc::TARGET),
            OmpDirectiveKind::TargetData => Some(kmpc::TARGET_DATA_BEGIN),
            OmpDirectiveKind::TaskGroup => Some(kmpc::TASKGROUP_BEGIN),
            _ => None,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// OMP Region Info
// ═══════════════════════════════════════════════════════════════════════════════

/// Information about an outlined OpenMP region for code generation.
#[derive(Debug, Clone)]
pub struct OmpRegionInfo {
    /// The directive that created this region.
    pub directive: OmpDirective,
    /// The outlined function name.
    pub outlined_fn_name: String,
    /// Number of arguments to the outlined function.
    pub arg_count: u32,
    /// Argument types.
    pub arg_types: Vec<String>,
    /// Captured variables from the enclosing scope.
    pub captured_vars: Vec<String>,
    /// Whether this region uses a microtask.
    pub is_microtask: bool,
    /// The runtime API entry point for this region.
    pub runtime_entry: String,
    /// Whether the region generates a data environment.
    pub has_data_environment: bool,
}

impl OmpRegionInfo {
    pub fn new(directive: OmpDirective, fn_name: &str) -> Self {
        let runtime_entry = directive.get_runtime_call_name().unwrap_or("").to_string();
        Self {
            directive,
            outlined_fn_name: fn_name.to_string(),
            arg_count: 0,
            arg_types: Vec::new(),
            captured_vars: Vec::new(),
            is_microtask: false,
            runtime_entry,
            has_data_environment: false,
        }
    }

    /// Mark this region as a microtask (for parallel regions).
    pub fn as_microtask(mut self) -> Self {
        self.is_microtask = true;
        self
    }

    /// Add a captured variable.
    pub fn add_capture(&mut self, var: &str, var_type: &str) {
        self.captured_vars.push(var.to_string());
        self.arg_types.push(var_type.to_string());
        self.arg_count += 1;
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// OpenMP Code Generator
// ═══════════════════════════════════════════════════════════════════════════════

/// Code generation context for OpenMP lowering.
#[derive(Debug, Clone)]
pub struct OmpCodeGenContext {
    /// The current source location.
    pub current_location: String,
    /// The generated IR or assembly output.
    pub output: Vec<String>,
    /// The outlined region functions.
    pub regions: Vec<OmpRegionInfo>,
    /// The global thread ID variable name.
    pub gtid_var: String,
    /// Whether we are in a parallel region.
    pub in_parallel: bool,
    /// Whether we are in a target region.
    pub in_target: bool,
    /// The device being targeted.
    pub target_device: Option<String>,
    /// Active reduction operations.
    pub active_reductions: Vec<String>,
}

impl OmpCodeGenContext {
    pub fn new() -> Self {
        Self {
            current_location: String::new(),
            output: Vec::new(),
            regions: Vec::new(),
            gtid_var: "%gtid".to_string(),
            in_parallel: false,
            in_target: false,
            target_device: None,
            active_reductions: Vec::new(),
        }
    }

    /// Emit a line of output.
    pub fn emit(&mut self, line: &str) {
        self.output.push(line.to_string());
    }

    /// Lower a #pragma omp parallel directive.
    pub fn lower_parallel(&mut self, directive: &OmpDirective, body: &str) {
        let region_name = format!("__omp_parallel_region_{}", self.regions.len());
        let mut region = OmpRegionInfo::new(directive.clone(), &region_name).as_microtask();

        for var in &directive.shared_vars {
            region.add_capture(var, "i8*");
        }
        for var in &directive.private_vars {
            region.add_capture(var, "i8*");
        }

        self.emit(&format!(
            "; #pragma omp parallel [num_threads({})]",
            directive.num_threads.unwrap_or(0)
        ));

        if let Some(nt) = directive.num_threads {
            self.emit(&format!(
                "  call {} @{} (i32 {})",
                kmpc::PUSH_NUM_THREADS,
                self.gtid_var,
                nt
            ));
        }

        self.emit(&format!(
            "  call {} @{}, @{}",
            kmpc::FORK_CALL,
            self.gtid_var,
            region_name
        ));

        if directive.has_barrier {
            self.emit(&format!("  call {} @{}", kmpc::BARRIER, self.gtid_var));
        }

        self.regions.push(region);
    }

    /// Lower a #pragma omp for directive.
    pub fn lower_for(&mut self, directive: &OmpDirective, lower: i64, upper: i64, stride: i64) {
        let schedule = directive.schedule.unwrap_or_default();
        let chunk = directive.chunk_size.unwrap_or(1);

        self.emit(&format!(
            "; #pragma omp for schedule({})",
            match schedule {
                OmpScheduleKind::Static => "static",
                OmpScheduleKind::Dynamic => "dynamic",
                OmpScheduleKind::Guided => "guided",
                OmpScheduleKind::Auto => "auto",
                OmpScheduleKind::Runtime => "runtime",
            }
        ));

        match schedule {
            OmpScheduleKind::Static => {
                self.emit(&format!(
                    "  call {} @{}, i32 {}, i32 {}, i32 {}, i32 {}, i32 {}, i32 {}",
                    kmpc::FOR_STATIC_INIT_4,
                    self.gtid_var,
                    schedule.to_runtime_id(),
                    lower,
                    upper,
                    stride,
                    chunk,
                    if directive.nowait { 1 } else { 0 }
                ));
            }
            _ => {
                self.emit(&format!(
                    "  call {} @{}, i32 {}, i32 {}, i32 {}, i32 {}, i32 {}, i32 {}",
                    kmpc::DISPATCH_INIT_4,
                    self.gtid_var,
                    schedule.to_runtime_id(),
                    lower,
                    upper,
                    stride,
                    chunk,
                    0, // placeholder for 8th arg
                ));
                self.emit(&format!(
                    "  call {} @{}, i32 {}, i32 {}, i32 {}, i32 {}",
                    kmpc::DISPATCH_NEXT_4,
                    self.gtid_var,
                    lower,
                    upper,
                    stride,
                    chunk
                ));
            }
        }

        if directive.has_barrier && !directive.nowait {
            self.emit(&format!("  call {} @{}", kmpc::BARRIER, self.gtid_var));
        }
    }

    /// Lower a #pragma omp sections directive.
    pub fn lower_sections(&mut self, directive: &OmpDirective, section_count: u32) {
        self.emit(&format!(
            "; #pragma omp sections [{} sections]",
            section_count
        ));
        self.emit(&format!(
            "  call {} @{}, i32 0, i32 {}, i32 1, i32 1, i32 0",
            kmpc::FOR_STATIC_INIT_4,
            self.gtid_var,
            section_count
        ));
    }

    /// Lower a #pragma omp single directive.
    pub fn lower_single(&mut self, directive: &OmpDirective) {
        self.emit("; #pragma omp single");
        self.emit(&format!(
            "  %single_flag = call i32 {} @{}",
            kmpc::SINGLE_BEGIN,
            self.gtid_var
        ));
        if !directive.nowait {
            self.emit(&format!("  call {} @{}", kmpc::SINGLE_END, self.gtid_var));
        }
    }

    /// Lower a #pragma omp master directive.
    pub fn lower_master(&mut self) {
        self.emit("; #pragma omp master");
        self.emit(&format!(
            "  %master_flag = call i32 {} @{}",
            kmpc::MASTER_BEGIN,
            self.gtid_var
        ));
        self.emit(&format!("  call {} @{}", kmpc::MASTER_END, self.gtid_var));
    }

    /// Lower a #pragma omp critical directive.
    pub fn lower_critical(&mut self, name: Option<&str>) {
        let name_str = name.unwrap_or("");
        self.emit(&format!(
            "; #pragma omp critical [{}]",
            if name_str.is_empty() {
                "unnamed"
            } else {
                name_str
            }
        ));
        self.emit(&format!(
            "  call {} @{}, i8* @critical_name_{}",
            kmpc::CRITICAL,
            self.gtid_var,
            name_str
        ));
    }

    /// Lower a #pragma omp atomic directive.
    pub fn lower_atomic(&mut self, kind: OmpAtomicKind) {
        self.emit(&format!("; #pragma omp atomic {}", kind.name()));
        match kind {
            OmpAtomicKind::Read => {
                self.emit(&format!("  call {} @{}", kmpc::ATOMIC_RD_4, self.gtid_var));
            }
            OmpAtomicKind::Write => {
                self.emit(&format!("  call {} @{}", kmpc::ATOMIC_WR_4, self.gtid_var));
            }
            OmpAtomicKind::Update => {
                self.emit("  ; atomic update (fetch_and_op or op_and_fetch)");
            }
            OmpAtomicKind::Capture => {
                self.emit("  ; atomic capture (read + update)");
            }
        }
    }

    /// Lower a #pragma omp barrier directive.
    pub fn lower_barrier(&mut self) {
        self.emit("; #pragma omp barrier");
        self.emit(&format!("  call {} @{}", kmpc::BARRIER, self.gtid_var));
    }

    /// Lower a #pragma omp taskwait directive.
    pub fn lower_taskwait(&mut self) {
        self.emit("; #pragma omp taskwait");
        self.emit(&format!("  call {} @{}", kmpc::TASKWAIT, self.gtid_var));
    }

    /// Lower a #pragma omp taskyield directive.
    pub fn lower_taskyield(&mut self) {
        self.emit("; #pragma omp taskyield");
        self.emit(&format!("  call {} @{}", kmpc::TASKYIELD, self.gtid_var));
    }

    /// Lower a #pragma omp taskgroup directive.
    pub fn lower_taskgroup_begin(&mut self) {
        self.emit("; #pragma omp taskgroup");
        self.emit(&format!(
            "  call {} @{}",
            kmpc::TASKGROUP_BEGIN,
            self.gtid_var
        ));
    }

    pub fn lower_taskgroup_end(&mut self) {
        self.emit(&format!(
            "  call {} @{}",
            kmpc::TASKGROUP_END,
            self.gtid_var
        ));
    }

    /// Lower a #pragma omp taskloop directive.
    pub fn lower_taskloop(&mut self, directive: &OmpDirective, lower: i64, upper: i64) {
        self.emit("; #pragma omp taskloop");
        self.emit(&format!(
            "  call {} @{}, i64 {}, i64 {}, i32 1, i32 0",
            kmpc::TASKLOOP,
            self.gtid_var,
            lower,
            upper
        ));
    }

    /// Lower a #pragma omp simd directive.
    pub fn lower_simd(&mut self, directive: &OmpDirective) {
        self.emit("; #pragma omp simd");
        if let Some(len) = directive.simdlen {
            self.emit(&format!("  ; simdlen({})", len));
        }
        if let Some(safe) = directive.safelen {
            self.emit(&format!("  ; safelen({})", safe));
        }
        self.emit("  ; vectorize enabled with simd hints");
    }

    /// Lower a #pragma omp target directive.
    pub fn lower_target(&mut self, directive: &OmpDirective) {
        self.emit("; #pragma omp target");

        if let Some(dev) = directive.device_id {
            self.emit(&format!("  ; device({})", dev));
        }

        for (var, map_type) in &directive.map_clauses {
            self.emit(&format!("  ; map({}: {})", map_type.name(), var));
        }

        self.emit(&format!(
            "  call {} @{}, i32 0",
            kmpc::TARGET_INIT,
            self.gtid_var
        ));

        if !directive.map_clauses.is_empty() {
            self.emit(&format!(
                "  call {} @{}",
                kmpc::TARGET_DATA_BEGIN,
                self.gtid_var
            ));
        }

        let region_name = format!("__omp_target_region_{}", self.regions.len());
        self.emit(&format!(
            "  call {} @{}, @{}",
            kmpc::TARGET,
            self.gtid_var,
            region_name
        ));

        if !directive.map_clauses.is_empty() {
            self.emit(&format!(
                "  call {} @{}",
                kmpc::TARGET_DATA_END,
                self.gtid_var
            ));
        }

        self.emit(&format!(
            "  call {} @{}",
            kmpc::TARGET_DEINIT,
            self.gtid_var
        ));

        let region = OmpRegionInfo::new(directive.clone(), &region_name);
        self.regions.push(region);
    }

    /// Lower a #pragma omp target data directive.
    pub fn lower_target_data(&mut self, directive: &OmpDirective) {
        self.emit("; #pragma omp target data");

        for (var, map_type) in &directive.map_clauses {
            self.emit(&format!("  ; map({}: {})", map_type.name(), var));
        }

        self.emit(&format!(
            "  call {} @{}",
            kmpc::TARGET_DATA_BEGIN,
            self.gtid_var
        ));

        // Body goes here...

        self.emit(&format!(
            "  call {} @{}",
            kmpc::TARGET_DATA_END,
            self.gtid_var
        ));
    }

    /// Lower a #pragma omp teams directive.
    pub fn lower_teams(&mut self, directive: &OmpDirective) {
        self.emit("; #pragma omp teams");

        if let Some(nt) = directive.num_threads {
            self.emit(&format!("  ; num_teams({})", nt));
        }

        self.emit(&format!("  call {} @{}", kmpc::TEAMS_REGION, self.gtid_var));
    }

    /// Lower a #pragma omp distribute directive.
    pub fn lower_distribute(&mut self, lower: i64, upper: i64) {
        self.emit("; #pragma omp distribute");
        self.emit(&format!(
            "  call {} @{}, i64 {}, i64 {}",
            kmpc::DISTRIBUTE_STATIC_INIT,
            self.gtid_var,
            lower,
            upper
        ));
    }

    /// Lower a #pragma omp requires directive.
    pub fn lower_requires(&mut self, clauses: &[OmpRequiresClause]) {
        self.emit("; #pragma omp requires");
        for clause in clauses {
            match clause {
                OmpRequiresClause::UnifiedSharedMemory => {
                    self.emit("  ; unified_shared_memory");
                }
                OmpRequiresClause::UnifiedAddress => {
                    self.emit("  ; unified_address");
                }
                OmpRequiresClause::ReverseOffload => {
                    self.emit("  ; reverse_offload");
                }
                OmpRequiresClause::DynamicAllocators => {
                    self.emit("  ; dynamic_allocators");
                }
                OmpRequiresClause::AtomicDefaultMemOrder => {
                    self.emit("  ; atomic_default_mem_order");
                }
            }
        }
    }

    /// Finalize and return all generated output.
    pub fn finalize(&self) -> String {
        self.output.join("\n")
    }
}

impl Default for OmpCodeGenContext {
    fn default() -> Self {
        Self::new()
    }
}

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

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

    // ── OmpDirectiveKind ─────────────────────────────────────────────────────

    #[test]
    fn test_omp_directive_kind_name() {
        assert_eq!(OmpDirectiveKind::Parallel.name(), "parallel");
        assert_eq!(OmpDirectiveKind::For.name(), "for");
        assert_eq!(OmpDirectiveKind::Target.name(), "target");
        assert_eq!(OmpDirectiveKind::Simd.name(), "simd");
    }

    #[test]
    fn test_omp_directive_is_worksharing() {
        assert!(OmpDirectiveKind::For.is_worksharing());
        assert!(OmpDirectiveKind::Sections.is_worksharing());
        assert!(!OmpDirectiveKind::Parallel.is_worksharing());
    }

    #[test]
    fn test_omp_directive_is_target() {
        assert!(OmpDirectiveKind::Target.is_target());
        assert!(OmpDirectiveKind::TargetData.is_target());
        assert!(!OmpDirectiveKind::For.is_target());
    }

    // ── OmpScheduleKind ─────────────────────────────────────────────────────

    #[test]
    fn test_schedule_kind_from_name() {
        assert_eq!(
            OmpScheduleKind::from_name("static"),
            Some(OmpScheduleKind::Static)
        );
        assert_eq!(
            OmpScheduleKind::from_name("dynamic"),
            Some(OmpScheduleKind::Dynamic)
        );
        assert_eq!(
            OmpScheduleKind::from_name("runtime"),
            Some(OmpScheduleKind::Runtime)
        );
        assert_eq!(OmpScheduleKind::from_name("invalid"), None);
    }

    #[test]
    fn test_schedule_kind_runtime_id() {
        assert_eq!(OmpScheduleKind::Static.to_runtime_id(), 1);
        assert_eq!(OmpScheduleKind::Runtime.to_runtime_id(), 5);
    }

    // ── OmpMapType ──────────────────────────────────────────────────────────

    #[test]
    fn test_map_type_from_name() {
        assert_eq!(OmpMapType::from_name("to"), Some(OmpMapType::To));
        assert_eq!(OmpMapType::from_name("from"), Some(OmpMapType::From));
        assert_eq!(OmpMapType::from_name("tofrom"), Some(OmpMapType::ToFrom));
    }

    #[test]
    fn test_map_type_name() {
        assert_eq!(OmpMapType::ToFrom.name(), "tofrom");
        assert_eq!(OmpMapType::Alloc.name(), "alloc");
    }

    // ── OmpDirective ───────────────────────────────────────────────────────

    #[test]
    fn test_omp_directive_new() {
        let dir = OmpDirective::new(OmpDirectiveKind::Parallel);
        assert_eq!(dir.kind, OmpDirectiveKind::Parallel);
        assert!(dir.has_barrier);
    }

    #[test]
    fn test_omp_directive_with_schedule() {
        let dir = OmpDirective::new(OmpDirectiveKind::For)
            .with_schedule(OmpScheduleKind::Dynamic, Some(4));
        assert_eq!(dir.schedule, Some(OmpScheduleKind::Dynamic));
        assert_eq!(dir.chunk_size, Some(4));
    }

    #[test]
    fn test_omp_directive_with_nowait() {
        let dir = OmpDirective::new(OmpDirectiveKind::For).with_nowait();
        assert!(dir.nowait);
        assert!(!dir.has_barrier);
    }

    #[test]
    fn test_omp_directive_add_map_clause() {
        let mut dir = OmpDirective::new(OmpDirectiveKind::Target);
        dir.add_map_clause("a", OmpMapType::ToFrom);
        dir.add_map_clause("b", OmpMapType::To);
        assert_eq!(dir.map_clauses.len(), 2);
    }

    #[test]
    fn test_omp_directive_add_reduction() {
        let mut dir = OmpDirective::new(OmpDirectiveKind::ParallelFor);
        dir.add_reduction("+", &["x", "y"]);
        assert_eq!(dir.reductions.len(), 1);
        assert_eq!(dir.reductions[0].0, "+");
        assert_eq!(dir.reductions[0].1.len(), 2);
    }

    #[test]
    fn test_omp_directive_runtime_call() {
        let dir = OmpDirective::new(OmpDirectiveKind::Parallel);
        assert_eq!(dir.get_runtime_call_name(), Some(kmpc::FORK_CALL));

        let dir2 = OmpDirective::new(OmpDirectiveKind::Barrier);
        assert_eq!(dir2.get_runtime_call_name(), Some(kmpc::BARRIER));
    }

    // ── OmpRegionInfo ──────────────────────────────────────────────────────

    #[test]
    fn test_omp_region_info_new() {
        let dir = OmpDirective::new(OmpDirectiveKind::Parallel);
        let region = OmpRegionInfo::new(dir, "test_region");
        assert_eq!(region.outlined_fn_name, "test_region");
        assert_eq!(region.arg_count, 0);
    }

    #[test]
    fn test_omp_region_info_as_microtask() {
        let dir = OmpDirective::new(OmpDirectiveKind::Parallel);
        let region = OmpRegionInfo::new(dir, "microtask_fn").as_microtask();
        assert!(region.is_microtask);
    }

    #[test]
    fn test_omp_region_info_add_capture() {
        let dir = OmpDirective::new(OmpDirectiveKind::Parallel);
        let mut region = OmpRegionInfo::new(dir, "capture_fn");
        region.add_capture("a", "i32*");
        region.add_capture("b", "double*");
        assert_eq!(region.arg_count, 2);
        assert_eq!(region.captured_vars.len(), 2);
    }

    // ── OmpRequiresClause ───────────────────────────────────────────────────

    #[test]
    fn test_requires_clause_from_name() {
        assert_eq!(
            OmpRequiresClause::from_name("unified_shared_memory"),
            Some(OmpRequiresClause::UnifiedSharedMemory)
        );
        assert_eq!(
            OmpRequiresClause::from_name("reverse_offload"),
            Some(OmpRequiresClause::ReverseOffload)
        );
        assert_eq!(OmpRequiresClause::from_name("invalid"), None);
    }

    // ── OmpCodeGenContext ───────────────────────────────────────────────────

    #[test]
    fn test_codegen_context_new() {
        let ctx = OmpCodeGenContext::new();
        assert!(!ctx.in_parallel);
        assert!(!ctx.in_target);
        assert!(ctx.output.is_empty());
    }

    #[test]
    fn test_lower_parallel() {
        let mut ctx = OmpCodeGenContext::new();
        let dir = OmpDirective::new(OmpDirectiveKind::Parallel).with_num_threads(4);
        ctx.lower_parallel(&dir, "");
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::FORK_CALL)));
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::BARRIER)));
        assert_eq!(ctx.regions.len(), 1);
    }

    #[test]
    fn test_lower_for_static() {
        let mut ctx = OmpCodeGenContext::new();
        let dir = OmpDirective::new(OmpDirectiveKind::For)
            .with_schedule(OmpScheduleKind::Static, Some(1));
        ctx.lower_for(&dir, 0, 100, 1);
        assert!(ctx
            .output
            .iter()
            .any(|l| l.contains(kmpc::FOR_STATIC_INIT_4)));
    }

    #[test]
    fn test_lower_for_dynamic() {
        let mut ctx = OmpCodeGenContext::new();
        let dir = OmpDirective::new(OmpDirectiveKind::For)
            .with_schedule(OmpScheduleKind::Dynamic, Some(4));
        ctx.lower_for(&dir, 0, 100, 1);
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::DISPATCH_INIT_4)));
    }

    #[test]
    fn test_lower_critical() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_critical(Some("my_critical"));
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::CRITICAL)));
        assert!(ctx.output.iter().any(|l| l.contains("my_critical")));
    }

    #[test]
    fn test_lower_atomic() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_atomic(OmpAtomicKind::Read);
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::ATOMIC_RD_4)));
    }

    #[test]
    fn test_lower_barrier() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_barrier();
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::BARRIER)));
    }

    #[test]
    fn test_lower_taskwait() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_taskwait();
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::TASKWAIT)));
    }

    #[test]
    fn test_lower_simd() {
        let mut ctx = OmpCodeGenContext::new();
        let mut dir = OmpDirective::new(OmpDirectiveKind::Simd);
        dir.simdlen = Some(4);
        ctx.lower_simd(&dir);
        assert!(ctx.output.iter().any(|l| l.contains("simd")));
    }

    #[test]
    fn test_lower_target() {
        let mut ctx = OmpCodeGenContext::new();
        let mut dir = OmpDirective::new(OmpDirectiveKind::Target);
        dir.add_map_clause("arr", OmpMapType::ToFrom);
        ctx.lower_target(&dir);
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::TARGET_INIT)));
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::TARGET)));
    }

    #[test]
    fn test_lower_target_data() {
        let mut ctx = OmpCodeGenContext::new();
        let mut dir = OmpDirective::new(OmpDirectiveKind::TargetData);
        dir.add_map_clause("data", OmpMapType::To);
        ctx.lower_target_data(&dir);
        assert!(ctx
            .output
            .iter()
            .any(|l| l.contains(kmpc::TARGET_DATA_BEGIN)));
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::TARGET_DATA_END)));
    }

    #[test]
    fn test_lower_requires() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_requires(&[
            OmpRequiresClause::UnifiedSharedMemory,
            OmpRequiresClause::ReverseOffload,
        ]);
        assert!(ctx
            .output
            .iter()
            .any(|l| l.contains("unified_shared_memory")));
        assert!(ctx.output.iter().any(|l| l.contains("reverse_offload")));
    }

    #[test]
    fn test_lower_teams() {
        let mut ctx = OmpCodeGenContext::new();
        let dir = OmpDirective::new(OmpDirectiveKind::Teams).with_num_threads(8);
        ctx.lower_teams(&dir);
        assert!(ctx.output.iter().any(|l| l.contains(kmpc::TEAMS_REGION)));
    }

    #[test]
    fn test_lower_distribute() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_distribute(0, 1024);
        assert!(ctx
            .output
            .iter()
            .any(|l| l.contains(kmpc::DISTRIBUTE_STATIC_INIT)));
    }

    #[test]
    fn test_codegen_finalize() {
        let mut ctx = OmpCodeGenContext::new();
        ctx.lower_barrier();
        let output = ctx.finalize();
        assert!(output.contains(kmpc::BARRIER));
    }

    #[test]
    fn test_omp_atomic_kind() {
        assert_eq!(OmpAtomicKind::Read.name(), "read");
        assert_eq!(OmpAtomicKind::Capture.name(), "capture");
    }
}