decuda 0.1.1

//! Source-text rewriting engine: applies IR-node span substitutions to a
//! source string for a given target backend.

use crate::cli::Target;
use crate::ir::{CudaQualifier, IrNode, TranslationUnit};
use crate::targets::TargetBackend;

/// Emit `unit` to `target`. Returns the rewritten source plus a list of
/// warnings produced during the rewrite (structurally the rewriter also has
/// the option of mutating the IR with `Warning` nodes up-front).
pub fn emit(target: Target, backend: &dyn TargetBackend, unit: &TranslationUnit) -> String {
    let source = unit.source.clone();
    let banner = backend.banner(&source);
    let replaced = apply_replacements(target, unit);
    format!("{banner}{replaced}")
}

/// Apply the span replacements for `target` against `unit.source`, returning
/// the rewritten text. Replacements are applied in source order, tracking
/// the cumulative byte shift so each edit lands at its original semantic
/// position regardless of earlier edits' length differences.
///
/// Overlapping edits are detected and the later-starting edit is dropped;
/// the earlier one wins. This protects against, e.g., a `RuntimeCall` for
/// `cuda_runtime` accidentally landing inside a `HeaderInclude` line.
pub fn apply_replacements(target: Target, unit: &TranslationUnit) -> String {
    let mut s = unit.source.clone();
    let mut edits: Vec<(usize, usize, String)> = Vec::with_capacity(unit.nodes.len());
    for node in &unit.nodes {
        if let Some(repl) = replacement_for(target, node) {
            let end = effective_end(target, node, &s);
            edits.push((node.start(), end, repl));
        }
    }
    // Drop overlapping edits: keep the one with the smallest start; among
    // ties, the longer span wins.
    edits.sort_by_key(|(start, end, _)| (*start, std::cmp::Reverse(*end)));
    let mut filtered: Vec<(usize, usize, String)> = Vec::with_capacity(edits.len());
    let mut last_end: usize = 0;
    for (start, end, text) in edits {
        if start >= last_end {
            filtered.push((start, end, text));
            last_end = end;
        } else {
            // Overlaps a previously kept edit: drop.
        }
    }
    edits = filtered;
    edits.sort_by_key(|(start, _, _)| *start);
    let mut acc_delta: i64 = 0;
    for (start, end, text) in edits {
        let orig_len = (end - start) as i64;
        let adj_start = (start as i64 + acc_delta).max(0) as usize;
        let adj_end = (end as i64 + acc_delta).max(0) as usize;
        if adj_end > s.len() || adj_start > adj_end {
            continue;
        }
        s.replace_range(adj_start..adj_end, &text);
        acc_delta += text.len() as i64 - orig_len;
    }
    s
}

/// Decide what (if anything) to replace for a given node and target. Returns
/// `None` if the node needs no rewrite for the given target.
pub fn replacement_for(target: Target, node: &IrNode) -> Option<String> {
    let slot = crate::ir::slot_index(target);
    match node {
        IrNode::QualifierDecl { qualifier, .. } => qualifier_rewrite(*qualifier, target),
        IrNode::KernelLaunch {
            kernel,
            grid,
            block,
            smem,
            stream,
            args,
            ..
        } => Some(launch_rewrite(target, kernel, grid, block, smem, stream, args)),
        IrNode::RuntimeCall { name: _, args, mappings, .. } => {
            // When a mapping exists for this target, swap in the mapped name
            // and adapt arguments (no arg renaming for v1). When there is no
            // mapping, leave the original call untouched: the migration report
            // carries the warning, and cluttering every unmapped call with a
            // `TODO` comment makes the output unreadable.
            if let Some(mapped) = mappings[slot].clone() {
                let args_joined = args.join(", ");
                Some(format!("{mapped}({args_joined})"))
            } else {
                None
            }
        }
        IrNode::BuiltinRef { kind, .. } => Some(builtin_rewrite(*kind, target)),
        IrNode::HeaderInclude { header, replacements, .. } => {
            replacements[slot].clone().map(|mapped| match target {
                Target::Rust => format!("// was: #include {header}  ->  {mapped}"),
                _ => format!("#include {mapped} /* was: {header} */"),
            })
        }
        IrNode::AtomicIntrinsic { name, .. } => Some(atomic_rewrite(name, target)),
        IrNode::KernelDef { .. } => None, // kernel body is rewritten in-place
        IrNode::Warning { .. } => None,
    }
}

