llvm-native-core 0.1.2

//! C++20 Coroutine Support — implements coroutine lowering, frame layout,
//! promise type interaction, and symmetric transfer.
//!
//! Covers the C++20 coroutine TS (Technical Specification) model:
//! - co_await, co_yield, co_return
//! - Coroutine frame allocation and layout
//! - Promise type interaction (initial_suspend, final_suspend, etc.)
//! - Symmetric transfer optimization
//! - Awaitable protocol (await_ready, await_suspend, await_resume)
//!
//! Clean-room behavioral reconstruction from:
//! - C++20 Standard §9.5.4 (Coroutines)
//! - Published LLVM coroutine documentation
//! - No LLVM/Clang source code is consulted.

use std::collections::HashMap;

use super::cpp_ast::*;

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine State Machine States
// ═══════════════════════════════════════════════════════════════════════════════

/// The state of a coroutine at any point in execution.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CoroutineState {
    /// Initial state: before initial_suspend.
    Initial,
    /// Suspended at initial_suspend point.
    InitialSuspended,
    /// Running (resumed after initial suspend).
    Running,
    /// Suspended at a co_await / co_yield point.
    Suspended,
    /// At final_suspend but not yet destroyed.
    FinalSuspended,
    /// Done (destroyed or at end).
    Done,
}

