decuda 0.1.1

//! CUDA parser: preprocessed source -> tree-sitter-cpp for validation +
//! lightweight regex-based scan over the linear source to collect IR nodes
//! about CUDA-specific constructs (qualifiers, runtime calls, built-in
//! variables, header includes).

use std::path::Path;
use std::sync::OnceLock;

use anyhow::Result;
use regex::Regex;
use tree_sitter::{Language, Parser};

use crate::cuda_db;
use crate::ir::{
    slot_index, BuiltinKind, CudaQualifier, IrNode, TranslationUnit,
};
use crate::preprocess::{preprocess, LaunchSite};

/// parse + walk -> TranslationUnit
pub fn translate_path(path: &Path) -> Result<TranslationUnit> {
    let raw = std::fs::read_to_string(path)
        .map_err(|e| anyhow::anyhow!("read {}: {e}", path.display()))?;
    Ok(translate_source(raw, path.to_path_buf()))
}

/// In-memory variant: useful for tests and small fixtures.
pub fn translate_source(source: String, path: std::path::PathBuf) -> TranslationUnit {
    let mut unit = TranslationUnit::new(source.clone(), path);
    let pre = preprocess(&source);

    // Validate that the preprocessed source parses. We do not use the AST
    // directly because tree-sitter-cpp does not understand CUDA-specific
    // identifiers (those still appear as identifiers in the AST); we rely
    // on a per-target regex sweep below to harvest the IR.
    //
    // IMPORTANT: all IR harvesting happens against the ORIGINAL source, not
    // the preprocessed one. The preprocessor is only used to capture
    // kernel launches (their positions, kernel name, grid, block, args).
    // Re-running the regex sweeps against the rewritten text would produce
    // positions that no longer line up with `unit.source`.
    if let Some(mut parser) = make_parser() {
        let _ = parser.parse(&pre.text, None);
    }

    let skip = skip_spans(&source);
    harvest_launches(&pre, &skip, &mut unit);
    harvest_qualifiers(&source, &skip, &mut unit);
    harvest_builtins(&source, &skip, &mut unit);
    harvest_runtime_calls(&source, &skip, &mut unit);
    harvest_atomics(&source, &skip, &mut unit);
    harvest_includes(&source, &skip, &mut unit);
    harvest_kernel_defs(&source, &skip, &mut unit);
    harvest_ptx_inline(&source, &skip, &mut unit);

    unit.nodes.sort_by_key(|n| n.start());
    unit
}

fn make_parser() -> Option<Parser> {
    let mut parser = Parser::new();
    let lang: Language = tree_sitter_cpp::LANGUAGE.into();
    parser.set_language(&lang).ok()?;
    Some(parser)
}

/// Compute (start, end) byte spans that should be ignored by the regex sweep.
fn skip_spans(source: &str) -> Vec<(usize, usize)> {
    let mut out = Vec::new();
    let bytes = source.as_bytes();
    let mut i = 0usize;
    while i < bytes.len() {
        let b = bytes[i];
        // line comment
        if b == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'/' {
            let start = i;
            while i < bytes.len() && bytes[i] != b'\n' {
                i += 1;
            }
            out.push((start, i));
            continue;
        }
        // block comment
        if b == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
            let start = i;
            i += 2;
            while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
                i += 1;
            }
            i = (i + 2).min(bytes.len());
            out.push((start, i));
            continue;
        }
        // raw string R"delim(...)delim"
        if b == b'R' && i + 6 < bytes.len() && bytes[i + 1] == b'"' {
            // Find the opening delim
            let mut j = i + 2;
            let delim_start = j;
            while j < bytes.len() && bytes[j] != b'(' {
                j += 1;
            }
            if j < bytes.len() {
                let delim = &bytes[delim_start..j];
                // Find closing )delim"
                let mut k = j + 1;
                let needle: Vec<u8> = {
                    let mut v = Vec::with_capacity(delim.len() + 2);
                    v.push(b')');
                    v.extend_from_slice(delim);
                    v.push(b'"');
                    v
                };
                while k + needle.len() <= bytes.len() {
                    if &bytes[k..k + needle.len()] == needle.as_slice() {
                        k += needle.len();
                        out.push((i, k));
                        i = k;
                        break;
                    }
                    k += 1;
                }
                if k + needle.len() > bytes.len() {
                    out.push((i, bytes.len()));
                    i = bytes.len();
                }
                continue;
            }
        }
        // string literal
        if b == b'"' {
            let start = i;
            i += 1;
            while i < bytes.len() {
                if bytes[i] == b'\\' {
                    i += 2;
                    continue;
                }
                if bytes[i] == b'"' {
                    i += 1;
                    break;
                }
                i += 1;
            }
            out.push((start, i));
            continue;
        }
        // char literal
        if b == b'\'' {
            let start = i;
            i += 1;
            while i < bytes.len() {
                if bytes[i] == b'\\' {
                    i += 2;
                    continue;
                }
                if bytes[i] == b'\'' {
                    i += 1;
                    break;
                }
                i += 1;
            }
            out.push((start, i));
            continue;
        }
        i += 1;
    }
    out
}

