decuda 0.1.1

//! Intermediate representation produced by the CUDA parser and consumed by
//! every target backend.

#![allow(dead_code)]

use std::path::PathBuf;

use crate::cli::Target;

/// A single translation unit (one source file).
#[derive(Debug, Clone)]
pub struct TranslationUnit {
    pub source: String,
    pub path: PathBuf,
    /// Nodes in source order; spans refer to byte offsets into `source`.
    pub nodes: Vec<IrNode>,
}

impl TranslationUnit {
    pub fn new(source: String, path: PathBuf) -> Self {
        Self {
            source,
            path,
            nodes: Vec::new(),
        }
    }

    pub fn push(&mut self, node: IrNode) {
        self.nodes.push(node);
    }

    /// Print a per-line summary of the IR for the `inspect` command.
    pub fn print_summary(&self) {
        if self.nodes.is_empty() {
            println!("  (empty IR)");
            return;
        }
        for n in &self.nodes {
            let line = line_for_offset(&self.source, n.start());
            println!("  line {line:>4}: {}", n.describe());
        }
    }
}

/// A single observation about the source. Carries the source span so the
/// rewriter can replace it byte-precisely.
#[derive(Debug, Clone)]
pub enum IrNode {
    /// Function/storage qualifier such as `__global__`.
    QualifierDecl {
        start: usize,
        end: usize,
        qualifier: CudaQualifier,
        surface: String,
    },
    /// Kernel definition body boundary markers; the rewriter will copy the
    /// kernel body into the target as-is, only transforming indices/qualifiers.
    KernelDef {
        start: usize,
        end: usize,
        name: String,
        params: Vec<String>,
    },
    /// CUDA kernel launch (`kernel<<<grid, block, smem, stream>>>(args)`) that
    /// has been normalized into an ordinary call to `__decuda_launch(...)`.
    KernelLaunch {
        start: usize,
        end: usize,
        kernel: String,
        grid: String,
        block: String,
        smem: Option<String>,
        stream: Option<String>,
        args: Vec<String>,
    },
    /// A known CUDA runtime / driver API call (e.g. `cudaMalloc`).
    RuntimeCall {
        start: usize,
        end: usize,
        name: String,
        args: Vec<String>,
        /// Per-target mappings, computed lazily.
        mappings: [Option<String>; 4],
    },
    /// Reference to a CUDA built-in variable (`threadIdx.x`, etc.) or
    /// synchronization intrinsic (`__syncthreads`).
    BuiltinRef {
        start: usize,
        end: usize,
        kind: BuiltinKind,
        /// True if the identifier is followed by a `.x`/`.y`/`.z` field
        /// access in the source. Targets whose replacement is a scalar
        /// function (e.g. OpenCL) need to also rewrite the trailing `.x`.
        has_field_access: bool,
    },
    /// Header include directive.
    HeaderInclude {
        start: usize,
        end: usize,
        header: String,
        /// Replacement header per target; index by Target::iter_real offset.
        replacements: [Option<String>; 4],
    },
    /// Atomic intrinsic such as `atomicAdd`.
    AtomicIntrinsic {
        start: usize,
        end: usize,
        name: String,
    },
    /// A CUDA construct we do not know how to translate: we copy the bytes
    /// through and emit a warning in the migration report.
    Warning {
        start: usize,
        end: usize,
        surface: String,
        message: String,
    },
}

impl IrNode {
    pub fn start(&self) -> usize {
        match self {
            IrNode::QualifierDecl { start, .. }
            | IrNode::KernelDef { start, .. }
            | IrNode::KernelLaunch { start, .. }
            | IrNode::RuntimeCall { start, .. }
            | IrNode::BuiltinRef { start, .. }
            | IrNode::HeaderInclude { start, .. }
            | IrNode::AtomicIntrinsic { start, .. }
            | IrNode::Warning { start, .. } => *start,
        }
    }

