/**
 * PoC: compute_numel() integer overflow on MSVC (Windows) builds
 *
 * PyTorch's compute_numel() uses two different code paths:
 *   - GCC/Clang desktop: safe_compute_numel() with overflow detection
 *   - MSVC (Windows) / Mobile: multiply_integers() with NO overflow detection
 *
 * This test demonstrates that multiply_integers() silently overflows,
 * producing a wrong numel value that leads to undersized allocations.
 *
 * Compile and run:
 *   g++ -std=c++17 -fsanitize=undefined -o test_msvc_overflow test_msvc_overflow.cpp
 *   ./test_msvc_overflow
 *
 * Expected: shows UBSan signed-integer-overflow (undefined behavior)
 * and demonstrates that multiply_integers returns a wrong (overflowed) value.
 */

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <numeric>
#include <vector>

// Exact copy of multiply_integers from c10/util/accumulate.h
// This is the UNSAFE path used on MSVC builds
int64_t multiply_integers(const std::vector<int64_t>& container) {
    return std::accumulate(
        container.begin(),
        container.end(),
        static_cast<int64_t>(1),
        std::multiplies<>());
}

// Simplified safe_multiplies_u64 from c10/util/safe_numerics.h
// This is used by safe_compute_numel on GCC/Clang desktop
bool safe_multiplies_u64(const std::vector<int64_t>& dims, uint64_t* out) {
    uint64_t prod = 1;
    int total_bits = 0;
    bool is_zero = false;
    for (auto d : dims) {
        auto x = static_cast<uint64_t>(d);
        prod *= x;
        is_zero |= (x == 0);
        // Count bits needed
        int bits = 0;
        uint64_t v = x;
        while (v > 1) { v >>= 1; bits++; }
        if (x > 0) total_bits += bits + 1;
    }
    *out = prod;
    return !is_zero && (total_bits >= 64);
}

int main() {
    printf("=== compute_numel() MSVC vs GCC/Clang Overflow Test ===\n\n");

    // Test case: dimensions that cause int64 overflow
    // shape = [4611686018427387904, 4] = [2^62, 4]
    // product = 2^62 * 4 = 2^64 = overflow!
    int64_t dim0 = INT64_C(4611686018427387904);  // 2^62
    int64_t dim1 = 4;
    std::vector<int64_t> dims = {dim0, dim1};

    printf("Dimensions: [%ld, %ld]\n", (long)dim0, (long)dim1);
    printf("Expected numel: 2^62 * 4 = 2^64 (overflows int64)\n\n");

    // Path 1: MSVC path (multiply_integers - NO overflow check)
    printf("--- MSVC Path: multiply_integers() ---\n");
    int64_t msvc_numel = multiply_integers(dims);
    printf("  Result: %ld\n", (long)msvc_numel);
    printf("  Overflow detected: NO (multiply_integers has no detection)\n");
    if (msvc_numel == 0) {
        printf("  numel wrapped to 0 → zero-size allocation!\n");
    } else {
        printf("  numel wrapped to %ld → undersized allocation!\n", (long)msvc_numel);
    }
    printf("\n");

    // Path 2: GCC/Clang path (safe_multiplies_u64 - HAS overflow check)
    printf("--- GCC/Clang Path: safe_multiplies_u64() ---\n");
    uint64_t safe_numel;
    bool overflow = safe_multiplies_u64(dims, &safe_numel);
    printf("  Result: %lu\n", (unsigned long)safe_numel);
    printf("  Overflow detected: %s\n", overflow ? "YES → would throw TORCH_CHECK" : "NO");
    printf("\n");

    // Demonstrate the security impact
    printf("=== Security Impact ===\n\n");
    printf("On MSVC (Windows):\n");
    printf("  compute_numel() → multiply_integers() → numel = %ld\n", (long)msvc_numel);
    printf("  Tensor claims shape [2^62, 4] but numel_ = %ld\n", (long)msvc_numel);
    printf("  Operations using numel_ allocate %ld * itemsize bytes\n", (long)msvc_numel);
    printf("  But actual tensor spans 2^64 elements → HEAP BUFFER OVERFLOW\n\n");

    printf("On GCC/Clang desktop:\n");
    printf("  compute_numel() → safe_compute_numel() → overflow detected!\n");
    printf("  TORCH_CHECK throws: 'numel: integer multiplication overflow'\n");
    printf("  Model fails to load safely.\n\n");

    // Show the conditional compilation
    printf("=== Root Cause (TensorImpl.h:2596-2604) ===\n\n");
    printf("  int64_t compute_numel() const {\n");
    printf("  #if C10_HAS_BUILTIN_OVERFLOW() && !defined(C10_MOBILE)\n");
    printf("      return safe_compute_numel();   // GCC/Clang: SAFE\n");
    printf("  #else\n");
    printf("      return multiply_integers(...);  // MSVC/Mobile: UNSAFE!\n");
    printf("  #endif\n");
    printf("  }\n\n");

    printf("  C10_HAS_BUILTIN_OVERFLOW is 0 on MSVC (safe_numerics.h:10)\n");
    printf("  → ALL Windows PyTorch users are vulnerable\n");

    return 0;
}