/// Effective end-byte of a node for the given target. For HIP, the
/// trailing `.x` / `.y` / `.z` field access is preserved (HIP's builtins
/// are dim3-typed and the field access is still valid). For every other
/// target the field access must be consumed because the replacement is
/// a scalar function call.
pub fn effective_end(target: Target, node: &IrNode, source: &str) -> usize {
    if let IrNode::BuiltinRef {
        end,
        has_field_access: true,
        ..
    } = node
    {
        if target != Target::Hip {
            let bytes = source.as_bytes();
            let mut p = *end;
            while p < bytes.len() && bytes[p].is_ascii_whitespace() {
                p += 1;
            }
            if p + 1 < bytes.len()
                && bytes[p] == b'.'
                && matches!(bytes[p + 1], b'x' | b'y' | b'z')
                && (p + 2 == bytes.len() || !(bytes[p + 2].is_ascii_alphanumeric() || bytes[p + 2] == b'_'))
            {
                return p + 2;
            }
        }
    }
    node.end()
}

fn qualifier_rewrite(q: CudaQualifier, target: Target) -> Option<String> {
    use CudaQualifier::*;
    let s = match (q, target) {
        // __global__ -> kernel qualifier in the target language.
        (Global, Target::Hip) => "__global__",
        (Global, Target::Sycl) => "// TODO(decuda): rewrite as SYCL kernel lambda\n",
        (Global, Target::Rust) => "// TODO(decuda): rewrite as rust-gpu kernel fn\n",
        (Global, Target::Opencl) => "__kernel",

        // __device__ -> device-side function qualifier.
        (Device, Target::Hip) => "__device__",
        (Device, Target::Sycl) => "// SYCL device function",
        (Device, Target::Rust) => "/* device function */",
        (Device, Target::Opencl) => "__device",

        // __host__ -> just a normal function in every target.
        (Host, Target::Hip) => "__host__",
        (Host, Target::Sycl) => "",
        (Host, Target::Rust) => "",
        (Host, Target::Opencl) => "",

        // __forceinline__ / __noinline__: language inline hints.
        (ForceInline, Target::Hip) => "__forceinline__",
        (ForceInline, Target::Sycl) => "[[clang::always_inline]]",
        (ForceInline, Target::Rust) => "#[inline(always)]",
        (ForceInline, Target::Opencl) => "",

        (NoInline, Target::Hip) => "__noinline__",
        (NoInline, Target::Sycl) => "[[gnu::noinline]]",
        (NoInline, Target::Rust) => "#[inline(never)]",
        (NoInline, Target::Opencl) => "",

        // Storage qualifiers.
        (Shared, Target::Hip) => "__shared__",
        (Shared, Target::Sycl) => "__shared__", // SYCL has accessors; we keep the keyword for visibility
        (Shared, Target::Rust) => "/* shared -> rust-gpu group_memory */",
        (Shared, Target::Opencl) => "__local",

        (Constant, Target::Hip) => "__constant__",
        (Constant, Target::Sycl) => "/* constant -> SYCL constant_accessor */",
        (Constant, Target::Rust) => "/* constant */",
        (Constant, Target::Opencl) => "__constant",

        (Managed, Target::Hip) => "__managed__",
        (Managed, Target::Sycl) => "/* managed -> SYCL USM */",
        (Managed, Target::Rust) => "/* managed -> unified memory */",
        (Managed, Target::Opencl) => "/* use SVM: clSVMAlloc */",

        (Restricted, Target::Hip) => "__restrict__",
        (Restricted, Target::Sycl) => "",
        (Restricted, Target::Rust) => "",
        (Restricted, Target::Opencl) => "restrict",

        (Texture, _) => return None,
        (Surface, _) => return None,
        (LaunchBounds, _) => return None,
        (ClusterDim, _) => return None,
        (GridDim, _) => return None,
        (Const, _) => "",
        (Pinned, _) => "",
        (_, Target::All) => "",
    };
    Some(s.to_string())
}