fn in_skip(skip: &[(usize, usize)], pos: usize) -> bool {
    skip.iter().any(|(s, e)| pos >= *s && pos < *e)
}

fn harvest_launches(
    pre: &crate::preprocess::PreprocessedSource,
    skip: &[(usize, usize)],
    unit: &mut TranslationUnit,
) {
    for LaunchSite {
        start,
        end,
        line: _,
        kernel,
        grid,
        block,
        smem,
        stream,
        args,
    } in &pre.launches
    {
        // Skip launches that fall inside comments or string literals.
        if in_skip(skip, *start) || in_skip(skip, end.saturating_sub(1)) {
            continue;
        }
        unit.push(IrNode::KernelLaunch {
            start: *start,
            end: *end,
            kernel: kernel.clone(),
            grid: grid.clone(),
            block: block.clone(),
            smem: smem.clone(),
            stream: stream.clone(),
            args: args.clone(),
        });
    }
    // Marker offsets are reserved for future use by rewriters that need to
    // know where the synthetic `__decuda_launch(...)` token sits after rewrite.
    let _ = pre.marker_offsets.len();
}

fn harvest_qualifiers(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    // Search for every CUDA function/storage qualifier. We use regex for
    // simplicity: each match is recorded as a qualifier IR node.
    let re = qualifier_re();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        let tok = caps.get(1).unwrap().as_str();
        if let Some(q) = CudaQualifier::from_token(tok).or_else(|| CudaQualifier::from_storage_class(tok))
        {
            unit.push(IrNode::QualifierDecl {
                start: m.start(),
                end: m.end(),
                qualifier: q,
                surface: tok.to_string(),
            });
        }
    }
}

fn qualifier_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(r"\b(__global__|__device__|__host__|__forceinline__|__noinline__|__shared__|__constant__|__managed__|__restrict__)\b")
            .expect("qualifier regex")
    })
}

fn harvest_builtins(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    let re = builtin_re();
    let bytes = text.as_bytes();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        let tok = caps.get(1).unwrap().as_str();
        if let Some(kind) = BuiltinKind::from_token(tok) {
            let mut end = m.end();
            let mut p = m.end();
            while p < bytes.len() && bytes[p].is_ascii_whitespace() {
                p += 1;
            }
            let mut has_field_access = false;
            if p < bytes.len() && bytes[p] == b'(' && kind.consumes_args() {
                // Sync intrinsics like `__syncthreads()` carry their own
                // argument list; include it so the replacement substitutes
                // the whole call expression (avoids `__syncthreads()();`).
                // Warp primitives (`__shfl_sync(...)`, etc.) do NOT consume
                // args — only the function name is rewritten, preserving the
                // argument list as-is.
                if let Some((close, _)) = crate::preprocess::balanced_parens(text, p) {
                    end = close + 1;
                }
            } else if p + 1 < bytes.len()
                && bytes[p] == b'.'
                && matches!(bytes[p + 1], b'x' | b'y' | b'z')
                && (p + 2 == bytes.len() || !is_id_char(bytes[p + 2]))
            {
                // `threadIdx.x` / `blockIdx.y` / etc. Note the field access
                // so the rewriter can decide whether to consume the `.x`
                // (OpenCL/SYCL/Rust drop it; HIP keeps it).
                has_field_access = true;
            }
            unit.push(IrNode::BuiltinRef {
                start: m.start(),
                end,
                kind,
                has_field_access,
            });
        }
    }
}

