lasm/

lasm_function.rs

use std::collections::HashMap;
use cranelift_jit::JITModule;
use cranelift_module::FuncId;
use super::{LasmError, Value};
use std::ffi::CStr;

#[repr(u8)]
#[derive(Clone, Copy, Debug)]
pub enum Tag {
    Int = 0,
    Float = 1,
    Ptr = 2,
    Null = 3,
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct Slot {
    pub tag: Tag,
    pub _pad: [u8; 7],
    pub data: u64,
}

pub struct LasmFunction {
    jit_module: Option<JITModule>,
    func_id: FuncId,
    variables: Vec<String>,
}

#[unsafe(no_mangle)]
pub extern "C" fn write_int(value: i64) -> i64 {
    print!("{}", value);
    0
}

#[unsafe(no_mangle)]
pub extern "C" fn write_float(value: f64) -> i64 {
    print!("{}", value);
    0
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_malloc(size: i64) -> i64 {
    use libc::malloc;
    if size <= 0 {
        return 0;
    }
    unsafe {
        malloc(size as usize) as i64
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_free(ptr: i64) {
    use libc::free;
    if ptr == 0 {
        return;
    }
    unsafe {
        free(ptr as *mut libc::c_void);
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_sleep(milliseconds: i64) {
    use std::thread;
    use std::time::Duration;
    thread::sleep(Duration::from_millis(milliseconds as u64));
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct TimeResult {
    pub seconds: i64,
    pub nanoseconds: i64,
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_time() -> i64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
    (now.as_secs() as i64 * 1_000_000_000) + now.subsec_nanos() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_rand(seed: i64) -> i64 {
    // Simple LCG: seed = (seed * 1103515245 + 12345) % (1 << 31)
    let mut s = seed as u32;
    s = ((s as u64 * 1103515245u64 + 12345) % (1u64 << 31)) as u32;
    let value = s as f64 / (1u32 << 31) as f64;
    value.to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_atoi(ptr: i64) -> i64 {
    use std::ffi::CStr;
    if ptr == 0 {
        return 0;
    }
    unsafe {
        let c_str = CStr::from_ptr(ptr as *const i8);
        if let Ok(s) = c_str.to_str() {
            s.parse().unwrap_or(0)
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_system(cmd_ptr: i64) -> i64 {
    use std::ffi::CStr;
    use std::process::Command;
    
    #[cfg(unix)]
    unsafe {
        let cmd_c = CStr::from_ptr(cmd_ptr as *const i8);
        if let Ok(cmd_s) = cmd_c.to_str() {
            match Command::new("sh").arg("-c").arg(cmd_s).status() {
                Ok(status) => status.code().unwrap_or(-1) as i64,
                Err(_) => -1,
            }
        } else {
            -1
        }
    }
    #[cfg(windows)]
    unsafe {
        let cmd_c = CStr::from_ptr(cmd_ptr as *const i8);
        if let Ok(cmd_s) = cmd_c.to_str() {
            match Command::new("cmd").arg("/C").arg(cmd_s).status() {
                Ok(status) => status.code().unwrap_or(-1) as i64,
                Err(_) => -1,
            }
        } else {
            -1
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_ends_with(str_ptr: i64, suffix_ptr: i64) -> i64 {
    use std::ffi::CStr;
    unsafe {
        let str_c = CStr::from_ptr(str_ptr as *const i8);
        let suffix_c = CStr::from_ptr(suffix_ptr as *const i8);
        
        if let (Ok(str_s), Ok(suffix_s)) = (str_c.to_str(), suffix_c.to_str()) {
            if str_s.ends_with(suffix_s) {
                1
            } else {
                0
            }
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_trim(str_ptr: i64, out_ptr: i64) -> i64 {
    use std::ffi::CStr;
    unsafe {
        let str_c = CStr::from_ptr(str_ptr as *const i8);
        if let Ok(s) = str_c.to_str() {
            let trimmed = s.trim();
            let bytes = trimmed.as_bytes();
            let len = bytes.len().min(255);
            std::ptr::copy_nonoverlapping(bytes.as_ptr(), out_ptr as *mut u8, len);
            *(out_ptr as *mut u8).add(len) = 0;
            len as i64
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_trim_start(str_ptr: i64, out_ptr: i64) -> i64 {
    use std::ffi::CStr;
    unsafe {
        let str_c = CStr::from_ptr(str_ptr as *const i8);
        if let Ok(s) = str_c.to_str() {
            let trimmed = s.trim_start();
            let bytes = trimmed.as_bytes();
            let len = bytes.len().min(255);
            std::ptr::copy_nonoverlapping(bytes.as_ptr(), out_ptr as *mut u8, len);
            *(out_ptr as *mut u8).add(len) = 0;
            len as i64
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_trim_end(str_ptr: i64, out_ptr: i64) -> i64 {
    use std::ffi::CStr;
    unsafe {
        let str_c = CStr::from_ptr(str_ptr as *const i8);
        if let Ok(s) = str_c.to_str() {
            let trimmed = s.trim_end();
            let bytes = trimmed.as_bytes();
            let len = bytes.len().min(255);
            std::ptr::copy_nonoverlapping(bytes.as_ptr(), out_ptr as *mut u8, len);
            *(out_ptr as *mut u8).add(len) = 0;
            len as i64
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_streq(str1_ptr: i64, str2_ptr: i64) -> i64 {
    use std::ffi::CStr;
    unsafe {
        let str1_c = CStr::from_ptr(str1_ptr as *const i8);
        let str2_c = CStr::from_ptr(str2_ptr as *const i8);
        
        if let (Ok(str1_s), Ok(str2_s)) = (str1_c.to_str(), str2_c.to_str()) {
            if str1_s == str2_s {
                1
            } else {
                0
            }
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_isws(str_ptr: i64) -> i64 {
    use std::ffi::CStr;
    unsafe {
        let str_c = CStr::from_ptr(str_ptr as *const i8);
        if let Ok(s) = str_c.to_str() {
            if s.trim().is_empty() {
                1
            } else {
                0
            }
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
#[allow(unused_variables)]
pub extern "C" fn lasm_syscall(number: i64, arg0: i64, arg1: i64, arg2: i64, arg3: i64, arg4: i64, arg5: i64) -> i64 {
    #[cfg(macos)]
    {
        unsafe { libc::syscall(number as i32, arg0, arg1, arg2, arg3, arg4, arg5) as i64 }
    }
    #[cfg(linux)]
    {
        unsafe { libc::syscall(number, arg0, arg1, arg2, arg3, arg4, arg5) }
    }
    #[cfg(windows)]
    {
        match number {
            1 => { // write
                if arg0 == 1 { // stdout
                    let buf = arg1 as *const u8;
                    let count = arg2 as usize;
                    unsafe {
                        let slice = std::slice::from_raw_parts(buf, count);
                        if let Ok(s) = std::str::from_utf8(slice) {
                            print!("{}", s);
                            std::io::Write::flush(&mut std::io::stdout()).ok();
                            count as i64
                        } else {
                            -1
                        }
                    }
                } else {
                    -1
                }
            }
            0 => { // read
                if arg0 == 0 { // stdin
                    let buf = arg1 as *mut u8;
                    let count = arg2 as usize;
                    unsafe {
                        let mut input = String::new();
                        if std::io::stdin().read_line(&mut input).is_ok() {
                            let bytes = input.as_bytes();
                            let len = bytes.len().min(count);
                            std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf, len);
                            len as i64
                        } else {
                            -1
                        }
                    }
                } else {
                    -1
                }
            }
            _ => -1
        }
    }
    #[cfg(not(any(unix, windows)))]
    {
        -1
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_print_int(value: i64) {
    print!("{}", value);
    use std::io::Write;
    let _ = std::io::stdout().flush();
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_print_float(bits: i64) {
    let value = f64::from_bits(bits as u64);
    print!("{}", value);
    use std::io::Write;
    let _ = std::io::stdout().flush();
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_print_str(ptr: i64) {
    use std::ffi::CStr;
    if ptr == 0 {
        return;
    }
    unsafe {
        let c_str = CStr::from_ptr(ptr as *const i8);
        if let Ok(s) = c_str.to_str() {
            print!("{}", s);
        }
    }
    use std::io::Write;
    let _ = std::io::stdout().flush();
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_itoa(value: i64, buf: i64, size: i64) -> i64 {
    if buf == 0 || size <= 0 {
        return 0;
    }
    let s = format!("{}", value);
    let bytes = s.as_bytes();
    let len = bytes.len().min(size as usize - 1);
    unsafe {
        std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, len);
        *(buf as *mut u8).add(len) = 0; // null terminate
    }
    len as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_fadd(a_bits: i64, b_bits: i64) -> i64 {
    let a = f64::from_bits(a_bits as u64);
    let b = f64::from_bits(b_bits as u64);
    (a + b).to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_fsub(a_bits: i64, b_bits: i64) -> i64 {
    let a = f64::from_bits(a_bits as u64);
    let b = f64::from_bits(b_bits as u64);
    (a - b).to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_fmul(a_bits: i64, b_bits: i64) -> i64 {
    let a = f64::from_bits(a_bits as u64);
    let b = f64::from_bits(b_bits as u64);
    (a * b).to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_fdiv(a_bits: i64, b_bits: i64) -> i64 {
    let a = f64::from_bits(a_bits as u64);
    let b = f64::from_bits(b_bits as u64);
    (a / b).to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_itof(value: i64) -> i64 {
    (value as f64).to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_ftoi(bits: i64) -> i64 {
    let f = f64::from_bits(bits as u64);
    f as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_fmod(a_bits: i64, b_bits: i64) -> i64 {
    let a = f64::from_bits(a_bits as u64);
    let b = f64::from_bits(b_bits as u64);
    (a % b).to_bits() as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_atof(ptr: i64) -> i64 {
    use std::ffi::CStr;
    if ptr == 0 {
        return 0.0f64.to_bits() as i64;
    }
    unsafe {
        let c_str = CStr::from_ptr(ptr as *const i8);
        if let Ok(s) = c_str.to_str() {
            if let Ok(f) = s.parse::<f64>() {
                f.to_bits() as i64
            } else {
                0.0f64.to_bits() as i64
            }
        } else {
            0.0f64.to_bits() as i64
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_ftoa(bits: i64, buf: i64, size: i64) -> i64 {
    if buf == 0 || size <= 0 {
        return 0;
    }
    let f = f64::from_bits(bits as u64);
    let s = format!("{}", f);
    let bytes = s.as_bytes();
    let len = bytes.len().min(size as usize - 1);
    unsafe {
        std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, len);
        *(buf as *mut u8).add(len) = 0; // null terminate
    }
    len as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_memcpy(dst: i64, src: i64, len: i64) {
    if dst == 0 || src == 0 || len <= 0 {
        return;
    }
    unsafe {
        std::ptr::copy_nonoverlapping(src as *const u8, dst as *mut u8, len as usize);
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_memset(dst: i64, val: i64, len: i64) {
    if dst == 0 || len <= 0 {
        return;
    }
    unsafe {
        std::ptr::write_bytes(dst as *mut u8, val as u8, len as usize);
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_fmt_time(total_nanos: i64, fmt_ptr: i64, buffer: i64) -> i64 {
    use std::ffi::CStr;
    use std::time::{SystemTime, UNIX_EPOCH};

    if buffer == 0 || fmt_ptr == 0 {
        return 0;
    }

    let fmt_str = unsafe {
        match CStr::from_ptr(fmt_ptr as *const i8).to_str() {
            Ok(s) => s,
            Err(_) => "%Y-%m-%d %H:%M:%S",
        }
    };
    
    let total_seconds = if total_nanos == 0 {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
        now.as_secs()
    } else {
        (total_nanos / 1_000_000_000) as u64
    };
    
    let _nanoseconds = if total_nanos == 0 {
        let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
        now.subsec_nanos()
    } else {
        (total_nanos % 1_000_000_000) as u32
    };
    
    let days_since_epoch = total_seconds / 86400;
    let mut year = 1970;
    let mut days = days_since_epoch;
    
    loop {
        let is_leap = (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0);
        let days_in_year = if is_leap { 366 } else { 365 };
        if days >= days_in_year {
            days -= days_in_year;
            year += 1;
        } else {
            break;
        }
    }
    
    let month_days = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];
    let is_leap = (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0);
    let mut month = 1;
    let mut day = days + 1;
    
    for (i, &days_in_month) in month_days.iter().enumerate() {
        let actual_days = if i == 1 && is_leap { days_in_month + 1 } else { days_in_month };
        if day > actual_days {
            day -= actual_days;
            month += 1;
        } else {
            break;
        }
    }
    
    let seconds_in_day = total_seconds % 86400;
    let hour = seconds_in_day / 3600;
    let minute = (seconds_in_day % 3600) / 60;
    let second = seconds_in_day % 60;
    
    let mut result = String::new();
    let mut chars = fmt_str.chars().peekable();
    
    while let Some(ch) = chars.next() {
        if ch == '%' {
            if let Some(spec) = chars.next() {
                match spec {
                    'Y' => result.push_str(&format!("{:04}", year)),
                    'm' => result.push_str(&format!("{:02}", month)),
                    'd' => result.push_str(&format!("{:02}", day)),
                    'H' => result.push_str(&format!("{:02}", hour)),
                    'M' => result.push_str(&format!("{:02}", minute)),
                    'S' => result.push_str(&format!("{:02}", second)),
                    '%' => result.push('%'),
                    _ => {
                        result.push('%');
                        result.push(spec);
                    }
                }
            } else {
                result.push('%');
            }
        } else {
            result.push(ch);
        }
    }
    
    let bytes = result.as_bytes();
    let len = bytes.len().min(255);
    
    unsafe {
        std::ptr::copy_nonoverlapping(bytes.as_ptr(), buffer as *mut u8, len);
        *(buffer as *mut u8).add(len) = 0;
    }
    
    len as i64
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_starts_with(str_ptr: i64, prefix_ptr: i64) -> i64 {
    use std::ffi::CStr;
    
    if str_ptr == 0 || prefix_ptr == 0 {
        return 0;
    }
    
    unsafe {
        let str_cstr = CStr::from_ptr(str_ptr as *const i8);
        let prefix_cstr = CStr::from_ptr(prefix_ptr as *const i8);
        
        if let (Ok(str), Ok(prefix)) = (str_cstr.to_str(), prefix_cstr.to_str()) {
            if str.starts_with(prefix) {
                1
            } else {
                0
            }
        } else {
            0
        }
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_str_len(str_ptr: i64) -> i64 {
    use std::ffi::CStr;
    
    if str_ptr == 0 {
        return 0;
    }
    
    unsafe {
        CStr::from_ptr(str_ptr as *const i8).to_bytes().len() as i64
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_strcmp(str1_ptr: i64, str2_ptr: i64) -> i64 {
    if str1_ptr == 0 || str2_ptr == 0 {
        return if str1_ptr == str2_ptr { 0 } else { 1 };
    }
    unsafe {
        let s1 = CStr::from_ptr(str1_ptr as *const i8).to_bytes();
        let s2 = CStr::from_ptr(str2_ptr as *const i8).to_bytes();
        s1.cmp(s2) as i64
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn lasm_strcpy(dst: i64, src: i64) -> i64 {
    if dst == 0 || src == 0 {
        return 0;
    }
    unsafe {
        let src_str = CStr::from_ptr(src as *const i8).to_bytes();
        let dst_slice = std::slice::from_raw_parts_mut(dst as *mut u8, src_str.len() + 1);
        dst_slice[..src_str.len()].copy_from_slice(src_str);
        dst_slice[src_str.len()] = 0;
        dst
    }
}


#[unsafe(no_mangle)]
pub extern "C" fn lasm_flush(fd: i64) {
    if fd == 1 { // stdout
        use std::io::Write;
        let _ = std::io::stdout().flush();
    } else if fd == 2 { // stderr
        use std::io::Write;
        let _ = std::io::stderr().flush();
    }
}

impl LasmFunction {
    pub fn new(jit_module: JITModule, func_id: FuncId, variables: Vec<String>) -> Self {
        LasmFunction {
            jit_module: Some(jit_module),
            func_id,
            variables,
        }
    }

    pub fn call(&self, variables: &HashMap<String, Value>) -> Result<(i32, HashMap<String, Value>), LasmError> {
        if let Some(ref jit) = self.jit_module {
            let code = jit.get_finalized_function(self.func_id);
            let mut vars_vec: Vec<(String, Value)> = self.variables.iter().map(|var| {
            let value = variables.get(var).cloned().unwrap_or(Value::Null);
                (var.clone(), value)
            }).collect();
            vars_vec.sort_by(|a, b| a.0.cmp(&b.0));
            let mut slots = vec![];
            for (_, value) in &vars_vec {
                let slot = match value {
                    Value::Int(i) => Slot { tag: Tag::Int, _pad: [0; 7], data: *i as u64 },
                    Value::Float(f) => Slot { tag: Tag::Float, _pad: [0; 7], data: f.to_bits() },
                    Value::String(s) => {
                        let c_str = std::ffi::CString::new(s.clone()).map_err(|_| LasmError::Runtime("Invalid string".to_string()))?;
                        let ptr = c_str.as_ptr() as *mut u8;
                        std::mem::forget(c_str);
                        Slot { tag: Tag::Ptr, _pad: [0; 7], data: ptr as u64 }
                    }
                    Value::Ptr(p) => Slot { tag: Tag::Ptr, _pad: [0; 7], data: *p as u64 },
                    Value::Null => Slot { tag: Tag::Null, _pad: [0; 7], data: 0 },
                };
                slots.push(slot);
            }

            let mut slot_array = slots.into_boxed_slice();
            let ptr = slot_array.as_mut_ptr();
            let len = slot_array.len();
            std::mem::forget(slot_array);

            let func: extern "C" fn(*mut Slot, usize) -> i32 = unsafe { std::mem::transmute(code) };
            let exit_code = func(ptr, len);

            let slots_slice = unsafe { std::slice::from_raw_parts(ptr, len) };
            let mut final_vars = HashMap::new();
            for i in 0..len {
                let slot = &slots_slice[i];
                let value = match slot.tag {
                    Tag::Int => Value::Int(slot.data as i64),
                    Tag::Float => Value::Float(f64::from_bits(slot.data)),
                    Tag::Ptr => Value::Ptr(slot.data as *mut u8),
                    Tag::Null => Value::Null,
                };
                final_vars.insert(vars_vec[i].0.clone(), value);
            }

            Ok((exit_code, final_vars))
        } else {
            Err(LasmError::Runtime("No JIT module available".to_string()))
        }
    }
}