alloc.rs

Overview

The `alloc.rs` file implements a **custom global memory allocator** for the Rust components of the project, which delegates all memory management operations to Python's native memory allocator APIs (`PyMem_Malloc`, `PyMem_Free`, and `PyMem_Realloc`). This integration ensures that memory allocations performed by Rust code remain compatible with Python's memory management system, preserving consistency, safety, and debugging capabilities when Rust interacts with Python objects.

By replacing Rust's default global allocator with this Python-aware allocator, the project avoids memory fragmentation and ownership conflicts between Rust and Python runtimes, which is critical for seamless interoperability in Python extension modules or embedded Rust code.

Detailed Description

PyMemAllocator Struct

• Purpose:
  Represents the custom allocator that forwards all allocation and deallocation requests to Python's memory management functions.

• Definition:

  struct PyMemAllocator {}
  

  This struct is a zero-sized type acting purely as a marker for implementing allocator traits.

Global Allocator Instance

• Static Instance:

  #[global_allocator]
  static ALLOCATOR: PyMemAllocator = PyMemAllocator {};
  

• Purpose:
  Marks ALLOCATOR as the global allocator for the Rust crate. All memory allocation/deallocation done by Rust code will route through this allocator.

Trait Implementations

Sync for PyMemAllocator

• Why needed:
  The global allocator must be thread-safe (implement Sync) because memory allocations may occur concurrently.

• Unsafe Implementation:

  unsafe impl Sync for PyMemAllocator {}
  

• Note:
  The implementation is marked unsafe because the compiler cannot verify thread safety automatically, but the underlying Python allocator is assumed thread-safe for these operations.

GlobalAlloc for PyMemAllocator

This trait defines the core allocator interface for the Rust global allocator.

unsafe impl GlobalAlloc for PyMemAllocator {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8;

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout);

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8;

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8;
}

alloc

• Signature:

  unsafe fn alloc(&self, layout: Layout) -> *mut u8
  

• Description:
  Allocates a memory block of layout.size() bytes by calling pyo3_ffi::PyMem_Malloc.

• Parameters:

  • layout: Describes size and alignment of the requested memory block.

• Returns:
  Pointer to the allocated memory block (*mut u8), or null on failure.

• Example Usage:

  let ptr = ALLOCATOR.alloc(Layout::from_size_align(32, 8).unwrap());
  

dealloc

• Signature:

  unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout)
  

• Description:
  Frees a previously allocated memory block pointed by ptr using pyo3_ffi::PyMem_Free.

• Parameters:

  • ptr: Pointer to memory to free.

  • _layout: The layout of the memory block (unused in this implementation).

• Example Usage:

  ALLOCATOR.dealloc(ptr, layout);
  

alloc_zeroed

• Signature:

  unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8
  

• Description:
  Allocates zero-initialized memory by first allocating the requested size via PyMem_Malloc and then explicitly zeroing the memory using core::ptr::write_bytes.

• Parameters:

  • layout: Memory size and alignment.

• Returns:
  Pointer to zero-initialized memory.

• Implementation Detail:
  Python's allocator does not provide a direct zeroing allocation, so zeroing is done manually.

• Example Usage:

  let zeroed_ptr = ALLOCATOR.alloc_zeroed(Layout::from_size_align(64, 8).unwrap());
  

realloc

• Signature:

  unsafe fn realloc(&self, ptr: *mut u8, _layout: Layout, new_size: usize) -> *mut u8
  

• Description:
  Changes the size of the memory block pointed by ptr to new_size bytes by calling pyo3_ffi::PyMem_Realloc.

• Parameters:

  • ptr: Pointer to the existing allocation.

  • _layout: Original layout (unused).

  • new_size: The new requested size.

• Returns:
  Pointer to the reallocated memory block (which may be different from ptr).

• Example Usage:

  let new_ptr = ALLOCATOR.realloc(ptr, old_layout, 128);
  

Important Implementation Details

• Use of Unsafe Code:
  All allocator methods are marked unsafe because:

  • They perform raw pointer manipulations.

  • They call into C API functions (PyMem_Malloc, PyMem_Free, PyMem_Realloc) which are inherently unsafe.

  The allocator assumes that these Python APIs are correctly implemented and thread-safe.

• Alignment Handling:
  The allocator ignores alignment requirements beyond size because Python's allocator does not expose alignment guarantees. This is acceptable because Python's allocator typically provides sufficiently aligned memory for general use.

• Zeroing Memory:
  Since Python's allocator does not provide zeroed allocation natively, zeroing is done explicitly with a call to core::ptr::write_bytes.

• Global Registration:
  By marking the static instance ALLOCATOR with #[global_allocator], the Rust compiler uses this allocator automatically for all heap allocations in this crate, without requiring manual invocation.

Interaction with Other System Components

• Integration with Python Runtime:
  Because Rust code in this project interacts heavily with Python objects (e.g., serialization/deserialization), using Python's memory allocator ensures that the memory lifecycle is consistent with Python's expectations.

• Compatibility with Deserialization and Serialization:
  The allocator supports other modules like deserialization key caches or JSON parsing by providing safe memory management that aligns with Python's memory tracking.

• FFI Boundary Safety:
  Avoids potential memory corruption or leaks that could arise if Rust and Python used different allocators for the same data, especially when passing pointers across FFI boundaries.

Usage Example

use core::alloc::{GlobalAlloc, Layout};

fn example_allocation() {
    let layout = Layout::from_size_align(256, 8).unwrap();
    unsafe {
        // Allocate 256 bytes
        let ptr = ALLOCATOR.alloc(layout);

        if !ptr.is_null() {
            // Use the memory...

            // Deallocate when done
            ALLOCATOR.dealloc(ptr, layout);
        }
    }
}

Mermaid Diagram: Memory Allocation Flow

flowchart TD
    RustCode[All Rust Code]
    MemoryOps[Memory Operations\n(alloc, dealloc, realloc, alloc_zeroed)]
    PyMemAPI[Python Memory API\n(PyMem_Malloc, PyMem_Free, PyMem_Realloc)]

    RustCode -->|requests memory| MemoryOps
    MemoryOps -->|calls| PyMemAPI

Summary

• alloc.rs replaces Rust’s default heap allocator with a Python-aware allocator.

• Implements GlobalAlloc trait forwarding memory operations to Python's allocator APIs.

• Ensures memory compatibility and safety between Rust and Python runtimes.

• Uses unsafe Rust code to interface with Python C APIs.

• Registered globally to transparently handle all heap allocations in the crate.

• Critical for stable, efficient, and safe integration of Rust code within Python extensions or embedded interpreters.

This design is fundamental for the project’s Python interoperability, preventing memory management issues when Rust and Python components share data and resources.