calltrace-rs 1.1.4

//! x86_64 Register Reading for Function Argument Capture
//!
//! This module implements the x86_64 System V ABI argument passing conventions
//! to extract function arguments from CPU registers and stack memory.

use crate::dwarf_analyzer::{ParameterInfo, TypeInfo};
use crate::error::{CallTraceError, Result};
use serde::{Deserialize, Serialize};
use std::arch::asm;

/// x86_64 register context captured at function entry
#[repr(C)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RegisterContext {
    // Integer argument registers (System V ABI)
    pub rdi: u64, // 1st integer/pointer argument
    pub rsi: u64, // 2nd integer/pointer argument
    pub rdx: u64, // 3rd integer/pointer argument
    pub rcx: u64, // 4th integer/pointer argument
    pub r8: u64,  // 5th integer/pointer argument
    pub r9: u64,  // 6th integer/pointer argument

    // Stack pointers
    pub rsp: u64, // Stack pointer
    pub rbp: u64, // Base pointer (frame pointer)

    // SSE registers for floating point arguments
    pub xmm: [f64; 8], // XMM0-XMM7

    // Additional context
    pub rip: u64,    // Instruction pointer (approximate)
    pub rflags: u64, // Flags register

    // Metadata
    pub valid: bool,
    pub timestamp: u64,

    // Return value registers (captured on function exit)
    pub return_rax: u64,    // Primary integer/pointer return value
    pub return_xmm0: f64,   // Primary floating point return value
    pub return_valid: bool, // Whether return values were captured
}

/// Argument classification according to x86_64 System V ABI
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum ArgumentClass {
    Integer,    // Passed in integer register or stack
    Sse,        // Passed in SSE register
    Memory,     // Passed on stack (large structures)
    X87,        // Passed in x87 register (long double)
    ComplexX87, // Complex numbers in x87
    NoClass,    // Empty or unclassified
}

/// Location where an argument is stored
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArgumentLocation {
    pub class: ArgumentClass,
    pub register_index: Option<usize>, // Which register (if applicable)
    pub stack_offset: Option<usize>,   // Stack offset (if applicable)
    pub size: usize,                   // Size in bytes
}

/// Captured argument value with type information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CapturedArgument {
    pub name: String,
    pub type_name: String,
    pub location: ArgumentLocation,
    pub value: ArgumentValue,
    pub valid: bool,
    pub error: Option<String>,
}

/// Union-like enum for different argument value types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ArgumentValue {
    Integer(u64),
    Float(f32),
    Double(f64),
    Pointer(u64),
    String(String),
    Raw(Vec<u8>),

    // Complex types
    Struct {
        type_name: String,
        members: Vec<StructMember>,
        size: usize,
    },
    Array {
        element_type: String,
        elements: Vec<ArgumentValue>,
        count: usize,
        element_size: usize,
    },
    Union {
        type_name: String,
        raw_data: Vec<u8>,
        size: usize,
    },

    // Special values
    Null,
    Unknown {
        type_name: String,
        raw_data: Vec<u8>,
        error: Option<String>,
    },
}

/// Structure member information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StructMember {
    pub name: String,
    pub type_name: String,
    pub offset: usize,
    pub size: usize,
    pub value: Box<ArgumentValue>,
}

