rialo-s-program-entrypoint 0.4.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

#![allow(clippy::bool_assert_comparison)]

use crate::allocator::*;

#[test]
fn test_basic_allocation() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();
        let ptr1 = alloc.alloc(layout);
        assert!(!ptr1.is_null());
        assert_eq!(ptr1.align_offset(8), 0);

        let ptr2 = alloc.alloc(layout);
        assert!(!ptr2.is_null());
        assert!(ptr2 as usize > ptr1 as usize);
    }
}

#[test]
fn test_dealloc_last() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();
        let ptr1 = alloc.alloc(layout);
        let used1 = alloc.used();

        let ptr2 = alloc.alloc(layout);
        let used2 = alloc.used();
        assert!(used2 > used1);

        // Deallocating last allocation should reclaim space
        alloc.dealloc(ptr2, layout);
        assert_eq!(alloc.used(), used1);

        // Deallocating non-last allocation does nothing
        alloc.alloc(layout);
        alloc.dealloc(ptr1, layout);
        assert!(alloc.used() > used1);
    }
}

#[test]
fn test_realloc_in_place() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();
        let ptr = alloc.alloc(layout);

        // Write some data
        core::ptr::write_bytes(ptr, 0x42, 32);

        // Reallocate in place (last allocation, growing)
        let new_ptr = alloc.realloc(ptr, layout, 64);
        assert_eq!(new_ptr, ptr); // Same pointer

        // Data should be preserved
        assert_eq!(*ptr, 0x42);

        // Reallocate in place (last allocation, shrinking)
        let shrink_ptr = alloc.realloc(new_ptr, Layout::from_size_align(64, 8).unwrap(), 32);
        assert_eq!(shrink_ptr, ptr); // Same pointer
    }
}

#[test]
fn test_realloc_non_last() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();
        let ptr1 = alloc.alloc(layout);
        core::ptr::write_bytes(ptr1, 0x11, 32);

        let _ptr2 = alloc.alloc(layout);

        // Reallocate non-last (shrinking should be no-op)
        let shrink_ptr = alloc.realloc(ptr1, layout, 16);
        assert_eq!(shrink_ptr, ptr1);
        assert_eq!(*ptr1, 0x11);

        // Reallocate non-last (growing requires copy)
        let grow_ptr = alloc.realloc(ptr1, layout, 64);
        assert_ne!(grow_ptr, ptr1); // Different pointer
        assert_eq!(*grow_ptr, 0x11); // Data copied
    }
}

#[test]
#[ignore] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
fn test_alignment() {
    let alloc = BumpAllocator::<()>::new_test(4096);

    unsafe {
        for align in [1, 2, 4, 8, 16, 32, 64, 128] {
            let layout = Layout::from_size_align(1, align).unwrap();
            let ptr = alloc.alloc(layout);
            assert!(!ptr.is_null());
            assert_eq!(ptr.align_offset(align), 0, "Failed alignment {}", align);
        }
    }
}

#[test]
fn test_oom() {
    let alloc = BumpAllocator::<()>::new_test(128);

    // First, let's understand what we have
    let header_size = core::mem::size_of::<Header<()>>();
    eprintln!("Header size: {}", header_size);
    eprintln!("Total heap: 128");
    eprintln!("Available after header: {}", 128 - header_size);

    unsafe {
        let layout = Layout::from_size_align(100, 8).unwrap();
        let ptr1 = alloc.alloc(layout);
        assert!(!ptr1.is_null());

        let used_after_first = alloc.used();
        eprintln!("Used after first alloc: {}", used_after_first);
        eprintln!("Remaining: {}", 128 - header_size - used_after_first);

        // Now allocate the remaining space
        let layout2 = Layout::from_size_align(50, 8).unwrap();
        let ptr2 = alloc.alloc(layout2);
        eprintln!("Second alloc result: {:?}", ptr2);
        assert!(ptr2.is_null(), "Second allocation should fail due to OOM");
    }
}