impl CoroutineState {
    pub fn as_int(&self) -> u32 {
        match self {
            CoroutineState::Initial => 0,
            CoroutineState::InitialSuspended => 1,
            CoroutineState::Running => 2,
            CoroutineState::Suspended => 3,
            CoroutineState::FinalSuspended => 4,
            CoroutineState::Done => 5,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Awaitable Protocol
// ═══════════════════════════════════════════════════════════════════════════════

/// The awaitable interface: `await_ready()`, `await_suspend()`, `await_resume()`.
#[derive(Debug, Clone)]
pub struct AwaitableDesc {
    /// Type of the awaitable expression.
    pub awaitable_type: String,
    /// Whether await_ready returns bool (as required).
    pub has_await_ready: bool,
    /// Whether await_suspend exists.
    pub has_await_suspend: bool,
    /// Whether await_resume exists.
    pub has_await_resume: bool,
    /// Whether this is a symmetric transfer awaitable.
    pub is_symmetric_transfer: bool,
}

impl AwaitableDesc {
    pub fn new(awaitable_type: &str) -> Self {
        Self {
            awaitable_type: awaitable_type.to_string(),
            has_await_ready: true,
            has_await_suspend: true,
            has_await_resume: true,
            is_symmetric_transfer: false,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Promise Type Descriptor
// ═══════════════════════════════════════════════════════════════════════════════

/// Describes the coroutine promise type interface.
#[derive(Debug, Clone)]
pub struct PromiseTypeDesc {
    /// The fully-qualified name of the promise type.
    pub type_name: String,
    /// Whether get_return_object is defined.
    pub has_get_return_object: bool,
    /// Whether initial_suspend is defined.
    pub has_initial_suspend: bool,
    /// Whether final_suspend is defined.
    pub has_final_suspend: bool,
    /// Whether unhandled_exception is defined.
    pub has_unhandled_exception: bool,
    /// Whether return_void is defined (for void coroutines).
    pub has_return_void: bool,
    /// Whether return_value is defined (for non-void coroutines).
    pub has_return_value: bool,
    /// Whether yield_value is defined.
    pub has_yield_value: bool,
    /// The type returned by get_return_object.
    pub return_object_type: String,
    /// The awaitable type returned by initial_suspend.
    pub initial_suspend_awaitable: Option<AwaitableDesc>,
    /// The awaitable type returned by final_suspend.
    pub final_suspend_awaitable: Option<AwaitableDesc>,
    /// Whether the allocator overload is defined.
    pub has_allocator: bool,
    /// Whether the deallocator is defined.
    pub has_deallocator: bool,
}

impl PromiseTypeDesc {
    pub fn new(type_name: &str, return_object_type: &str) -> Self {
        Self {
            type_name: type_name.to_string(),
            has_get_return_object: true,
            has_initial_suspend: true,
            has_final_suspend: true,
            has_unhandled_exception: true,
            has_return_void: true,
            has_return_value: false,
            has_yield_value: false,
            return_object_type: return_object_type.to_string(),
            initial_suspend_awaitable: Some(AwaitableDesc::new("std::suspend_always")),
            final_suspend_awaitable: Some(AwaitableDesc::new("std::suspend_always")),
            has_allocator: true,
            has_deallocator: true,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Frame Layout
// ═══════════════════════════════════════════════════════════════════════════════

/// The layout of a coroutine frame in memory.
#[derive(Debug, Clone)]
pub struct CoroutineFrame {
    /// The promise object (at a fixed offset).
    pub promise_type: String,
    /// Offset of the promise in the frame.
    pub promise_offset: usize,
    /// Saved local variables stored in the frame.
    pub locals: Vec<FrameLocal>,
    /// The resume function pointer.
    pub resume_fn_ptr_offset: usize,
    /// The destroy function pointer.
    pub destroy_fn_ptr_offset: usize,
    /// The current coroutine state index.
    pub state_index_offset: usize,
    /// Total frame size in bytes.
    pub frame_size: usize,
    /// Frame alignment.
    pub frame_alignment: usize,
}

/// A local variable stored in the coroutine frame.
#[derive(Debug, Clone)]
pub struct FrameLocal {
    /// Name of the original local.
    pub name: String,
    /// LLVM type.
    pub llvm_type: String,
    /// Offset in the frame.
    pub offset: usize,
    /// Size in bytes.
    pub size: usize,
    /// Whether this local spans a suspend point.
    pub spans_suspend: bool,
    /// Whether this local needs destruction.
    pub needs_destruction: bool,
}

impl CoroutineFrame {
    pub fn new(promise_type: &str) -> Self {
        Self {
            promise_type: promise_type.to_string(),
            promise_offset: 0,
            locals: Vec::new(),
            resume_fn_ptr_offset: 0,
            destroy_fn_ptr_offset: 8,
            state_index_offset: 16,
            frame_size: 32, // minimum header + pointer fields
            frame_alignment: 8,
        }
    }

    /// Add a local variable to the frame.
    pub fn add_local(&mut self, name: &str, llvm_type: &str, size: usize, spans_suspend: bool) {
        let offset = self.frame_size;
        self.frame_size += (size + self.frame_alignment - 1) & !(self.frame_alignment - 1);
        self.locals.push(FrameLocal {
            name: name.to_string(),
            llvm_type: llvm_type.to_string(),
            offset,
            size,
            spans_suspend,
            needs_destruction: false,
        });
    }

    /// Generate LLVM IR for the coroutine frame type.
    pub fn to_ir_type(&self) -> String {
        let mut ir = String::new();
        ir.push_str(&format!(
            "%coro.frame.ty = type {{ i32, ; state\n  ptr, ; resume fn\n  ptr, ; destroy fn\n"
        ));
        ir.push_str(&format!("  {}, ; promise\n", self.promise_type));
        for local in &self.locals {
            ir.push_str(&format!("  {}, ; {}\n", local.llvm_type, local.name));
        }
        ir.push_str("}\n");
        ir
    }

    /// Generate the LLVM IR for allocating the frame.
    pub fn gen_alloc_ir(&self, coro_name: &str) -> String {
        format!(
            "  %frame.ptr{} = call ptr @llvm.coro.begin(ptr null, i32 0, ptr null, ptr null)\n",
            coro_name
        )
    }

    /// Generate the LLVM IR for freeing the frame.
    pub fn gen_free_ir(&self, coro_name: &str) -> String {
        format!(
            "  call void @llvm.coro.free(ptr %frame.ptr{}, ptr null)\n",
            coro_name
        )
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Lowering
// ═══════════════════════════════════════════════════════════════════════════════

/// Generates LLVM IR for coroutine constructs.
pub struct CoroutineLowering {
    /// The promise type description.
    promise: PromiseTypeDesc,
    /// The coroutine frame layout.
    frame: CoroutineFrame,
    /// Suspend point counter.
    suspend_point_counter: u32,
    /// The current coroutine function name.
    coroutine_name: String,
}

impl CoroutineLowering {
    pub fn new(promise: PromiseTypeDesc, coroutine_name: &str) -> Self {
        Self {
            frame: CoroutineFrame::new(&promise.type_name),
            promise,
            suspend_point_counter: 0,
            coroutine_name: coroutine_name.to_string(),
        }
    }

    /// Get the next suspend point ID.
    fn next_suspend_point(&mut self) -> u32 {
        let id = self.suspend_point_counter;
        self.suspend_point_counter += 1;
        id
    }

    /// Generate the coroutine entry/allocation IR.
    pub fn gen_coroutine_entry_ir(&self) -> String {
        let mut ir = String::new();
        ir.push_str(&format!("; Coroutine '{}' entry\n", self.coroutine_name));

        // Allocate frame
        ir.push_str(&format!(
            "  %frame = call ptr @llvm.coro.begin(ptr null, i32 0, ptr null, ptr null)\n"
        ));

        // Create promise
        ir.push_str(&format!(
            "  %promise = getelementptr inbounds {}, ptr %frame, i32 0, i32 3\n",
            self.frame.promise_type
        ));
        ir.push_str(&format!(
            "  call void @llvm.coro.promise(ptr %promise, i32 {}, i1 false)\n",
            self.frame.promise_alignment()
        ));

        // Initialize promise
        if self.promise.has_get_return_object {
            ir.push_str("  ; promise.get_return_object()\n");
        }
        if self.promise.has_initial_suspend {
            ir.push_str("  ; initial_suspend\n");
        }

        ir
    }

    /// Generate IR for a `co_await expr` expression.
    pub fn gen_co_await_ir(&mut self, expr: &str) -> String {
        let id = self.next_suspend_point();
        let mut ir = String::new();

        ir.push_str(&format!("; co_await {} (suspend point {})\n", expr, id));

        // save state
        ir.push_str(&format!(
            "  %state.ptr{} = getelementptr inbounds i32, ptr %frame, i32 0, i32 0\n",
            id
        ));
        ir.push_str(&format!(
            "  store i32 {}, ptr %state.ptr{}\n",
            CoroutineState::Suspended.as_int(),
            id
        ));

        // save coroutine frame (if symmetric transfer)
        ir.push_str(&format!(
            "  %save{} = call i1 @llvm.coro.save(ptr %frame)\n",
            id
        ));
        ir.push_str(&format!(
            "  %suspend{} = call i8 @llvm.coro.suspend(ptr %frame, i1 false)\n",
            id
        ));

        // switch on suspend result
        ir.push_str(&format!(
            "  switch i8 %suspend{}, label %%coro.suspend{} [\n",
            id, id
        ));
        ir.push_str("    i8 0, label %coro.ret\n");
        ir.push_str("    i8 1, label %coro.cont\n");
        ir.push_str("  ]\n\n");

        // Resume point
        ir.push_str(&format!("coro.resume{}:\n", id));
        ir.push_str("  ; await_resume() called here\n");
        ir.push_str(&format!("  br label %%coro.cont{}\n\n", id));

        ir.push_str(&format!("coro.cont{}:\n", id));

        ir
    }

    /// Generate IR for a `co_yield expr` expression.
    pub fn gen_co_yield_ir(&mut self, expr: &str) -> String {
        let id = self.next_suspend_point();
        let mut ir = String::new();

        ir.push_str(&format!("; co_yield {} (suspend point {})\n", expr, id));

        // Call promise.yield_value(expr)
        if self.promise.has_yield_value {
            ir.push_str(&format!("  ; promise.yield_value({})\n", expr));
        }

        // Save state
        ir.push_str(&format!(
            "  store i32 {}, ptr %%state.ptr{}\n",
            CoroutineState::Suspended.as_int(),
            id
        ));

        // Suspend
        ir.push_str(&format!(
            "  %suspend_y{} = call i8 @llvm.coro.suspend(ptr %%frame, i1 false)\n",
            id
        ));
        ir.push_str(&format!(
            "  switch i8 %suspend_y{}, label %%coro.suspend{} [\n",
            id, id
        ));
        ir.push_str("    i8 0, label %coro.ret\n");
        ir.push_str(&format!("    i8 1, label %%coro.resume_y{}\n", id));
        ir.push_str("  ]\n\n");

        ir.push_str(&format!("coro.resume_y{}:\n", id));
        ir
    }

    /// Generate IR for `co_return` or `co_return expr`.
    pub fn gen_co_return_ir(&mut self, expr: Option<&str>) -> String {
        let mut ir = String::new();
        ir.push_str("; co_return\n");

        // Call promise.return_void() or promise.return_value(expr)
        if let Some(val) = expr {
            if self.promise.has_return_value {
                ir.push_str(&format!("  ; promise.return_value({})\n", val));
            }
        } else if self.promise.has_return_void {
            ir.push_str("  ; promise.return_void()\n");
        }

        // Go to final suspend
        ir.push_str("  br label %coro.final\n");
        ir
    }

    /// Generate the coroutine finalization sequence.
    pub fn gen_coroutine_final_ir(&self) -> String {
        let mut ir = String::new();

        ir.push_str("coro.final:\n");
        ir.push_str("  ; final_suspend\n");

        // Save state as final suspended
        ir.push_str(&format!(
            "  store i32 {}, ptr %state.ptr\n",
            CoroutineState::FinalSuspended.as_int()
        ));

        // Final suspend
        ir.push_str("  %save_final = call i1 @llvm.coro.save(ptr %frame)\n");
        ir.push_str("  %suspend_final = call i8 @llvm.coro.suspend(ptr %frame, i1 true)\n");

        // Cleanup
        ir.push_str("coro.cleanup:\n");
        ir.push_str("  ; destroy frame\n");
        ir.push_str("  call void @llvm.coro.destroy(ptr %frame)\n");
        ir.push_str("  ret void\n");

        // Suspend cleanup
        ir.push_str("\ncoro.suspend:\n");
        ir.push_str("  call void @llvm.coro.destroy(ptr %frame)\n");
        ir.push_str("  ret void\n");

        ir
    }

    /// Generate the coroutine end marker.
    pub fn gen_coroutine_end_ir(&self) -> String {
        format!("  %end = call i1 @llvm.coro.end(ptr %frame, i1 false, ptr null)\n")
    }

    /// Generate declarations for the LLVM coroutine intrinsics.
    pub fn gen_intrinsic_decls_ir() -> &'static str {
        r#"declare ptr @llvm.coro.begin(ptr, i32, ptr, ptr)
declare ptr @llvm.coro.free(ptr, ptr)
declare i1 @llvm.coro.end(ptr, i1, ptr)
declare i8 @llvm.coro.suspend(ptr, i1)
declare i1 @llvm.coro.save(ptr)
declare ptr @llvm.coro.promise(ptr, i32, i1)
declare void @llvm.coro.destroy(ptr)
declare ptr @llvm.coro.frame()
declare i1 @llvm.coro.done(ptr)
declare ptr @llvm.coro.subfn.addr(ptr, i8)
"#
    }
}

impl CoroutineFrame {
    fn promise_alignment(&self) -> usize {
        self.frame_alignment
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Symmetric Transfer
// ═══════════════════════════════════════════════════════════════════════════════

/// Information about a symmetric transfer between coroutines.
#[derive(Debug, Clone)]
pub struct SymmetricTransfer {
    /// The source coroutine.
    pub source: String,
    /// The target coroutine.
    pub target: String,
    /// Whether this is a tail-call transfer.
    pub is_tail_call: bool,
}

impl SymmetricTransfer {
    /// Generate IR for a symmetric transfer from source to target coroutine.
    pub fn gen_ir(&self) -> String {
        format!(
            r#"; Symmetric transfer: {} -> {}
  %target.handle = call ptr @llvm.coro.subfn.addr(ptr %frame.target, i8 0)
  musttail call fastcc void %target.handle(ptr %frame.target)
  ret void
"#,
            self.source, self.target
        )
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// co_await / co_yield / co_return Parsing and Lowering
// ═══════════════════════════════════════════════════════════════════════════════

/// Recognized coroutine keyword.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CoroutineKeyword {
    CoAwait,
    CoYield,
    CoReturn,
}

/// A parsed coroutine expression.
#[derive(Debug, Clone)]
pub struct CoroutineExpr {
    /// Which keyword was used.
    pub keyword: CoroutineKeyword,
    /// The operand expression (None for bare `co_return;`).
    pub operand: Option<String>,
    /// The source location.
    pub source_loc: usize,
}

/// Parse a coroutine keyword expression from source text.
pub fn parse_coroutine_expr(source: &str) -> Option<CoroutineExpr> {
    let trimmed = source.trim();

    if let Some(rest) = trimmed.strip_prefix("co_await ") {
        Some(CoroutineExpr {
            keyword: CoroutineKeyword::CoAwait,
            operand: Some(rest.trim_end_matches(';').to_string()),
            source_loc: 0,
        })
    } else if let Some(rest) = trimmed.strip_prefix("co_yield ") {
        Some(CoroutineExpr {
            keyword: CoroutineKeyword::CoYield,
            operand: Some(rest.trim_end_matches(';').to_string()),
            source_loc: 0,
        })
    } else if trimmed == "co_return;" {
        Some(CoroutineExpr {
            keyword: CoroutineKeyword::CoReturn,
            operand: None,
            source_loc: 0,
        })
    } else if let Some(rest) = trimmed.strip_prefix("co_return ") {
        Some(CoroutineExpr {
            keyword: CoroutineKeyword::CoReturn,
            operand: Some(rest.trim_end_matches(';').to_string()),
            source_loc: 0,
        })
    } else {
        None
    }
}

/// Lower a coroutine expression to LLVM IR.
pub fn lower_coroutine_expr(expr: &CoroutineExpr, lowering: &mut CoroutineLowering) -> String {
    match expr.keyword {
        CoroutineKeyword::CoAwait => {
            let operand = expr.operand.as_deref().unwrap_or("unknown");
            lowering.gen_co_await_ir(operand)
        }
        CoroutineKeyword::CoYield => {
            let operand = expr.operand.as_deref().unwrap_or("unknown");
            lowering.gen_co_yield_ir(operand)
        }
        CoroutineKeyword::CoReturn => lowering.gen_co_return_ir(expr.operand.as_deref()),
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Promise Type Discovery
// ═══════════════════════════════════════════════════════════════════════════════

/// Discovers the promise type from a coroutine return type.
/// Per C++20 §9.5.4.1: the promise type is found via
/// `std::coroutine_traits<ReturnType, Args...>::promise_type`.
pub struct PromiseTypeDiscovery {
    /// The coroutine return type.
    pub return_type: String,
    /// The discovered promise type name.
    pub promise_type: Option<String>,
    /// The return object type.
    pub return_object_type: Option<String>,
}

impl PromiseTypeDiscovery {
    pub fn new(return_type: &str) -> Self {
        Self {
            return_type: return_type.to_string(),
            promise_type: None,
            return_object_type: None,
        }
    }

    /// Attempt to discover the promise type from the return type.
    /// Returns `true` if a promise type was found.
    pub fn discover(&mut self) -> bool {
        // Convention: for `generator<T>`, promise_type is `generator<T>::promise_type`
        // For `task<T>`, promise_type is `task<T>::promise_type`
        let promise_candidate = format!("{}::promise_type", self.return_type);
        self.promise_type = Some(promise_candidate.clone());
        self.return_object_type = Some(self.return_type.clone());
        true
    }

    /// Try to find `get_return_object` return type.
    pub fn discover_return_object(&self) -> Option<String> {
        self.return_object_type.clone()
    }

    /// Build a PromiseTypeDesc from the discovered types.
    pub fn to_promise_desc(&self) -> Option<PromiseTypeDesc> {
        let type_name = self.promise_type.as_ref()?;
        let ret_obj = self
            .discover_return_object()
            .unwrap_or_else(|| "unknown".to_string());
        Some(PromiseTypeDesc::new(type_name, &ret_obj))
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Awaitable Protocol — Full Implementation
// ═══════════════════════════════════════════════════════════════════════════════

/// The three methods of the awaitable protocol.
#[derive(Debug, Clone)]
pub struct AwaitableProtocol {
    /// The awaitable type.
    pub awaitable_type: String,
    /// `bool await_ready()` → whether the awaitable is immediately ready.
    pub await_ready_returns: bool,
    /// `void/void* await_suspend(coroutine_handle<>)` → what happens on suspend.
    pub await_suspend_kind: AwaitSuspendKind,
    /// `T await_resume()` → the resumed value type.
    pub await_resume_type: String,
}

/// The kind of `await_suspend`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AwaitSuspendKind {
    /// `void await_suspend(coroutine_handle<>)` — returns void.
    ReturnsVoid,
    /// `bool await_suspend(coroutine_handle<>)` — returns bool (custom reschedule).
    ReturnsBool,
    /// `coroutine_handle<> await_suspend(coroutine_handle<>)` — symmetric transfer.
    SymmetricTransfer,
}

impl AwaitableProtocol {
    pub fn new(awaitable_type: &str) -> Self {
        Self {
            awaitable_type: awaitable_type.to_string(),
            await_ready_returns: false,
            await_suspend_kind: AwaitSuspendKind::ReturnsVoid,
            await_resume_type: "void".to_string(),
        }
    }

    /// Check if this is a noexcept awaitable.
    pub fn is_noexcept(&self) -> bool {
        // All three methods should be noexcept for best optimization
        true
    }

    /// Generate IR for `await_ready()` check.
    pub fn gen_await_ready_ir(&self, awaitable_var: &str) -> String {
        format!(
            "  %ready = call i1 @_ZNK{}11await_readyEv(ptr {})\n",
            self.awaitable_type.replace("::", ""),
            awaitable_var
        )
    }

    /// Generate IR for `await_suspend()` call.
    pub fn gen_await_suspend_ir(&self, awaitable_var: &str, handle: &str) -> String {
        match self.await_suspend_kind {
            AwaitSuspendKind::ReturnsVoid => format!(
                "  call void @_ZNK{}13await_suspendEv(ptr {}, ptr {})\n",
                self.awaitable_type.replace("::", ""),
                awaitable_var,
                handle
            ),
            AwaitSuspendKind::ReturnsBool => format!(
                "  %should_suspend = call i1 @_ZNK{}13await_suspendEv(ptr {}, ptr {})\n",
                self.awaitable_type.replace("::", ""),
                awaitable_var,
                handle
            ),
            AwaitSuspendKind::SymmetricTransfer => format!(
                "  %target_handle = call ptr @_ZNK{}13await_suspendEv(ptr {}, ptr {})\n",
                self.awaitable_type.replace("::", ""),
                awaitable_var,
                handle
            ),
        }
    }

    /// Generate IR for `await_resume()`.
    pub fn gen_await_resume_ir(&self, awaitable_var: &str) -> String {
        format!(
            "  %resume_val = call {} @_ZNK{}12await_resumeEv(ptr {})\n",
            self.await_resume_type,
            self.awaitable_type.replace("::", ""),
            awaitable_var
        )
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Frame Layout: Promise, Params, Locals, Spill Slots
// ═══════════════════════════════════════════════════════════════════════════════

/// Extended coroutine frame layout including spill slots.
#[derive(Debug, Clone)]
pub struct ExtendedCoroutineFrame {
    /// Base frame info.
    pub base: CoroutineFrame,
    /// ABI-preserved register spill slots.
    pub spill_slots: Vec<SpillSlot>,
    /// The coroutine parameters saved in the frame.
    pub saved_params: Vec<FrameLocal>,
    /// Total frame size including spill slots.
    pub total_frame_size: usize,
}

/// A register spill slot in the coroutine frame.
#[derive(Debug, Clone)]
pub struct SpillSlot {
    /// The register name being spilled.
    pub register: String,
    /// Offset in the frame.
    pub offset: usize,
    /// Size in bytes.
    pub size: usize,
}

impl ExtendedCoroutineFrame {
    pub fn new(promise_type: &str) -> Self {
        Self {
            base: CoroutineFrame::new(promise_type),
            spill_slots: Vec::new(),
            saved_params: Vec::new(),
            total_frame_size: 32, // matches base
        }
    }

    /// Add a spill slot for a register.
    pub fn add_spill_slot(&mut self, register: &str, size: usize) {
        let offset = self.total_frame_size;
        self.total_frame_size += size;
        self.spill_slots.push(SpillSlot {
            register: register.to_string(),
            offset,
            size,
        });
    }

    /// Add a saved parameter to the frame.
    pub fn add_saved_param(&mut self, name: &str, size: usize) {
        let offset = self.total_frame_size;
        self.total_frame_size += size;
        self.saved_params.push(FrameLocal {
            name: name.to_string(),
            llvm_type: format!("i{}", size * 8),
            offset,
            size,
            spans_suspend: true,
            needs_destruction: false,
        });
    }

    /// Generate full frame layout including spill slots.
    pub fn gen_frame_layout_ir(&self) -> String {
        let mut ir = self.base.to_ir_type();
        if !self.spill_slots.is_empty() {
            ir.push_str("; Spill slots:\n");
            for slot in &self.spill_slots {
                ir.push_str(&format!(
                    ";   {} at offset {} ({} bytes)\n",
                    slot.register, slot.offset, slot.size
                ));
            }
        }
        ir
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Splitting: Ramp, Resume, Destroy Functions
// ═══════════════════════════════════════════════════════════════════════════════

/// A split coroutine with three functions: ramp (alloc+init+suspend),
/// resume, and destroy.
#[derive(Debug, Clone)]
pub struct SplitCoroutine {
    /// The coroutine name.
    pub name: String,
    /// The ramp function (initial creation).
    pub ramp_fn: String,
    /// The resume function.
    pub resume_fn: String,
    /// The destroy function.
    pub destroy_fn: String,
    /// The cleanup function (called during unwinding).
    pub cleanup_fn: String,
}

impl SplitCoroutine {
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            ramp_fn: format!("{}_ramp", name),
            resume_fn: format!("{}_resume", name),
            destroy_fn: format!("{}_destroy", name),
            cleanup_fn: format!("{}_cleanup", name),
        }
    }

    /// Generate the ramp function IR.
    pub fn gen_ramp_fn_ir(
        &self,
        frame: &ExtendedCoroutineFrame,
        promise: &PromiseTypeDesc,
    ) -> String {
        let mut ir = String::new();
        ir.push_str(&format!(
            "define ptr @{}(ptr %%coro.promise) {{\n",
            self.ramp_fn
        ));
        ir.push_str("entry:\n");
        // Allocate frame
        ir.push_str(&format!(
            "  %%frame = call ptr @llvm.coro.begin(ptr null, i32 0, ptr null, ptr null)\n"
        ));
        // Initialize promise
        ir.push_str(&format!(
            "  %%promise.ptr = getelementptr inbounds {}, ptr %%frame, i32 0, i32 3\n",
            frame.base.promise_type
        ));
        if promise.has_get_return_object {
            ir.push_str("  %return_obj = call ptr @promise_get_return_object(ptr %promise.ptr)\n");
        }
        if promise.has_initial_suspend {
            ir.push_str("  ; initial_suspend\n");
            ir.push_str("  call i8 @llvm.coro.suspend(ptr %frame, i1 false)\n");
        }
        ir.push_str("  ret ptr %return_obj\n");
        ir.push_str("}\n\n");
        ir
    }

    /// Generate the resume function IR.
    pub fn gen_resume_fn_ir(&self, lowerer: &CoroutineLowering) -> String {
        let mut ir = String::new();
        ir.push_str(&format!(
            "define void @{}(ptr %%frame) {{\n",
            self.resume_fn
        ));
        ir.push_str("entry:\n");
        ir.push_str(&format!("  {}_resume_entry:\n", self.name));
        ir.push_str(&lowerer.gen_coroutine_end_ir());
        ir.push_str("  ret void\n");
        ir.push_str("}\n\n");
        ir
    }

    /// Generate the destroy function IR.
    pub fn gen_destroy_fn_ir(&self) -> String {
        format!(
            "define void @{}(ptr %frame) {{\nentry:\n  call void @llvm.coro.destroy(ptr %frame)\n  ret void\n}}\n",
            self.destroy_fn
        )
    }

    /// Generate the cleanup function IR (for exception unwinding).
    pub fn gen_cleanup_fn_ir(&self, has_exceptions: bool) -> String {
        let mut ir = String::new();
        ir.push_str(&format!(
            "define void @{}(ptr %%frame) {{\n",
            self.cleanup_fn
        ));
        ir.push_str("entry:\n");
        if has_exceptions {
            ir.push_str("  ; destroy locals with destructors\n");
        }
        ir.push_str("  call void @llvm.coro.destroy(ptr %frame)\n");
        ir.push_str("  ret void\n");
        ir.push_str("}\n");
        ir
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Symmetric Transfer — Tail-Calling Between Coroutines
// ═══════════════════════════════════════════════════════════════════════════════

/// Extended symmetric transfer supporting arbitrary target coroutines.
pub struct SymmetricTransferBuilder {
    /// The source frame address.
    pub source_frame: String,
    /// The target handle.
    pub target_handle: String,
    /// Whether this is a musttail call.
    pub is_musttail: bool,
    /// Spill slots to restore before transfer.
    pub spill_restores: Vec<String>,
}

impl SymmetricTransferBuilder {
    pub fn new(source: &str, target: &str) -> Self {
        Self {
            source_frame: source.to_string(),
            target_handle: target.to_string(),
            is_musttail: true,
            spill_restores: Vec::new(),
        }
    }

    /// Build the symmetric transfer with spill restoration.
    pub fn build(&self) -> SymmetricTransfer {
        SymmetricTransfer {
            source: self.source_frame.clone(),
            target: self.target_handle.clone(),
            is_tail_call: self.is_musttail,
        }
    }

    /// Generate the full IR for symmetric transfer with spill slots.
    pub fn gen_full_ir(&self, spill_slots: &[SpillSlot]) -> String {
        let mut ir = format!(
            "; Full symmetric transfer: {} -> {}\n",
            self.source_frame, self.target_handle
        );

        // Restore spill slots
        for slot in spill_slots {
            ir.push_str(&format!(
                "  %restored_{} = load i{}, ptr getelementptr inbounds (i8, ptr %frame, i64 {})\n",
                slot.register,
                slot.size * 8,
                slot.offset
            ));
        }

        // Tail call target
        ir.push_str(&format!(
            "  %target.fn = call ptr @llvm.coro.subfn.addr(ptr {}, i8 0)\n",
            self.target_handle
        ));
        ir.push_str(&format!(
            "  musttail call fastcc void %target.fn(ptr {})\n",
            self.target_handle
        ));
        ir.push_str("  ret void\n");
        ir
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// HALO (Heap Allocation eLision Optimization)
// ═══════════════════════════════════════════════════════════════════════════════

/// HALO analysis: determines whether a coroutine frame can be allocated
/// on the caller's stack instead of the heap.
#[derive(Debug, Clone)]
pub struct HaloAnalysis {
    /// Whether the coroutine can be stack-allocated.
    pub can_elide_heap: bool,
    /// Whether the frame never escapes the caller.
    pub frame_does_not_escape: bool,
    /// The maximum live range of the coroutine (in basic blocks).
    pub max_live_range: usize,
    /// Whether a heap allocation is required (e.g., for cross-thread use).
    pub requires_heap: bool,
}

impl HaloAnalysis {
    pub fn new() -> Self {
        Self {
            can_elide_heap: false,
            frame_does_not_escape: true,
            max_live_range: 0,
            requires_heap: false,
        }
    }

    /// Analyze whether heap allocation can be elided.
    /// Returns `true` if the coroutine can go on the stack.
    pub fn analyze(
        &mut self,
        frame_is_escaped: bool,
        is_always_inline: bool,
        is_cross_thread: bool,
    ) -> bool {
        // Conditions for HALO:
        // 1. The frame does not escape to another coroutine/fiber/thread
        // 2. The coroutine's lifetime is bounded by the caller
        // 3. The frame size is not prohibitively large
        self.frame_does_not_escape = !frame_is_escaped && !is_cross_thread;
        self.requires_heap = is_cross_thread || frame_is_escaped;

        if self.frame_does_not_escape && (is_always_inline || !self.requires_heap) {
            self.can_elide_heap = true;
        }

        self.can_elide_heap
    }

    /// Generate the frame allocation IR considering HALO.
    pub fn gen_halo_alloc_ir(&self, coro_name: &str, frame_size: usize) -> String {
        if self.can_elide_heap {
            format!(
                "  ; HALO: stack-allocated frame for '{}'\n  %frame.{} = alloca i8, i64 {}\n",
                coro_name, coro_name, frame_size
            )
        } else {
            format!(
                "  ; Heap-allocated frame for '{}'\n  %frame.{} = call ptr @malloc(i64 {})\n",
                coro_name, coro_name, frame_size
            )
        }
    }

    /// Escape analysis: check if any path leaks the frame pointer.
    pub fn check_escape(
        uses_move_to_heap: bool,
        is_passed_to_function: bool,
        is_stored_globally: bool,
    ) -> bool {
        uses_move_to_heap || is_passed_to_function || is_stored_globally
    }
}

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

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Handle Types
// ═══════════════════════════════════════════════════════════════════════════════

/// Represents `std::coroutine_handle<Promise>`.
#[derive(Debug, Clone)]
pub struct CoroutineHandle {
    /// The promise type (or void for `coroutine_handle<>`).
    pub promise_type: Option<String>,
    /// The raw frame pointer address.
    pub address: Option<String>,
    /// Whether this is a noop coroutine handle.
    pub is_noop: bool,
}

impl CoroutineHandle {
    /// Create a typed coroutine_handle<Promise>.
    pub fn typed(promise_type: &str) -> Self {
        Self {
            promise_type: Some(promise_type.to_string()),
            address: None,
            is_noop: false,
        }
    }

    /// Create an untyped coroutine_handle<> (void promise).
    pub fn untyped() -> Self {
        Self {
            promise_type: None,
            address: None,
            is_noop: false,
        }
    }

    /// Create a noop coroutine handle.
    pub fn noop() -> Self {
        Self {
            promise_type: None,
            address: None,
            is_noop: true,
        }
    }

    /// The LLVM type of the coroutine handle.
    pub fn llvm_type(&self) -> &str {
        "ptr"
    }

    /// Generate `coroutine_handle::from_promise(promise)`.
    pub fn gen_from_promise_ir(&self, promise_ref: &str) -> String {
        format!(
            "  %handle = call ptr @llvm.coro.promise(ptr {}, i32 0, i1 false)\n",
            promise_ref
        )
    }

    /// Generate `coroutine_handle::from_address(void* addr)`.
    pub fn gen_from_address_ir(addr: &str) -> String {
        format!("  %handle = ptrtoint ptr {} to ptr\n", addr)
    }

    /// Generate `coroutine_handle::address()`.
    pub fn gen_address_ir(&self, handle: &str) -> String {
        format!("  %addr = ptrtoint ptr {} to i64\n", handle)
    }

    /// Generate `coroutine_handle::resume()`.
    pub fn gen_resume_ir(&self, handle: &str) -> String {
        if self.is_noop {
            "; noop coroutine handle — resume is a no-op\n".to_string()
        } else {
            format!(
                "  %resume_fn = call ptr @llvm.coro.subfn.addr(ptr {}, i8 0)\n  call fastcc void %resume_fn(ptr {})\n",
                handle, handle
            )
        }
    }

    /// Generate `coroutine_handle::destroy()`.
    pub fn gen_destroy_ir(&self, handle: &str) -> String {
        if self.is_noop {
            "; noop coroutine handle — destroy is a no-op\n".to_string()
        } else {
            format!("  call void @llvm.coro.destroy(ptr {})\n", handle)
        }
    }

    /// Generate `coroutine_handle::done()` check.
    pub fn gen_done_ir(&self, handle: &str) -> String {
        format!("  %%done = call i1 @llvm.coro.done(ptr {})\n", handle)
    }

    /// Generate `coroutine_handle::promise()` accessor.
    pub fn gen_promise_ir(&self, handle: &str) -> String {
        if let Some(ref pty) = self.promise_type {
            format!(
                "  %promise.ptr = getelementptr inbounds {}, ptr {}, i32 0, i32 3\n",
                pty, handle
            )
        } else {
            "; untyped handle has no promise access\n".to_string()
        }
    }
}

/// A noop coroutine handle — used for symmetric transfer where no
/// actual coroutine switching occurs.
#[derive(Debug, Clone)]
pub struct NoopCoroutineHandle;

impl NoopCoroutineHandle {
    /// The noop coroutine handle address (well-known constant).
    pub const ADDRESS: &str = "@__noop_coro_frame";

    /// Generate the noop coroutine frame declaration.
    pub fn gen_decl_ir() -> &'static str {
        "@__noop_coro_frame = internal constant i8 0\n"
    }

    /// Check if a handle refers to the noop coroutine.
    pub fn is_noop(handle_addr: &str) -> bool {
        handle_addr == "__noop_coro_frame" || handle_addr.contains("noop")
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Lifecycle Management
// ═══════════════════════════════════════════════════════════════════════════════

/// Manages the full lifecycle of a coroutine from creation to destruction.
#[derive(Debug, Clone)]
pub struct CoroutineLifecycle {
    /// The coroutine frame address.
    pub frame_ptr: Option<String>,
    /// The current state.
    pub state: CoroutineState,
    /// Whether the coroutine is heap-allocated.
    pub is_heap_allocated: bool,
    /// Whether the coroutine has completed.
    pub is_done: bool,
    /// Suspend point labels for structured concurrency.
    pub suspend_points: Vec<String>,
}

impl CoroutineLifecycle {
    pub fn new() -> Self {
        Self {
            frame_ptr: None,
            state: CoroutineState::Initial,
            is_heap_allocated: false,
            is_done: false,
            suspend_points: Vec::new(),
        }
    }

    /// Transition to the next state.
    pub fn advance(&mut self, new_state: CoroutineState) {
        self.state = new_state;
        if matches!(new_state, CoroutineState::Done) {
            self.is_done = true;
        }
    }

    /// Record a suspend point label.
    pub fn add_suspend_point(&mut self, label: &str) {
        self.suspend_points.push(label.to_string());
    }

    /// Generate the IR for transitioning between states.
    pub fn gen_state_transition_ir(&self, from: CoroutineState, to: CoroutineState) -> String {
        format!(
            "  ; state transition: {:?} → {:?}\n  store i32 {}, ptr getelementptr inbounds (i32, ptr %frame, i32 0, i32 0)\n",
            from,
            to,
            to.as_int()
        )
    }

    /// Whether the coroutine can be safely destroyed.
    pub fn can_destroy(&self) -> bool {
        matches!(
            self.state,
            CoroutineState::Suspended | CoroutineState::FinalSuspended | CoroutineState::Done
        )
    }

    /// Generate custom allocator IR if the promise type defines operator new.
    pub fn gen_custom_allocator_ir(promise_type: &str, size: usize) -> String {
        format!(
            "  ; custom allocator for promise {}\n  %alloc = call ptr @{}_operator_new(i64 {})\n",
            promise_type, promise_type, size
        )
    }

    /// Generate custom deallocator IR if the promise type defines operator delete.
    pub fn gen_custom_deallocator_ir(promise_type: &str, frame: &str) -> String {
        format!(
            "  ; custom deallocator for promise {}\n  call void @{}_operator_delete(ptr {})\n",
            promise_type, promise_type, frame
        )
    }
}

// ═══════════════════════════════════════════════════════════════════════════════
// Coroutine Parameter Passing
// ═══════════════════════════════════════════════════════════════════════════════

/// Handles passing parameters into a coroutine's frame.
#[derive(Debug, Clone)]
pub struct CoroutineParamPassing {
    /// Parameter values to copy into the frame.
    pub params: Vec<CoroutineParam>,
    /// Whether parameters are passed by value (copied) or reference.
    pub pass_by_value: bool,
}

/// A parameter that must be saved in the coroutine frame.
#[derive(Debug, Clone)]
pub struct CoroutineParam {
    /// Parameter name.
    pub name: String,
    /// Parameter type.
    pub ty: String,
    /// Whether the parameter is used after a suspend point.
    pub spans_suspend: bool,
    /// Whether the parameter needs destruction at frame destroy.
    pub needs_cleanup: bool,
}

impl CoroutineParamPassing {
    pub fn new() -> Self {
        Self {
            params: Vec::new(),
            pass_by_value: true,
        }
    }

    /// Add a parameter to be saved in the frame.
    pub fn add_param(&mut self, name: &str, ty: &str, spans_suspend: bool) {
        self.params.push(CoroutineParam {
            name: name.to_string(),
            ty: ty.to_string(),
            spans_suspend,
            needs_cleanup: false,
        });
    }

    /// Generate IR to copy parameters into the frame.
    pub fn gen_param_copy_ir(&self, frame: &ExtendedCoroutineFrame) -> String {
        let mut ir = String::new();
        ir.push_str("; Copy coroutine parameters into frame\n");
        for param in &self.params {
            if param.spans_suspend {
                ir.push_str(&format!(
                    "  %param.{} = getelementptr inbounds {}, ptr %frame, i32 0, i32 {}\n",
                    param.name,
                    param.ty,
                    frame.saved_params.len() + 4 // after header
                ));
                ir.push_str(&format!(
                    "  store {} %{}, ptr %param.{}\n",
                    param.ty, param.name, param.name
                ));
            }
        }
        ir
    }

    /// Mark a parameter as needing cleanup (has a destructor).
    pub fn mark_needs_cleanup(&mut self, param_name: &str) {
        if let Some(param) = self.params.iter_mut().find(|p| p.name == param_name) {
            param.needs_cleanup = true;
        }
    }
}

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

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

    #[test]
    fn test_coroutine_states() {
        assert_eq!(CoroutineState::Initial.as_int(), 0);
        assert_eq!(CoroutineState::Running.as_int(), 2);
        assert_eq!(CoroutineState::Done.as_int(), 5);
    }

    #[test]
    fn test_promise_type_desc() {
        let promise = PromiseTypeDesc::new("MyPromise", "MyTask");
        assert!(promise.has_get_return_object);
        assert!(promise.has_initial_suspend);
        assert!(promise.has_final_suspend);
        assert!(promise.has_return_void);
        assert!(!promise.has_return_value);
        assert_eq!(promise.return_object_type, "MyTask");
    }

    #[test]
    fn test_coroutine_frame_basic() {
        let frame = CoroutineFrame::new("MyPromise");
        assert!(frame.frame_size >= 32);
        assert_eq!(frame.promise_type, "MyPromise");
    }

    #[test]
    fn test_coroutine_frame_add_local() {
        let mut frame = CoroutineFrame::new("Promise");
        let initial_size = frame.frame_size;
        frame.add_local("x", "i32", 4, true);
        assert!(frame.frame_size > initial_size);
        assert_eq!(frame.locals.len(), 1);
        assert_eq!(frame.locals[0].name, "x");
        assert!(frame.locals[0].spans_suspend);
    }

    #[test]
    fn test_coroutine_lowering_entry() {
        let promise = PromiseTypeDesc::new("GenPromise", "Generator<int>");
        let lowering = CoroutineLowering::new(promise, "my_coro");
        let ir = lowering.gen_coroutine_entry_ir();
        assert!(ir.contains("coro.begin"));
        assert!(ir.contains("coro.promise"));
        assert!(ir.contains("my_coro"));
    }

    #[test]
    fn test_coroutine_lowering_co_await() {
        let promise = PromiseTypeDesc::new("Promise", "Task");
        let mut lowering = CoroutineLowering::new(promise, "await_test");
        let ir = lowering.gen_co_await_ir("%awaiter");
        assert!(ir.contains("co_await"));
        assert!(ir.contains("coro.save"));
        assert!(ir.contains("coro.suspend"));
    }

    #[test]
    fn test_coroutine_lowering_co_yield() {
        let promise = PromiseTypeDesc::new("Promise", "Generator<int>");
        let mut lowering = CoroutineLowering::new(promise, "yield_test");
        let ir = lowering.gen_co_yield_ir("%val");
        assert!(ir.contains("co_yield"));
        assert!(ir.contains("coro.suspend"));
    }

    #[test]
    fn test_coroutine_lowering_co_return() {
        let promise = PromiseTypeDesc::new("Promise", "Task");
        let mut lowering = CoroutineLowering::new(promise, "return_test");
        let ir = lowering.gen_co_return_ir(None);
        assert!(ir.contains("co_return"));
        assert!(ir.contains("coro.final"));
    }

    #[test]
    fn test_coroutine_final_ir() {
        let promise = PromiseTypeDesc::new("Promise", "Task");
        let lowering = CoroutineLowering::new(promise, "final_test");
        let ir = lowering.gen_coroutine_final_ir();
        assert!(ir.contains("coro.cleanup"));
        assert!(ir.contains("coro.destroy"));
    }

    #[test]
    fn test_symmetric_transfer_ir() {
        let transfer = SymmetricTransfer {
            source: "A".into(),
            target: "B".into(),
            is_tail_call: true,
        };
        let ir = transfer.gen_ir();
        assert!(ir.contains("musttail"));
        assert!(ir.contains("A -> B"));
    }

    #[test]
    fn test_intrinsic_decls() {
        let ir = CoroutineLowering::gen_intrinsic_decls_ir();
        assert!(ir.contains("llvm.coro.begin"));
        assert!(ir.contains("llvm.coro.end"));
        assert!(ir.contains("llvm.coro.suspend"));
        assert!(ir.contains("llvm.coro.destroy"));
    }
}