impl RegisterContext {
    /// Capture current register context using inline assembly
    ///
    /// # Safety
    /// This function uses inline assembly and must be called very early
    /// in the function entry hook before registers are modified.
    pub unsafe fn capture() -> Result<Self> {
        let mut context = RegisterContext {
            rdi: 0,
            rsi: 0,
            rdx: 0,
            rcx: 0,
            r8: 0,
            r9: 0,
            rsp: 0,
            rbp: 0,
            rip: 0,
            rflags: 0,
            xmm: [0.0; 8],
            valid: false,
            timestamp: 0,
            return_rax: 0,
            return_xmm0: 0.0,
            return_valid: false,
        };

        #[cfg(target_arch = "x86_64")]
        {
            // Capture integer registers
            asm!(
                "mov {rdi}, rdi",
                "mov {rsi}, rsi",
                "mov {rdx}, rdx",
                "mov {rcx}, rcx",
                "mov {r8}, r8",
                "mov {r9}, r9",
                "mov {rsp}, rsp",
                "mov {rbp}, rbp",
                "pushfq",
                "pop {rflags}",
                rdi = out(reg) context.rdi,
                rsi = out(reg) context.rsi,
                rdx = out(reg) context.rdx,
                rcx = out(reg) context.rcx,
                r8 = out(reg) context.r8,
                r9 = out(reg) context.r9,
                rsp = out(reg) context.rsp,
                rbp = out(reg) context.rbp,
                rflags = out(reg) context.rflags,
                options(nostack, preserves_flags)
            );

            // Capture SSE registers for floating point arguments
            asm!(
                "movsd {xmm0}, xmm0",
                "movsd {xmm1}, xmm1",
                "movsd {xmm2}, xmm2",
                "movsd {xmm3}, xmm3",
                "movsd {xmm4}, xmm4",
                "movsd {xmm5}, xmm5",
                "movsd {xmm6}, xmm6",
                "movsd {xmm7}, xmm7",
                xmm0 = out(xmm_reg) context.xmm[0],
                xmm1 = out(xmm_reg) context.xmm[1],
                xmm2 = out(xmm_reg) context.xmm[2],
                xmm3 = out(xmm_reg) context.xmm[3],
                xmm4 = out(xmm_reg) context.xmm[4],
                xmm5 = out(xmm_reg) context.xmm[5],
                xmm6 = out(xmm_reg) context.xmm[6],
                xmm7 = out(xmm_reg) context.xmm[7],
                options(nostack, preserves_flags)
            );

            // Get approximate instruction pointer
            asm!(
                "lea {rip}, [rip]",
                rip = out(reg) context.rip,
                options(nostack, preserves_flags)
            );

            context.valid = true;
            context.timestamp = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as u64;

            Ok(context)
        }

        #[cfg(not(target_arch = "x86_64"))]
        {
            Err(CallTraceError::NotSupported)
        }
    }

    /// Get integer register value by index (0-5 for RDI, RSI, RDX, RCX, R8, R9)
    #[inline]
    pub fn get_integer_register(&self, index: usize) -> Option<u64> {
        match index {
            0 => Some(self.rdi),
            1 => Some(self.rsi),
            2 => Some(self.rdx),
            3 => Some(self.rcx),
            4 => Some(self.r8),
            5 => Some(self.r9),
            _ => None,
        }
    }

    /// Get SSE register value by index (0-7 for XMM0-XMM7)
    #[inline]
    pub fn get_sse_register(&self, index: usize) -> Option<f64> {
        self.xmm.get(index).copied()
    }
}

/// Classify argument according to x86_64 System V ABI
#[inline]
pub fn classify_argument(type_name: &str, size: usize, arg_index: usize) -> ArgumentLocation {
    const MAX_INTEGER_REGS: usize = 6;
    const MAX_SSE_REGS: usize = 8;

    // Basic classification based on type and size
    if size <= 8 {
        if type_name.contains("float") || type_name.contains("double") {
            // Floating point argument
            if arg_index < MAX_SSE_REGS {
                ArgumentLocation {
                    class: ArgumentClass::Sse,
                    register_index: Some(arg_index),
                    stack_offset: None,
                    size,
                }
            } else {
                ArgumentLocation {
                    class: ArgumentClass::Memory,
                    register_index: None,
                    stack_offset: Some((arg_index - MAX_SSE_REGS) * 8 + 8),
                    size,
                }
            }
        } else {
            // Integer/pointer argument
            if arg_index < MAX_INTEGER_REGS {
                ArgumentLocation {
                    class: ArgumentClass::Integer,
                    register_index: Some(arg_index),
                    stack_offset: None,
                    size,
                }
            } else {
                ArgumentLocation {
                    class: ArgumentClass::Memory,
                    register_index: None,
                    stack_offset: Some((arg_index - MAX_INTEGER_REGS) * 8 + 8),
                    size,
                }
            }
        }
    } else {
        // Large structures go on stack
        ArgumentLocation {
            class: ArgumentClass::Memory,
            register_index: None,
            stack_offset: Some(arg_index * 8 + 8),
            size,
        }
    }
}