#[test]
fn test_many_small_allocations() {
    let alloc = BumpAllocator::<()>::new_test(4096);

    unsafe {
        let layout = Layout::from_size_align(8, 8).unwrap();
        let mut ptrs = Vec::new();

        // Allocate many small blocks
        for i in 0..100 {
            let ptr = alloc.alloc(layout);
            assert!(!ptr.is_null(), "Failed at allocation {}", i);
            core::ptr::write(ptr as *mut u64, i);
            ptrs.push(ptr);
        }

        // Verify data integrity
        for (i, &ptr) in ptrs.iter().enumerate() {
            assert_eq!(*(ptr as *const u64), i as u64);
        }
    }
}

#[test]
fn test_overflow_detection() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        // Allocate to near the end
        let layout = Layout::from_size_align(900, 8).unwrap();
        let ptr1 = alloc.alloc(layout);
        assert!(!ptr1.is_null());

        // This should properly detect overflow
        let huge_layout = Layout::from_size_align(u32::MAX as usize, 8).unwrap();
        let ptr2 = alloc.alloc(huge_layout);
        assert!(ptr2.is_null());
    }
}

#[test]
fn test_old_alloc_entire_heap() {
    // Old test: "alloc the entire"
    // Verifies we can allocate up to heap capacity (minus overhead)
    let heap_size = 128;
    let alloc = BumpAllocator::<()>::new_test(heap_size);

    unsafe {
        let layout = Layout::from_size_align(1, size_of::<u8>()).unwrap();
        let mut allocations = 0;

        // Allocate as much as possible
        loop {
            let ptr = alloc.alloc(layout);
            if ptr.is_null() {
                break;
            }
            allocations += 1;
        }

        // Should be able to allocate most of the heap
        // (accounting for header overhead in new allocator)
        let header_size = core::mem::size_of::<Header<()>>();
        let expected_min = heap_size - header_size - 8; // some slack
        assert!(
            allocations >= expected_min,
            "Only allocated {} bytes out of {} available",
            allocations,
            heap_size - header_size
        );

        // Next allocation should fail
        assert_eq!(alloc.alloc(layout), core::ptr::null_mut());
    }
}

#[test]
#[ignore] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
fn test_old_check_alignment() {
    // Old test: "check alignment"
    // Verifies all standard alignments work correctly
    let heap_size = 128;
    let alloc = BumpAllocator::<()>::new_test(heap_size);

    unsafe {
        // Test u8 alignment
        let ptr = alloc.alloc(Layout::from_size_align(1, size_of::<u8>()).unwrap());
        assert_eq!(0, ptr.align_offset(size_of::<u8>()));

        // Test u16 alignment
        let ptr = alloc.alloc(Layout::from_size_align(1, size_of::<u16>()).unwrap());
        assert_eq!(0, ptr.align_offset(size_of::<u16>()));

        // Test u32 alignment
        let ptr = alloc.alloc(Layout::from_size_align(1, size_of::<u32>()).unwrap());
        assert_eq!(0, ptr.align_offset(size_of::<u32>()));

        // Test u64 alignment
        let ptr = alloc.alloc(Layout::from_size_align(1, size_of::<u64>()).unwrap());
        assert_eq!(0, ptr.align_offset(size_of::<u64>()));

        // Test u128 alignment (BPF uses 8-byte alignment)
        let ptr = alloc.alloc(Layout::from_size_align(1, size_of::<u128>()).unwrap());
        assert_eq!(0, ptr.align_offset(size_of::<u128>()));

        // Test 64-byte alignment
        let ptr = alloc.alloc(Layout::from_size_align(1, 64).unwrap());
        assert_eq!(0, ptr.align_offset(64));
    }
}

#[test]
fn test_old_alloc_entire_block() {
    // Old test: "alloc entire block (minus the pos ptr)"
    // Verifies a single large allocation works
    let heap_size = 128;
    let alloc = BumpAllocator::<()>::new_test(heap_size);

    unsafe {
        // Try to allocate most of the heap in one go
        // Account for header size in new allocator
        let header_size = core::mem::size_of::<Header<()>>();
        let alloc_size = heap_size - header_size - 16; // more slack for 8-byte alignment

        // Request 8-byte alignment since we're checking for it
        let ptr = alloc.alloc(Layout::from_size_align(alloc_size, 8).unwrap());
        assert_ne!(ptr, core::ptr::null_mut());
        assert_eq!(0, ptr.align_offset(8));
    }
}