fn is_id_char(b: u8) -> bool {
    b.is_ascii_alphanumeric() || b == b'_'
}

fn builtin_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(r"\b(threadIdx|blockIdx|blockDim|gridDim|warpSize|__syncthreads|__syncwarp|__sync_block|__laneid|__warp_id|__shfl_sync|__ballot_sync|__any_sync|__all_sync|__activemask)\b")
            .expect("builtin regex")
    })
}

fn harvest_runtime_calls(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    let re = runtime_re();
    let bytes = text.as_bytes();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        let name = caps.name("name").unwrap().as_str().to_string();

        // After the name, possibly whitespace, then expect `(` for the args.
        // We do not include the args inside the matched span; we extend it
        // forward through a balanced-paren walk to capture the full call.
        let mut p = m.end();
        while p < bytes.len() && bytes[p].is_ascii_whitespace() {
            p += 1;
        }
        let end = if p < bytes.len() && bytes[p] == b'(' {
            match crate::preprocess::balanced_parens(text, p) {
                Some((close, _)) => close + 1,
                None => m.end(),
            }
        } else {
            m.end()
        };

        // For argument extraction we want the substring between the matched
        // name and the call's closing paren.
        let args_text = if p < bytes.len() && bytes[p] == b'(' {
            crate::preprocess::balanced_parens(text, p)
                .map(|(_, inner)| inner)
                .unwrap_or("")
        } else {
            ""
        };
        let args = crate::preprocess::split_call_args_pub(args_text)
            .unwrap_or_else(|| vec![args_text.to_string()]);

        let mappings = cuda_mappings_for(&name);
        unit.push(IrNode::RuntimeCall {
            start: m.start(),
            end,
            name,
            args,
            mappings,
        });
    }
}

fn runtime_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(r"\b(?P<name>cuda[A-Za-z0-9_]*)").expect("runtime regex")
    })
}

fn cuda_mappings_for(name: &str) -> [Option<String>; 4] {
    let mut out = [None, None, None, None];
    if let Some(info) = cuda_db::lookup(name) {
        out[slot_index(crate::cli::Target::Hip)] =
            info.hip.map(str::to_string);
        out[slot_index(crate::cli::Target::Sycl)] =
            info.sycl.map(str::to_string);
        out[slot_index(crate::cli::Target::Rust)] =
            info.rust.map(str::to_string);
        out[slot_index(crate::cli::Target::Opencl)] =
            info.opencl.map(str::to_string);
    }
    out
}

fn harvest_atomics(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    let re = atomic_re();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        let name = caps.get(1).unwrap().as_str();
        unit.push(IrNode::AtomicIntrinsic {
            start: m.start(),
            end: m.end(),
            name: name.to_string(),
        });
    }
}

fn atomic_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(r"\b(atomic(Add|Sub|CAS|Exch|Min|Max|Or|And|Xor|Inc|Dec|CompareAndSwap))\b")
            .expect("atomic regex")
    })
}

fn harvest_includes(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    let re = include_re();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        let header = caps.get(1).unwrap().as_str().to_string();
        let replacements = include_replacements(&header);
        unit.push(IrNode::HeaderInclude {
            start: m.start(),
            end: m.end(),
            header,
            replacements,
        });
    }
}

fn include_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(r#"(?m)^\s*#\s*include\s*[<"]([^>"]+)[>]"#).expect("include regex")
    })
}