/// Extract argument value from register context
#[inline]
pub fn extract_argument(
    context: &RegisterContext,
    location: &ArgumentLocation,
    type_name: &str,
    is_pointer: bool,
) -> Result<ArgumentValue> {
    match location.class {
        ArgumentClass::Integer => {
            if let Some(reg_index) = location.register_index {
                if let Some(reg_value) = context.get_integer_register(reg_index) {
                    if is_pointer {
                        // Try to read string if it's a char*
                        if type_name.contains("char") && is_valid_pointer(reg_value as *const u8) {
                            match read_string_safe(reg_value as *const u8, 256) {
                                Ok(s) => return Ok(ArgumentValue::String(s)),
                                Err(_) => return Ok(ArgumentValue::Pointer(reg_value)),
                            }
                        } else {
                            return Ok(ArgumentValue::Pointer(reg_value));
                        }
                    } else {
                        return Ok(ArgumentValue::Integer(reg_value));
                    }
                }
            }
            Err(CallTraceError::RegisterError(
                "Invalid integer register".to_string(),
            ))
        }

        ArgumentClass::Sse => {
            if let Some(reg_index) = location.register_index {
                if let Some(reg_value) = context.get_sse_register(reg_index) {
                    if type_name.contains("double") {
                        return Ok(ArgumentValue::Double(reg_value));
                    } else if type_name.contains("float") {
                        return Ok(ArgumentValue::Float(reg_value as f32));
                    } else {
                        return Ok(ArgumentValue::Double(reg_value));
                    }
                }
            }
            Err(CallTraceError::RegisterError(
                "Invalid SSE register".to_string(),
            ))
        }

        ArgumentClass::Memory => {
            if let Some(stack_offset) = location.stack_offset {
                // Read from stack - this is quite dangerous and may not always work
                unsafe {
                    let stack_ptr = context.rsp as *const u8;
                    let value_ptr = stack_ptr.add(stack_offset) as *const u64;
                    if is_valid_pointer(value_ptr as *const u8) {
                        let value = *value_ptr;
                        if is_pointer {
                            Ok(ArgumentValue::Pointer(value))
                        } else {
                            Ok(ArgumentValue::Integer(value))
                        }
                    } else {
                        Err(CallTraceError::RegisterError(
                            "Invalid stack memory".to_string(),
                        ))
                    }
                }
            } else {
                Err(CallTraceError::RegisterError("No stack offset".to_string()))
            }
        }

        _ => Err(CallTraceError::NotSupported),
    }
}

/// Check if a pointer is likely valid for reading
#[inline]
fn is_valid_pointer(ptr: *const u8) -> bool {
    if ptr.is_null() {
        return false;
    }

    let addr = ptr as usize;

    // Basic sanity checks
    if addr < 0x1000 {
        return false; // Null page
    }

    if addr >= 0xffff_8000_0000_0000 {
        return false; // Kernel space
    }

    true
}

/// Safely read a null-terminated string from memory
fn read_string_safe(ptr: *const u8, max_len: usize) -> Result<String> {
    if !is_valid_pointer(ptr) {
        return Err(CallTraceError::RegisterError(
            "Invalid string pointer".to_string(),
        ));
    }

    unsafe {
        let mut bytes = Vec::new();
        let mut current = ptr;

        for _ in 0..max_len {
            let byte = *current;
            if byte == 0 {
                break;
            }
            bytes.push(byte);
            current = current.add(1);
        }

        String::from_utf8(bytes)
            .map_err(|_| CallTraceError::RegisterError("Invalid UTF-8 string".to_string()))
    }
}

/// Enhanced argument extraction with DWARF type information
#[inline]
pub fn extract_argument_with_type_info(
    context: &RegisterContext,
    location: &ArgumentLocation,
    type_info: &TypeInfo,
    param_info: &ParameterInfo,
) -> Result<ArgumentValue> {
    // Fast path: handle null pointers for integer registers
    if location.class == ArgumentClass::Integer && type_info.is_pointer {
        if let Some(reg_index) = location.register_index {
            if let Some(reg_value) = context.get_integer_register(reg_index) {
                if reg_value == 0 {
                    return Ok(ArgumentValue::Null);
                }
            }
        }
    }

    // Optimized dispatch based on type information using match for better branch prediction
    match (
        type_info.is_array,
        type_info.is_struct,
        type_info.is_pointer,
    ) {
        (true, _, _) => extract_array_argument(context, location, type_info, param_info),
        (false, true, _) => extract_struct_argument(context, location, type_info, param_info),
        (false, false, true) => extract_pointer_argument(context, location, type_info, param_info),
        _ => extract_basic_argument(context, location, type_info, param_info),
    }
}