fn launch_rewrite(
    target: Target,
    kernel: &str,
    grid: &str,
    block: &str,
    smem: &Option<String>,
    stream: &Option<String>,
    args: &[String],
) -> String {
    let args_joined = args.join(", ");
    match target {
        Target::Hip => {
            let smem = smem.as_deref().unwrap_or("0");
            let stream = stream.as_deref().unwrap_or("0");
            format!(
                "hipLaunchKernelGGL({kernel}, dim3({grid}), dim3({block}), {smem}, {stream}, {args_joined})"
            )
        }
        Target::Sycl => {
            // SYCL replaces kernel launches with a queue.submit + parallel_for;
            // emit a starter block the user can adapt to their accessors.
            let smem = smem.as_deref().unwrap_or("none");
            let stream = stream.as_deref().unwrap_or("default");
            format!(
                "{{ /* decuda SYCL launch: queue.submit([&](sycl::handler& h) {{ h.parallel_for(sycl::range<3>{{{grid}}}, [=](sycl::item<3> it) {{ /* kernel `{kernel}` body with thread indices from it.get_*() */ }}); }}); smem={smem} stream={stream} args={args_joined} */ }}"
            )
        }
        Target::Rust => {
            // cust kernel launch.
            let smem = smem.as_deref().unwrap_or("0");
            let stream_v = stream.as_deref().unwrap_or("default");
            format!(
                "{{ /* decuda cust launch */ let _kernel = modules.get_function(\"{kernel}\"); unsafe {{ let _ = launch!( _kernel<<<{grid} as grid_size, {block} as block_size, {smem} as usize, {stream_v}>>>({args_joined}) ); }} }}"
            )
        }
        Target::Opencl => {
            format!(
                "clEnqueueNDRangeKernel(queue, {kernel}_kernel, 1, NULL, (size_t[1]){{{grid}}}, (size_t[1]){{{block}}}, 0, NULL, NULL) /* args: {args_joined} */",
            )
        }
        Target::All => format!("{kernel}<<<{grid}, {block}>>>({args_joined})"),
    }
}

fn builtin_rewrite(kind: crate::ir::BuiltinKind, target: Target) -> String {
    use crate::ir::BuiltinKind::*;
    let s = match (kind, target) {
        (ThreadIdx, Target::Hip) => "threadIdx",
        (ThreadIdx, Target::Sycl) => "item.get_local_id()",
        (ThreadIdx, Target::Rust) => "thread_idx",
        (ThreadIdx, Target::Opencl) => "get_local_id(0)",

        (BlockIdx, Target::Hip) => "blockIdx",
        (BlockIdx, Target::Sycl) => "item.get_group(0)",
        (BlockIdx, Target::Rust) => "block_idx",
        (BlockIdx, Target::Opencl) => "get_group_id(0)",

        (BlockDim, Target::Hip) => "blockDim",
        (BlockDim, Target::Sycl) => "item.get_local_range()",
        (BlockDim, Target::Rust) => "block_dim",
        (BlockDim, Target::Opencl) => "get_local_size(0)",

        (GridDim, Target::Hip) => "gridDim",
        (GridDim, Target::Sycl) => "item.get_global_range() / item.get_local_range()",
        (GridDim, Target::Rust) => "grid_dim",
        (GridDim, Target::Opencl) => "get_num_groups(0)",

        (WarpSize, Target::Hip) => "warpSize",
        (WarpSize, Target::Sycl) => "32 /* warpSize */",
        (WarpSize, Target::Rust) => "WARP_SIZE",
        (WarpSize, Target::Opencl) => "warpSize /* CL_DEVICE_WARP_SIZE_NV */",

        (SyncThreads, Target::Hip) => "__syncthreads()",
        (SyncThreads, Target::Sycl) => "item.barrier(sycl::access::fence_space::global_space)",
        (SyncThreads, Target::Rust) => "group.sync()",
        (SyncThreads, Target::Opencl) => "barrier(CLK_LOCAL_MEM_FENCE)",

        (SyncWarp, Target::Hip) => "__syncwarp(0xFFFFFFFFu)",
        (SyncWarp, Target::Sycl) => "/* no native syncwarp on SYCL */ item.barrier() /* fallback */",
        (SyncWarp, Target::Rust) => "/* syncwarp: only lane=0 of warp at once */",
        (SyncWarp, Target::Opencl) => "barrier(CLK_LOCAL_MEM_FENCE) /* approx */",

        (FsyncBlock, Target::Hip) => "__sync_block()",
        (FsyncBlock, Target::Sycl) => "item.barrier()",
        (FsyncBlock, Target::Rust) => "group.sync()",
        (FsyncBlock, Target::Opencl) => "barrier(CLK_LOCAL_MEM_FENCE)",

        (LaneId, Target::Hip) => "__laneid()",
        (LaneId, Target::Sycl) => "item.get_sub_group().get_local_id()",
        (LaneId, Target::Rust) => "lane_id",
        (LaneId, Target::Opencl) => "get_sub_group_id() * get_sub_group_size() + get_sub_group_local_id()",

        (WarpId, Target::Hip) => "__warp_id()",
        (WarpId, Target::Sycl) => "item.get_sub_group().get_group_id()",
        (WarpId, Target::Rust) => "warp_id",
        (WarpId, Target::Opencl) => "get_sub_group_id()",

        // Warp primitives: HIP is identical to CUDA (no-op rewrite).
        // OpenCL/SYCL swap the function name but argument order/count may
        // differ — the migration report flags this for manual review.
        // Rust has no direct equivalent.
        (ShflSync, Target::Hip) => "__shfl_sync",
        (ShflSync, Target::Sycl) => "sycl::sub_group::shuffle",
        (ShflSync, Target::Rust) => "/* TODO: __shfl_sync */",
        (ShflSync, Target::Opencl) => "sub_group_shuffle",

        (BallotSync, Target::Hip) => "__ballot_sync",
        (BallotSync, Target::Sycl) => "sycl::sub_group::ballot",
        (BallotSync, Target::Rust) => "/* TODO: __ballot_sync */",
        (BallotSync, Target::Opencl) => "sub_group_ballot",

        (AnySync, Target::Hip) => "__any_sync",
        (AnySync, Target::Sycl) => "sycl::sub_group::any",
        (AnySync, Target::Rust) => "/* TODO: __any_sync */",
        (AnySync, Target::Opencl) => "sub_group_any",

        (AllSync, Target::Hip) => "__all_sync",
        (AllSync, Target::Sycl) => "sycl::sub_group::all",
        (AllSync, Target::Rust) => "/* TODO: __all_sync */",
        (AllSync, Target::Opencl) => "sub_group_all",

        (ActiveMask, Target::Hip) => "__activemask",
        (ActiveMask, Target::Sycl) => "sycl::sub_group::get_local_range",
        (ActiveMask, Target::Rust) => "/* TODO: __activemask */",
        (ActiveMask, Target::Opencl) => "get_sub_group_size",

        (_, Target::All) => "/* decuda: unsupported builtin */",
    };
    s.to_string()
}