    pub fn end(&self) -> usize {
        match self {
            IrNode::QualifierDecl { end, .. }
            | IrNode::KernelDef { end, .. }
            | IrNode::KernelLaunch { end, .. }
            | IrNode::RuntimeCall { end, .. }
            | IrNode::BuiltinRef { end, .. }
            | IrNode::HeaderInclude { end, .. }
            | IrNode::AtomicIntrinsic { end, .. }
            | IrNode::Warning { end, .. } => *end,
        }
    }

    pub fn kind_label(&self) -> &'static str {
        match self {
            IrNode::QualifierDecl { .. } => "qualifier",
            IrNode::KernelDef { .. } => "kernel_def",
            IrNode::KernelLaunch { .. } => "kernel_launch",
            IrNode::RuntimeCall { .. } => "runtime_call",
            IrNode::BuiltinRef { .. } => "builtin",
            IrNode::HeaderInclude { .. } => "include",
            IrNode::AtomicIntrinsic { .. } => "atomic",
            IrNode::Warning { .. } => "warning",
        }
    }

    pub fn describe(&self) -> String {
        match self {
            IrNode::QualifierDecl { qualifier, surface, .. } => {
                format!("qualifier {qualifier:?} `{surface}`")
            }
            IrNode::KernelDef { name, .. } => format!("kernel_def `{name}`"),
            IrNode::KernelLaunch {
                kernel, grid, block, ..
            } => format!("launch `{kernel}` <<<{grid}, {block}>>>"),
            IrNode::RuntimeCall { name, .. } => format!("runtime `{name}`"),
            IrNode::BuiltinRef { kind, .. } => format!("builtin {kind:?}"),
            IrNode::HeaderInclude { header, .. } => format!("include `{header}`"),
            IrNode::AtomicIntrinsic { name, .. } => format!("atomic `{name}`"),
            IrNode::Warning { message, .. } => format!("WARN: {message}"),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CudaQualifier {
    Global,
    Device,
    Host,
    ForceInline,
    NoInline,
    Const,
    Restricted,
    Shared,
    Constant,
    Managed,
    Pinned,
    Texture,
    Surface,
    LaunchBounds,
    ClusterDim,
    GridDim,
}

impl CudaQualifier {
    pub fn from_token(tok: &str) -> Option<Self> {
        Some(match tok {
            "__global__" => CudaQualifier::Global,
            "__device__" => CudaQualifier::Device,
            "__host__" => CudaQualifier::Host,
            "__forceinline__" => CudaQualifier::ForceInline,
            "__noinline__" => CudaQualifier::NoInline,
            "__restrict__" => CudaQualifier::Restricted,
            "__shared__" => CudaQualifier::Shared,
            "__constant__" => CudaQualifier::Constant,
            "__managed__" => CudaQualifier::Managed,
            // __launch_bounds__ and __cluster_dim__ are parsed but flagged
            // because they need careful, target-specific rewriting.
            _ => return None,
        })
    }

    pub fn from_storage_class(tok: &str) -> Option<Self> {
        match tok {
            "__shared__" => Some(CudaQualifier::Shared),
            "__constant__" => Some(CudaQualifier::Constant),
            "__managed__" => Some(CudaQualifier::Managed),
            "__restrict__" => Some(CudaQualifier::Restricted),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BuiltinKind {
    ThreadIdx,
    BlockIdx,
    BlockDim,
    GridDim,
    WarpSize,
    SyncThreads,
    SyncWarp,
    FsyncBlock,
    LaneId,
    WarpId,
    // Warp-level primitives (function-like, arguments preserved).
    ShflSync,
    BallotSync,
    AnySync,
    AllSync,
    ActiveMask,
}

impl BuiltinKind {
    pub fn from_token(tok: &str) -> Option<Self> {
        Some(match tok {
            "threadIdx" => BuiltinKind::ThreadIdx,
            "blockIdx" => BuiltinKind::BlockIdx,
            "blockDim" => BuiltinKind::BlockDim,
            "gridDim" => BuiltinKind::GridDim,
            "warpSize" => BuiltinKind::WarpSize,
            "__syncthreads" => BuiltinKind::SyncThreads,
            "__syncwarp" => BuiltinKind::SyncWarp,
            "__sync_block" => BuiltinKind::FsyncBlock,
            "__laneid" => BuiltinKind::LaneId,
            "__warp_id" => BuiltinKind::WarpId,
            "__shfl_sync" => BuiltinKind::ShflSync,
            "__ballot_sync" => BuiltinKind::BallotSync,
            "__any_sync" => BuiltinKind::AnySync,
            "__all_sync" => BuiltinKind::AllSync,
            "__activemask" => BuiltinKind::ActiveMask,
            _ => return None,
        })
    }

    /// Returns true if this builtin is a sync intrinsic whose argument list
    /// (if any) should be consumed and replaced wholesale (e.g.
    /// `__syncthreads()` → `barrier(CLK_LOCAL_MEM_FENCE)`).
    ///
    /// Returns false for warp primitives whose arguments are meaningful and
    /// must be preserved — only the function name is rewritten.
    pub fn consumes_args(self) -> bool {
        matches!(
            self,
            BuiltinKind::SyncThreads | BuiltinKind::SyncWarp | BuiltinKind::FsyncBlock
        )
    }
}

fn line_for_offset(source: &str, offset: usize) -> usize {
    let off = offset.min(source.len());
    1 + source[..off].bytes().filter(|b| *b == b'\n').count()
}

/// Lookup helper: given a Target, index into the per-target slots used by some
/// IR node variants.
pub fn slot_index(target: Target) -> usize {
    match target {
        Target::Hip => 0,
        Target::Sycl => 1,
        Target::Rust => 2,
        Target::Opencl => 3,
        Target::All => 0,
    }
}

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

    #[test]
    fn cuda_qualifier_round_trip() {
        assert_eq!(
            CudaQualifier::from_token("__global__"),
            Some(CudaQualifier::Global)
        );
        assert_eq!(CudaQualifier::from_token("nope"), None);
        assert_eq!(
            CudaQualifier::from_storage_class("__shared__"),
            Some(CudaQualifier::Shared)
        );
    }

    #[test]
    fn builtin_kind_round_trip() {
        assert_eq!(BuiltinKind::from_token("threadIdx"), Some(BuiltinKind::ThreadIdx));
        assert_eq!(BuiltinKind::from_token("nope"), None);
    }

    #[test]
    fn warp_primitive_builtins_round_trip() {
        assert_eq!(BuiltinKind::from_token("__shfl_sync"), Some(BuiltinKind::ShflSync));
        assert_eq!(BuiltinKind::from_token("__ballot_sync"), Some(BuiltinKind::BallotSync));
        assert_eq!(BuiltinKind::from_token("__any_sync"), Some(BuiltinKind::AnySync));
        assert_eq!(BuiltinKind::from_token("__all_sync"), Some(BuiltinKind::AllSync));
        assert_eq!(BuiltinKind::from_token("__activemask"), Some(BuiltinKind::ActiveMask));
    }

    #[test]
    fn consumes_args_only_for_sync_intrinsics() {
        // Sync intrinsics consume their argument list.
        assert!(BuiltinKind::SyncThreads.consumes_args());
        assert!(BuiltinKind::SyncWarp.consumes_args());
        assert!(BuiltinKind::FsyncBlock.consumes_args());
        // Warp primitives do NOT — only the function name is rewritten.
        assert!(!BuiltinKind::ShflSync.consumes_args());
        assert!(!BuiltinKind::BallotSync.consumes_args());
        assert!(!BuiltinKind::AnySync.consumes_args());
        assert!(!BuiltinKind::AllSync.consumes_args());
        assert!(!BuiltinKind::ActiveMask.consumes_args());
        // Variables like threadIdx also don't consume args.
        assert!(!BuiltinKind::ThreadIdx.consumes_args());
    }

    #[test]
    fn line_for_offset_basic() {
        let s = "a\nb\nc";
        assert_eq!(line_for_offset(s, 0), 1);
        assert_eq!(line_for_offset(s, 2), 2);
        assert_eq!(line_for_offset(s, 4), 3);
    }

    #[test]
    fn slot_indexes_unique() {
        let s: std::collections::HashSet<_> = Target::iter_real()
            .iter()
            .map(|t| slot_index(*t))
            .collect();
        assert_eq!(s.len(), 4);
    }
}