fn include_replacements(header: &str) -> [Option<String>; 4] {
    let mut out = [None, None, None, None];
    let h = header.trim();
    if h == "cuda_runtime.h" {
        out[slot_index(crate::cli::Target::Hip)] = Some("<hip/hip_runtime.h>".into());
        out[slot_index(crate::cli::Target::Sycl)] = Some("<sycl/sycl.hpp>".into());
        out[slot_index(crate::cli::Target::Rust)] =
            Some("cust::cuda_build_setup() /* TODO: import cust crate */".into());
        out[slot_index(crate::cli::Target::Opencl)] = Some("<CL/cl.h>".into());
    } else if h == "cuda.h" {
        out[slot_index(crate::cli::Target::Hip)] = Some("<hip/hip_runtime.h>".into());
        out[slot_index(crate::cli::Target::Sycl)] = Some("<sycl/sycl.hpp>".into());
        out[slot_index(crate::cli::Target::Rust)] = Some("cust::cuda_build_setup()".into());
        out[slot_index(crate::cli::Target::Opencl)] = Some("<CL/cl.h>".into());
    } else if h == "device_functions.h" {
        out[slot_index(crate::cli::Target::Hip)] = Some("<hip/device_functions.h>".into());
        out[slot_index(crate::cli::Target::Sycl)] = Some("<sycl/sycl.hpp>".into());
        out[slot_index(crate::cli::Target::Rust)] = None;
        out[slot_index(crate::cli::Target::Opencl)] = None;
    } else if h.contains("cuda_") {
        // Generic cuda_*.h header; mark with a placeholder.
        for s in &mut out {
            *s = Some(format!("/* TODO: replace with target equivalent of {h} */"));
        }
    }
    out
}

fn harvest_kernel_defs(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    // Look for `qualifier(s) type name ( params )` followed by `{` to detect
    // kernel definitions. This is intentionally simple: we only flag them so
    // the emit step can wrap them in target boilerplate.
    let re = kernel_def_re();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        let qual = caps.name("q").map(|m| m.as_str()).unwrap_or("");
        if qual.contains("__global__") {
            let name = caps.name("n").map(|m| m.as_str()).unwrap_or("").to_string();
            let params = caps
                .name("p")
                .map(|m| m.as_str())
                .unwrap_or("")
                .trim()
                .to_string();
            unit.push(IrNode::KernelDef {
                start: m.start(),
                end: m.end(),
                name,
                params: vec![params],
            });
        }
    }
}

fn kernel_def_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(
            r"(?xs)
            (?P<q>(?:__global__|__device__|__host__|__forceinline__|__noinline__)\s+)+
            [A-Za-z_][A-Za-z0-9_:,*\s&]+
            \s+
            (?P<n>[A-Za-z_][A-Za-z0-9_]*)
            \s*
            (?P<p>\([^()]*\))
            \s*
            \{
        ",
        )
        .expect("kernel def regex")
    })
}

/// Detect inline PTX assembly (`asm("...")` or `asm volatile("...")`) and
/// emit a `Warning` node so the migration report flags it for manual review.
/// PTX is NVIDIA-specific and has no equivalent in HIP/SYCL/OpenCL/Rust.
fn harvest_ptx_inline(text: &str, skip: &[(usize, usize)], unit: &mut TranslationUnit) {
    let re = ptx_re();
    for caps in re.captures_iter(text) {
        let m = caps.get(0).unwrap();
        if in_skip(skip, m.start()) {
            continue;
        }
        // Extend the span to include the full asm statement (semicolon or
        // closing brace, whichever comes first on the same line).
        let bytes = text.as_bytes();
        let mut end = m.end();
        // Consume the asm string literal and any trailing args.
        if end < bytes.len() && bytes[end] == b'(' {
            if let Some((close, _)) = crate::preprocess::balanced_parens(text, end) {
                end = close + 1;
            }
        }
        let surface = text[m.start()..end.min(text.len())].to_string();
        unit.push(IrNode::Warning {
            start: m.start(),
            end,
            surface,
            message: "inline PTX assembly (`asm(...)`) is NVIDIA-specific; \
                      manual rewrite required for all targets"
                .to_string(),
        });
    }
}