#[test]
fn test_old_dealloc_does_nothing() {
    // Old allocator had: "I'm a bump allocator, I don't free"
    // Verify dealloc of non-last allocations doesn't reclaim space
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();

        let ptr1 = alloc.alloc(layout);
        let ptr2 = alloc.alloc(layout);
        let ptr3 = alloc.alloc(layout);
        let used_after_3 = alloc.used();

        // Deallocating ptr1 (not last) should do nothing
        alloc.dealloc(ptr1, layout);
        assert_eq!(alloc.used(), used_after_3);

        // Deallocating ptr2 (not last) should do nothing
        alloc.dealloc(ptr2, layout);
        assert_eq!(alloc.used(), used_after_3);

        // Only deallocating ptr3 (last) should reclaim space
        alloc.dealloc(ptr3, layout);
        assert!(alloc.used() < used_after_3);
    }
}

#[test]
fn test_pointer_arithmetic_correctness() {
    // Verify to_offset and from_offset are inverses
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();
        let ptr = alloc.alloc(layout);
        assert!(!ptr.is_null());

        // Convert to offset and back
        let offset = alloc.to_offset(ptr);
        let reconstructed = alloc.from_offset(offset);

        assert_eq!(ptr, reconstructed, "Pointer reconstruction failed");
    }
}

#[test]
fn test_alignment_overflow() {
    // Test that alignment calculations don't overflow
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        // Allocate most of the space
        let layout = Layout::from_size_align(900, 8).unwrap();
        let ptr1 = alloc.alloc(layout);
        assert!(!ptr1.is_null());

        // Try to allocate with large alignment near the end
        // This should fail gracefully without overflow
        let layout2 = Layout::from_size_align(16, 64).unwrap();

        // May succeed or fail depending on remaining space, but shouldn't crash
        let _ptr2 = alloc.alloc(layout2);
    }
}

#[test]
fn test_size_overflow_detection() {
    // Test that size calculations detect overflow
    let alloc = BumpAllocator::<()>::new_test(256);

    unsafe {
        // Try to allocate something that would overflow u32
        let huge_size = (u32::MAX as usize) + 1;
        let layout = Layout::from_size_align(huge_size, 8).unwrap();
        let ptr = alloc.alloc(layout);

        // Should return null, not wrap around
        assert!(ptr.is_null(), "Should reject oversized allocation");
    }
}

#[test]
fn test_alignment_plus_size_overflow() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        // Allocate to near u32::MAX offset would be reached
        // In test environment, we can't actually reach u32::MAX
        // So just verify graceful failure with remaining space
        let layout1 = Layout::from_size_align(900, 8).unwrap();
        let ptr1 = alloc.alloc(layout1);
        assert!(!ptr1.is_null());

        // Try allocation that might cause alignment overflow
        let layout2 = Layout::from_size_align(200, 128).unwrap();

        // Should either succeed or fail gracefully, not crash
        // The exact result depends on available space after alignment
        let _ptr2 = alloc.alloc(layout2);
    }
}

#[test]
fn test_global_state_functionality() {
    #[derive(bytemuck::Zeroable)]
    struct GlobalData {
        counter: Cell<u32>,
        flag: Cell<bool>,
    }

    let alloc = BumpAllocator::<GlobalData>::new_test(1024);

    // Access and modify global state
    let global = alloc.global();
    assert_eq!(global.counter.get(), 0);
    assert_eq!(global.flag.get(), false);

    global.counter.set(42);
    global.flag.set(true);

    // Verify persistence
    let global2 = alloc.global();
    assert_eq!(global2.counter.get(), 42);
    assert_eq!(global2.flag.get(), true);

    // Verify it's the same memory location
    global.counter.set(100);
    assert_eq!(global2.counter.get(), 100);
}