/// Extract array argument with element analysis
fn extract_array_argument(
    context: &RegisterContext,
    location: &ArgumentLocation,
    type_info: &TypeInfo,
    _param_info: &ParameterInfo,
) -> Result<ArgumentValue> {
    let array_size = type_info.array_size.unwrap_or(0) as usize;
    let element_size = if let Some(ref base_type) = type_info.base_type {
        base_type.size.unwrap_or(1) as usize
    } else {
        1
    };

    // For arrays, we need to read from memory (usually stack or pointer)
    match location.class {
        ArgumentClass::Integer | ArgumentClass::Memory => {
            let data_ptr = if location.class == ArgumentClass::Integer {
                // Array passed as pointer
                if let Some(reg_index) = location.register_index {
                    if let Some(ptr_value) = context.get_integer_register(reg_index) {
                        ptr_value as *const u8
                    } else {
                        return Ok(ArgumentValue::Unknown {
                            type_name: type_info.name.clone(),
                            raw_data: vec![],
                            error: Some("Failed to read array pointer register".to_string()),
                        });
                    }
                } else {
                    return Ok(ArgumentValue::Unknown {
                        type_name: type_info.name.clone(),
                        raw_data: vec![],
                        error: Some("No register for array pointer".to_string()),
                    });
                }
            } else {
                // Array on stack
                if let Some(stack_offset) = location.stack_offset {
                    unsafe { (context.rsp as *const u8).add(stack_offset) }
                } else {
                    return Ok(ArgumentValue::Unknown {
                        type_name: type_info.name.clone(),
                        raw_data: vec![],
                        error: Some("No stack offset for array".to_string()),
                    });
                }
            };

            if !is_valid_pointer(data_ptr) {
                return Ok(ArgumentValue::Unknown {
                    type_name: type_info.name.clone(),
                    raw_data: vec![],
                    error: Some("Invalid array pointer".to_string()),
                });
            }

            // Extract array elements (limit to reasonable size)
            let max_elements = std::cmp::min(array_size, 64);
            let mut elements = Vec::new();

            for i in 0..max_elements {
                unsafe {
                    let element_ptr = data_ptr.add(i * element_size);
                    if !is_valid_pointer(element_ptr) {
                        break;
                    }

                    let element_value = extract_primitive_from_memory(
                        element_ptr,
                        element_size,
                        type_info
                            .base_type
                            .as_ref()
                            .map(|t| t.name.as_str())
                            .unwrap_or("unknown"),
                    )?;
                    elements.push(element_value);
                }
            }

            Ok(ArgumentValue::Array {
                element_type: type_info
                    .base_type
                    .as_ref()
                    .map(|t| t.name.clone())
                    .unwrap_or_else(|| "unknown".to_string()),
                elements,
                count: array_size,
                element_size,
            })
        }
        _ => Ok(ArgumentValue::Unknown {
            type_name: type_info.name.clone(),
            raw_data: vec![],
            error: Some("Unsupported array argument class".to_string()),
        }),
    }
}