fn ptx_re() -> &'static Regex {
    static R: OnceLock<Regex> = OnceLock::new();
    R.get_or_init(|| {
        Regex::new(r"\basm\s*(?:volatile\s*)?\(")
            .expect("ptx inline asm regex")
    })
}

/// Public helper for tests: exposed so other modules don't have to fight with
/// the private `make_parser`.
#[allow(dead_code)]
pub fn validate(source: &str) -> bool {
    let mut parser = match make_parser() {
        Some(p) => p,
        None => return false,
    };
    parser.parse(source, None).is_some()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ir::{BuiltinKind, CudaQualifier, IrNode};

    fn translate_str(src: &str) -> TranslationUnit {
        translate_source(src.to_string(), "input.cu".into())
    }

    #[test]
    fn finds_global_qualifier() {
        let u = translate_str("__global__ void k() {}");
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::QualifierDecl { qualifier: CudaQualifier::Global, .. }
        )));
    }

    #[test]
    fn finds_syncthreads() {
        let u = translate_str("__syncthreads();");
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::BuiltinRef { kind: BuiltinKind::SyncThreads, .. }
        )));
    }

    #[test]
    fn finds_cuda_runtime_call() {
        let u = translate_str("cudaMalloc((void**)&p, n);");
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::RuntimeCall { name, .. } if name == "cudaMalloc"
        )));
    }

    #[test]
    fn skips_comments_and_strings() {
        let u = translate_str(r#"// __global__ void fake() {}"#);
        assert!(u.nodes.is_empty());
        let u = translate_str(r#"const char* msg = "__global__";"#);
        assert!(u.nodes.is_empty());
    }

    #[test]
    fn finds_kernel_definition() {
        let u = translate_str("__global__ void saxpy(float a, float* x, float* y) { y[0] = a * x[0]; }");
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::KernelDef { name, .. } if name == "saxpy"
        )));
    }

    #[test]
    fn finds_atomic_call() {
        let u = translate_str("atomicAdd(&acc, 1);");
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::AtomicIntrinsic { name, .. } if name == "atomicAdd"
        )));
    }

    #[test]
    fn finds_include() {
        let u = translate_str("#include <cuda_runtime.h>\nint main(){}\n");
        let h = u.nodes.iter().find_map(|n| match n {
            IrNode::HeaderInclude { header, .. } => Some(header.clone()),
            _ => None,
        });
        assert_eq!(h.as_deref(), Some("cuda_runtime.h"));
    }

    #[test]
    fn kernel_launch_is_captured() {
        let u = translate_str("foo<<<grid, block>>>(p);");
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::KernelLaunch { kernel, grid, block, .. }
            if kernel == "foo" && grid == "grid" && block == "block"
        )));
    }

    #[test]
    fn finds_inline_ptx_asm() {
        let u = translate_str(r#"
            __global__ void k() {
                asm("mov.u32 %0, %1;" : "=r"(x) : "r"(y));
            }
        "#);
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::Warning { message, .. }
            if message.contains("PTX")
        )), "expected a PTX warning node");
    }

    #[test]
    fn finds_asm_volatile() {
        let u = translate_str(r#"
            __global__ void k() {
                asm volatile("membar.gl;");
            }
        "#);
        assert!(u.nodes.iter().any(|n| matches!(
            n,
            IrNode::Warning { message, .. }
            if message.contains("PTX")
        )), "expected a PTX warning for asm volatile");
    }

    #[test]
    fn skips_asm_in_comment() {
        let u = translate_str(r#"
            // asm("this is a comment");
            __global__ void k() {}
        "#);
        assert!(!u.nodes.iter().any(|n| matches!(
            n,
            IrNode::Warning { message, .. }
            if message.contains("PTX")
        )), "should not flag asm inside comments");
    }
}