#[test]
fn test_global_state_survives_allocations() {
    #[derive(bytemuck::Zeroable)]
    struct Counter {
        value: Cell<usize>,
    }

    let alloc = BumpAllocator::<Counter>::new_test(2048); // Larger heap

    // Set global state
    alloc.global().value.set(1);

    unsafe {
        // Do some allocations
        for i in 0..10 {
            let layout = Layout::from_size_align(32, 8).unwrap();
            let ptr = alloc.alloc(layout);
            assert!(!ptr.is_null(), "Allocation {} failed", i);

            // Global state should be unaffected
            let expected = i + 1;
            let actual = alloc.global().value.get();
            assert_eq!(
                actual, expected,
                "Global state corrupted at iteration {}",
                i
            );

            // Update global state
            alloc.global().value.set(i + 2);
        }

        // Final value should be preserved
        assert_eq!(alloc.global().value.get(), 11);
    }
}

#[test]
fn test_zero_alignment_rejected() {
    // Layout with zero alignment returns an error
    let result = Layout::from_size_align(32, 0);
    assert!(result.is_err(), "Zero alignment should be rejected");
}

#[test]
fn test_non_power_of_two_alignment_rejected() {
    // Layout with non-power-of-2 alignment returns an error
    let result = Layout::from_size_align(32, 3);
    assert!(
        result.is_err(),
        "Non-power-of-2 alignment should be rejected"
    );
}

#[test]
fn test_large_heap_simulation() {
    // Simulate a larger heap (1MB) to test offset calculations
    let alloc = BumpAllocator::<()>::new_test(1024 * 1024);

    unsafe {
        let layout = Layout::from_size_align(8, 8).unwrap();
        let mut last_ptr: *mut u8 = core::ptr::null_mut();

        // Allocate many small blocks
        for _ in 0..1000 {
            let ptr = alloc.alloc(layout);
            assert!(!ptr.is_null());

            if !last_ptr.is_null() {
                // Verify pointers are increasing (upward growth)
                assert!(ptr as usize > last_ptr as usize);
            }
            last_ptr = ptr;
        }

        // Verify we used a reasonable amount of space
        let used = alloc.used();
        assert!(used >= 8000); // At least 8 bytes * 1000 allocations
        assert!(used < 20000); // But not excessive due to alignment
    }
}

#[test]
fn test_realloc_overflow_protection() {
    let alloc = BumpAllocator::<()>::new_test(1024);

    unsafe {
        let layout = Layout::from_size_align(32, 8).unwrap();
        let ptr = alloc.alloc(layout);
        assert!(!ptr.is_null());

        // Try to realloc to a size that would overflow
        let huge_size = u32::MAX as usize;
        let new_ptr = alloc.realloc(ptr, layout, huge_size);

        // Should fail gracefully
        assert!(new_ptr.is_null());
    }
}

#[test]
fn test_header_alignment_verification() {
    // Verify header is properly aligned
    #[derive(bytemuck::Zeroable)]
    struct UnalignedGlobal {
        _a: u8,
        b: u64, // Requires 8-byte alignment
    }

    let alloc = BumpAllocator::<UnalignedGlobal>::new_test(1024);
    let global = alloc.global();

    // Verify the u64 field is properly aligned
    let ptr = &global.b as *const _ as usize;
    assert_eq!(ptr % 8, 0, "u64 field should be 8-byte aligned");
}

#[test]
#[ignore] // TODO(SUB-1547): Investigate how to re-enable tests after importing Agave crates
fn test_extreme_alignment_requirements() {
    let alloc = BumpAllocator::<()>::new_test(8192); // Larger heap for extreme alignments

    unsafe {
        // Test various large alignments
        for align in [64, 128, 256, 512] {
            // Reduced from 1024
            let layout = Layout::from_size_align(16, align).unwrap();
            let ptr = alloc.alloc(layout);

            if !ptr.is_null() {
                assert_eq!(
                    ptr.align_offset(align),
                    0,
                    "Failed {}-byte alignment",
                    align
                );
            }
        }
    }
}