fn atomic_rewrite(name: &str, target: Target) -> String {
    match target {
        // HIP and OpenCL keep the CUDA atomic intrinsic names unchanged.
        Target::Hip | Target::Opencl => name.to_string(),
        // SYCL has a different atomic API; rust-gpu also needs different
        // syntax. Leave the name as-is so the file at least parses; the
        // migration report flags the line for manual work.
        Target::Sycl | Target::Rust => name.to_string(),
        Target::All => name.to_string(),
    }
}

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

    #[test]
    fn global_qualifier_to_opencl_kernel() {
        let s = qualifier_rewrite(CudaQualifier::Global, Target::Opencl).unwrap();
        assert_eq!(s, "__kernel");
    }

    #[test]
    fn threadidx_rewrite_per_target() {
        assert_eq!(builtin_rewrite(BuiltinKind::ThreadIdx, Target::Hip), "threadIdx");
        assert_eq!(
            builtin_rewrite(BuiltinKind::ThreadIdx, Target::Opencl),
            "get_local_id(0)"
        );
    }

    #[test]
    fn syncthreads_rewrites_globally() {
        assert!(builtin_rewrite(BuiltinKind::SyncThreads, Target::Sycl).contains("barrier"));
    }

    #[test]
    fn shfl_sync_per_target() {
        // HIP: identical to CUDA (no-op name swap).
        assert_eq!(builtin_rewrite(BuiltinKind::ShflSync, Target::Hip), "__shfl_sync");
        // OpenCL: sub_group_shuffle (args preserved, may need manual fixup).
        assert_eq!(builtin_rewrite(BuiltinKind::ShflSync, Target::Opencl), "sub_group_shuffle");
        // SYCL: sycl::sub_group::shuffle.
        assert!(builtin_rewrite(BuiltinKind::ShflSync, Target::Sycl).contains("shuffle"));
        // Rust: no equivalent.
        assert!(builtin_rewrite(BuiltinKind::ShflSync, Target::Rust).contains("TODO"));
    }

    #[test]
    fn ballot_sync_per_target() {
        assert_eq!(builtin_rewrite(BuiltinKind::BallotSync, Target::Hip), "__ballot_sync");
        assert_eq!(builtin_rewrite(BuiltinKind::BallotSync, Target::Opencl), "sub_group_ballot");
    }

    #[test]
    fn activemask_per_target() {
        assert_eq!(builtin_rewrite(BuiltinKind::ActiveMask, Target::Hip), "__activemask");
        assert_eq!(builtin_rewrite(BuiltinKind::ActiveMask, Target::Opencl), "get_sub_group_size");
    }
}