/// Extract struct argument with member analysis
fn extract_struct_argument(
    context: &RegisterContext,
    location: &ArgumentLocation,
    type_info: &TypeInfo,
    _param_info: &ParameterInfo,
) -> Result<ArgumentValue> {
    let struct_size = type_info.size.unwrap_or(0) as usize;

    // Read raw struct data
    let _raw_data = match location.class {
        ArgumentClass::Integer => {
            // Small struct passed by value in register
            if let Some(reg_index) = location.register_index {
                if let Some(reg_value) = context.get_integer_register(reg_index) {
                    reg_value.to_le_bytes().to_vec()
                } else {
                    return Ok(ArgumentValue::Unknown {
                        type_name: type_info.name.clone(),
                        raw_data: vec![],
                        error: Some("Failed to read struct register".to_string()),
                    });
                }
            } else {
                return Ok(ArgumentValue::Unknown {
                    type_name: type_info.name.clone(),
                    raw_data: vec![],
                    error: Some("No register for struct".to_string()),
                });
            }
        }
        ArgumentClass::Memory => {
            // Struct on stack
            if let Some(stack_offset) = location.stack_offset {
                unsafe {
                    let struct_ptr = (context.rsp as *const u8).add(stack_offset);
                    if is_valid_pointer(struct_ptr) && struct_size <= 256 {
                        std::slice::from_raw_parts(struct_ptr, struct_size).to_vec()
                    } else {
                        return Ok(ArgumentValue::Unknown {
                            type_name: type_info.name.clone(),
                            raw_data: vec![],
                            error: Some("Invalid struct memory".to_string()),
                        });
                    }
                }
            } else {
                return Ok(ArgumentValue::Unknown {
                    type_name: type_info.name.clone(),
                    raw_data: vec![],
                    error: Some("No stack offset for struct".to_string()),
                });
            }
        }
        _ => {
            return Ok(ArgumentValue::Unknown {
                type_name: type_info.name.clone(),
                raw_data: vec![],
                error: Some("Unsupported struct argument class".to_string()),
            });
        }
    };

    // For now, return basic struct info
    // TODO: In future, parse DWARF member information
    Ok(ArgumentValue::Struct {
        type_name: type_info.name.clone(),
        members: vec![], // TODO: Extract member info from DWARF
        size: struct_size,
    })
}

/// Extract pointer argument with enhanced dereferencing
fn extract_pointer_argument(
    context: &RegisterContext,
    location: &ArgumentLocation,
    type_info: &TypeInfo,
    _param_info: &ParameterInfo,
) -> Result<ArgumentValue> {
    // Get pointer value
    let pointer_value = match location.class {
        ArgumentClass::Integer => {
            if let Some(reg_index) = location.register_index {
                context.get_integer_register(reg_index).unwrap_or(0)
            } else {
                return Ok(ArgumentValue::Unknown {
                    type_name: type_info.name.clone(),
                    raw_data: vec![],
                    error: Some("No register for pointer".to_string()),
                });
            }
        }
        ArgumentClass::Memory => {
            if let Some(stack_offset) = location.stack_offset {
                unsafe {
                    let ptr_ptr = (context.rsp as *const u8).add(stack_offset) as *const u64;
                    if is_valid_pointer(ptr_ptr as *const u8) {
                        *ptr_ptr
                    } else {
                        return Ok(ArgumentValue::Unknown {
                            type_name: type_info.name.clone(),
                            raw_data: vec![],
                            error: Some("Invalid pointer memory".to_string()),
                        });
                    }
                }
            } else {
                return Ok(ArgumentValue::Unknown {
                    type_name: type_info.name.clone(),
                    raw_data: vec![],
                    error: Some("No stack offset for pointer".to_string()),
                });
            }
        }
        _ => {
            return Ok(ArgumentValue::Unknown {
                type_name: type_info.name.clone(),
                raw_data: vec![],
                error: Some("Unsupported pointer argument class".to_string()),
            });
        }
    };

    if pointer_value == 0 {
        return Ok(ArgumentValue::Null);
    }

    // Try to dereference based on pointed-to type
    if let Some(ref base_type) = type_info.base_type {
        if base_type.name.contains("char") {
            // String pointer
            // Fall through to pointer value on error
            if let Ok(s) = read_string_safe(pointer_value as *const u8, 256) {
                return Ok(ArgumentValue::String(s));
            }
        }
        // TODO: Handle other pointed-to types (struct*, int*, etc.)
    }

    Ok(ArgumentValue::Pointer(pointer_value))
}

/// Extract basic argument (non-complex types)
#[inline]
fn extract_basic_argument(
    context: &RegisterContext,
    location: &ArgumentLocation,
    type_info: &TypeInfo,
    _param_info: &ParameterInfo,
) -> Result<ArgumentValue> {
    // Use existing extract_argument function
    extract_argument(context, location, &type_info.name, type_info.is_pointer)
}

