#!/bin/bash
# CUDA Kernel Guard — detects CPU/host leaks in .cu kernel files.
#
# Production H100 system. CUDA kernels must be pure device code.
# This guard catches patterns that force CPU synchronization,
# host memory access, or debug output in production kernels.
#
# Usage:
#   scripts/cuda-kernel-guard.sh <file.cu>
#
# Exit codes:
#   0 = clean
#   1 = violation detected
#
# Suppress reviewed lines with:  // cuda-ok: <reason>

set -euo pipefail

FILE="${1:-}"
if [ -z "$FILE" ] || [ ! -f "$FILE" ]; then
    exit 0
fi

# Only check .cu and .cuh files
case "$FILE" in
    *.cu|*.cuh) ;;
    *) exit 0 ;;
esac

LEAK_PATTERNS=(
    # ─── printf in device code (forces implicit sync, kills perf) ───
    # One printf in a kernel with 1024 threads = 1024 serial host writes.
    'printf\s*\('

    # ─── Host API calls from device code (illegal or sync-forcing) ───
    'cudaMalloc'
    'cudaFree'
    'cudaMemcpy\s*\('
    'cudaMemset\s*\('
    'cudaDeviceSynchronize'
    'cudaStreamSynchronize'
    'cudaEventSynchronize'

    # ─── Host memory access patterns ───
    # malloc/free/new/delete in device code = UB or host fallback
    '\bmalloc\s*\('
    '\bfree\s*\('
    '\bnew\s+\w+[\[\(]'
    '\bdelete\s+\w+'

    # ─── Host function markers in kernel files ───
    # __host__ functions shouldn't be in kernel .cu files (they belong in .rs host code)
    '__host__\s+[^_]'

    # ─── CPU-side file I/O (absolutely forbidden in kernels) ───
    '\bfopen\s*\('
    '\bfwrite\s*\('
    '\bfprintf\s*\('
    '\bstdout\b'
    '\bstderr\b'

    # ─── assert() in device code (triggers host trap, kills kernel) ───
    # Use __trap() or return early instead
    '\bassert\s*\('

    # ─── Double-precision math in f32 kernels (use f suffix) ───
    # exp/log/sin/cos are DOUBLE precision (4x slower on H100 tensor cores).
    # expf/logf/sinf/cosf are single-precision (correct for f32 kernels).
    # __expf/__logf are fast-math (~2 ULP) — optional, not required.
    # Only flag the double-precision versions (no f suffix, not prefixed with __)
    '[^_a-z]\bexp\s*\([^f]'
    '[^_a-z]\blog\s*\([^f]'
    '[^_a-z]\bsin\s*\([^f]'
    '[^_a-z]\bcos\s*\([^f]'
    '[^_a-z]\bsqrt\s*\([^f]'
)

# Patterns that are OK (suppress false positives)
SAFE_PATTERNS=(
    '// cuda-ok:'
    '// debug-only:'
    '#if 0'
    '#ifdef CUDA_DEBUG'
)

found=0

for pattern in "${LEAK_PATTERNS[@]}"; do
    hits=$(grep -nE "$pattern" "$FILE" 2>/dev/null || true)

    # Filter out safe patterns
    for safe in "${SAFE_PATTERNS[@]}"; do
        if [ -n "$hits" ]; then
            hits=$(echo "$hits" | grep -v "$safe" || true)
        fi
    done

    # Filter out comments: full-line (//, /*, *) and trailing /* */ comments
    if [ -n "$hits" ]; then
        hits=$(echo "$hits" | grep -v '[0-9]*:\s*//' \
                            | grep -v '[0-9]*:\s*\*' \
                            | grep -v '[0-9]*:\s*/\*' \
                            || true)
        # For trailing comments: if the match is ONLY inside /* ... */, skip it
        if [ -n "$hits" ]; then
            hits=$(echo "$hits" | while IFS= read -r line; do
                # Strip trailing /* comment */ from the line, then re-check pattern
                stripped=$(echo "$line" | sed 's|/\*[^*]*\*/||g')
                if echo "$stripped" | grep -qE "$pattern" 2>/dev/null; then
                    echo "$line"
                fi
            done || true)
        fi
    fi

    if [ -n "$hits" ]; then
        if [ "$found" -eq 0 ]; then
            echo "CUDA KERNEL LEAK: $FILE"
            found=1
        fi
        echo "$hits" | while IFS= read -r line; do
            echo "  $line"
        done
    fi
done

if [ "$found" -eq 1 ]; then
    echo ""
    echo "⛔ CPU/host operations detected in CUDA kernel."
    echo "   printf → remove or guard with #ifdef CUDA_DEBUG"
    echo "   cudaMalloc/Free → pre-allocate in host .rs code"
    echo "   sinf/cosf/expf/logf → use __sinf/__cosf/__expf/__logf"
    echo "   assert → use __trap() or return"
    exit 1
fi

exit 0