Add HW 6

homework/06-membership/tests/src/python_memory.rs

@@ -0,0 +1,83 @@
+//! Observe actual CPython allocation failures before AnySystem replaces a
+//! Python exception with a generic panic. Student stdout is not a result channel.
+use std::ffi::c_void;
+use std::mem::MaybeUninit;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Once;
+
+use pyo3::ffi::{self, PyMemAllocatorDomain, PyMemAllocatorEx};
+use pyo3::Python;
+
+static INSTALLED: Once = Once::new();
+static FAILED: AtomicBool = AtomicBool::new(false);
+
+fn observe(pointer: *mut c_void, nonzero: bool) -> *mut c_void {
+    if nonzero && pointer.is_null() {
+        FAILED.store(true, Ordering::Relaxed);
+    }
+    pointer
+}
+
+extern "C" fn malloc(ctx: *mut c_void, size: usize) -> *mut c_void {
+    // SAFETY: ctx points to an immutable original allocator, leaked for the
+    // interpreter's entire lifetime. Calls retain CPython's allocator contract.
+    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
+    observe((allocator.malloc.unwrap())(allocator.ctx, size), size != 0)
+}
+extern "C" fn calloc(ctx: *mut c_void, count: usize, size: usize) -> *mut c_void {
+    // SAFETY: same lifetime and callback contract as malloc above.
+    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
+    observe(
+        (allocator.calloc.unwrap())(allocator.ctx, count, size),
+        count != 0 && size != 0,
+    )
+}
+extern "C" fn realloc(ctx: *mut c_void, pointer: *mut c_void, size: usize) -> *mut c_void {
+    // SAFETY: same lifetime and callback contract as malloc above.
+    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
+    observe(
+        (allocator.realloc.unwrap())(allocator.ctx, pointer, size),
+        size != 0,
+    )
+}
+extern "C" fn free(ctx: *mut c_void, pointer: *mut c_void) {
+    // SAFETY: same lifetime and callback contract as malloc above.
+    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
+    (allocator.free.unwrap())(allocator.ctx, pointer);
+}
+
+pub fn install() {
+    Python::attach(|_| {
+        INSTALLED.call_once(|| {
+            for domain in [
+                PyMemAllocatorDomain::PYMEM_DOMAIN_RAW,
+                PyMemAllocatorDomain::PYMEM_DOMAIN_MEM,
+                PyMemAllocatorDomain::PYMEM_DOMAIN_OBJ,
+            ] {
+                // SAFETY: installation holds the GIL before any student code runs.
+                // CPython requires post-initialization hooks to wrap the previous
+                // allocator. Its four callbacks are supplied by the interpreter.
+                unsafe {
+                    let mut original = MaybeUninit::uninit();
+                    ffi::PyMem_GetAllocator(domain, original.as_mut_ptr());
+                    let original = Box::into_raw(Box::new(original.assume_init()));
+                    let mut wrapper = PyMemAllocatorEx {
+                        ctx: original.cast(),
+                        malloc: Some(malloc),
+                        calloc: Some(calloc),
+                        realloc: Some(realloc),
+                        free: Some(free),
+                    };
+                    ffi::PyMem_SetAllocator(domain, &mut wrapper);
+                }
+            }
+        })
+    });
+}
+
+pub fn reset() {
+    FAILED.store(false, Ordering::Relaxed);
+}
+pub fn failed() -> bool {
+    FAILED.load(Ordering::Relaxed)
+}