/// Extract primitive value from memory
#[inline]
unsafe fn extract_primitive_from_memory(
    ptr: *const u8,
    size: usize,
    type_name: &str,
) -> Result<ArgumentValue> {
    if !is_valid_pointer(ptr) {
        return Err(CallTraceError::RegisterError(
            "Invalid memory pointer".to_string(),
        ));
    }

    // Optimized match with frequently used sizes first
    match size {
        8 => {
            if type_name.contains("double") {
                let value = *(ptr as *const f64);
                Ok(ArgumentValue::Double(value))
            } else {
                let value = *(ptr as *const u64);
                Ok(ArgumentValue::Integer(value))
            }
        }
        4 => {
            if type_name.contains("float") {
                let value = *(ptr as *const f32);
                Ok(ArgumentValue::Float(value))
            } else {
                let value = *(ptr as *const u32) as u64;
                Ok(ArgumentValue::Integer(value))
            }
        }
        1 => {
            let value = *ptr as u64;
            Ok(ArgumentValue::Integer(value))
        }
        2 => {
            let value = *(ptr as *const u16) as u64;
            Ok(ArgumentValue::Integer(value))
        }
        _ => {
            // Unknown size, read as raw bytes (limit to 64 bytes for safety)
            let read_size = std::cmp::min(size, 64);
            let data = std::slice::from_raw_parts(ptr, read_size).to_vec();
            Ok(ArgumentValue::Raw(data))
        }
    }
}

/// Capture return values from function exit context
///
/// # Safety
/// This function uses inline assembly and must be called at function exit
/// before return registers are modified.
pub unsafe fn capture_return_values() -> Result<RegisterContext> {
    let mut context = RegisterContext {
        rdi: 0,
        rsi: 0,
        rdx: 0,
        rcx: 0,
        r8: 0,
        r9: 0,
        rsp: 0,
        rbp: 0,
        rip: 0,
        rflags: 0,
        xmm: [0.0; 8],
        valid: false,
        timestamp: 0,
        return_rax: 0,
        return_xmm0: 0.0,
        return_valid: false,
    };

    #[cfg(target_arch = "x86_64")]
    {
        // Capture return value registers
        asm!(
            "mov {rax}, rax",        // Capture RAX (integer/pointer return value)
            "movsd {xmm0}, xmm0",    // Capture XMM0 (floating point return value)
            "mov {rsp}, rsp",        // Also capture stack pointer for context
            rax = out(reg) context.return_rax,
            xmm0 = out(xmm_reg) context.return_xmm0,
            rsp = out(reg) context.rsp,
            options(nostack, preserves_flags)
        );

        context.return_valid = true;
        context.valid = true;
        context.timestamp = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_micros() as u64;

        Ok(context)
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        Err(CallTraceError::NotSupported)
    }
}

/// Extract return value based on function type information
#[inline]
pub fn extract_return_value(
    context: &RegisterContext,
    return_type: Option<&TypeInfo>,
) -> Option<ArgumentValue> {
    if !context.return_valid {
        return None;
    }

    // If we don't have type information, return raw RAX value
    let type_info = return_type?;

    // Handle different return types according to x86_64 ABI
    if type_info.name.contains("void") {
        None // void functions don't return values
    } else if type_info.name.contains("float") && type_info.size == Some(4) {
        // float return value in XMM0
        Some(ArgumentValue::Float(context.return_xmm0 as f32))
    } else if type_info.name.contains("double")
        || (type_info.name.contains("float") && type_info.size == Some(8))
    {
        // double return value in XMM0
        Some(ArgumentValue::Double(context.return_xmm0))
    } else if type_info.is_pointer {
        // Pointer return value in RAX
        if context.return_rax == 0 {
            Some(ArgumentValue::Null)
        } else {
            // Try to dereference if it's a string pointer
            if let Some(ref base_type) = type_info.base_type {
                if base_type.name.contains("char") {
                    // Fall through to pointer value on error
                    if let Ok(s) = read_string_safe(context.return_rax as *const u8, 256) {
                        return Some(ArgumentValue::String(s));
                    }
                }
            }
            Some(ArgumentValue::Pointer(context.return_rax))
        }
    } else if type_info.size.is_some_and(|s| s <= 8) {
        // Integer return value in RAX (up to 64-bit)
        Some(ArgumentValue::Integer(context.return_rax))
    } else {
        // Large structures or unknown types - return raw RAX value
        Some(ArgumentValue::Unknown {
            type_name: type_info.name.clone(),
            raw_data: context.return_rax.to_le_bytes().to_vec(),
            error: Some("Complex return type not fully supported".to_string()),
        })
    }
}