use once_cell::sync::Lazy;
use parking_lot::RwLock;
use std::ptr;
use std::sync::atomic::{AtomicPtr, Ordering};
static HEAP_BASE: AtomicPtr<u8> = AtomicPtr::new(ptr::null_mut());
pub const MAX_COMPRESSED_HEAP_SIZE: usize = 4 * 1024 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct CompressedPtr {
offset: u32,
}
impl CompressedPtr {
pub fn null() -> Self {
Self { offset: 0 }
}
pub fn is_null(&self) -> bool {
self.offset == 0
}
pub fn compress(ptr: *const u8) -> Self {
if ptr.is_null() {
return Self::null();
}
let heap_base = HEAP_BASE.load(Ordering::Relaxed);
if heap_base.is_null() {
initialize_heap_base(ptr);
let heap_base = HEAP_BASE.load(Ordering::Relaxed);
if heap_base.is_null() {
log::warn!("Failed to initialize heap base for pointer compression");
return Self::null();
}
}
let base_addr = heap_base as usize;
let ptr_addr = ptr as usize;
if ptr_addr < base_addr {
return registry_store(ptr);
}
let offset = ptr_addr - base_addr;
if offset > (u32::MAX as usize - 1) {
return registry_store(ptr);
}
Self {
offset: (offset as u32) + 1,
}
}
pub fn decompress(&self) -> *const u8 {
if self.is_null() {
return ptr::null();
}
if is_registry_tag(self.offset) {
let idx = registry_index(self.offset);
return registry_get(idx);
}
let heap_base = HEAP_BASE.load(Ordering::Relaxed);
if heap_base.is_null() {
log::error!("Heap base not initialized for decompression");
return ptr::null();
}
let base_addr = heap_base as usize;
let actual_offset = (self.offset as usize).saturating_sub(1);
let ptr_addr = base_addr + actual_offset;
ptr_addr as *const u8
}
pub fn offset(&self) -> u32 {
if self.is_null() {
0
} else {
self.offset - 1
}
}
pub unsafe fn from_offset(offset: u32) -> Self {
Self { offset }
}
pub fn add_offset(&self, additional_offset: usize) -> Option<Self> {
let base = if self.is_null() {
0usize
} else {
(self.offset as usize) - 1
};
let new_actual = base.checked_add(additional_offset)?;
if new_actual <= (u32::MAX as usize - 1) {
Some(Self {
offset: (new_actual as u32) + 1,
})
} else {
None
}
}
pub fn offset_from(&self, other: &CompressedPtr) -> Option<isize> {
if is_registry_tag(self.offset) || is_registry_tag(other.offset) {
return None;
}
let a = if self.is_null() { 0 } else { self.offset - 1 } as isize;
let b = if other.is_null() { 0 } else { other.offset - 1 } as isize;
if a >= b {
Some(a - b)
} else {
Some(-(b - a))
}
}
}
const REGISTRY_TAG: u32 = 1 << 31;
fn is_registry_tag(v: u32) -> bool {
(v & REGISTRY_TAG) != 0
}
fn registry_index(v: u32) -> usize {
((v & !REGISTRY_TAG) as usize).saturating_sub(1)
}
static POINTER_REGISTRY: Lazy<RwLock<Vec<usize>>> = Lazy::new(|| RwLock::new(Vec::new()));
fn registry_store(ptr: *const u8) -> CompressedPtr {
let mut reg = POINTER_REGISTRY.write();
let idx = reg.len();
if idx >= (u32::MAX as usize - 1) {
log::error!("Pointer registry exhausted; returning null compressed ptr");
return CompressedPtr::null();
}
reg.push(ptr as usize);
let stored = REGISTRY_TAG | ((idx as u32) + 1);
CompressedPtr { offset: stored }
}
fn registry_get(idx: usize) -> *const u8 {
let reg = POINTER_REGISTRY.read();
if idx < reg.len() {
reg[idx] as *const u8
} else {
ptr::null()
}
}
impl Default for CompressedPtr {
fn default() -> Self {
Self::null()
}
}
fn initialize_heap_base(first_ptr: *const u8) {
let page_size = page_size::get();
let aligned_base = (first_ptr as usize) & !(page_size - 1);
let base_ptr = aligned_base as *mut u8;
match HEAP_BASE.compare_exchange(
ptr::null_mut(),
base_ptr,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => {
log::info!("Initialized pointer compression heap base at {base_ptr:p}");
}
Err(existing) => {
log::debug!(
"Heap base already initialized at {existing:p}, first ptr was {first_ptr:p}"
);
}
}
}
pub fn get_heap_base() -> *const u8 {
HEAP_BASE.load(Ordering::Relaxed) as *const u8
}
#[cfg(test)]
pub fn set_heap_base_for_test(base: *mut u8) {
HEAP_BASE.store(base, Ordering::Relaxed);
}
#[cfg(test)]
pub fn reset_heap_base_for_test() {
HEAP_BASE.store(ptr::null_mut(), Ordering::Relaxed);
}
pub fn is_compression_beneficial() -> bool {
std::mem::size_of::<*const u8>() == 8
}
pub fn calculate_compression_savings(num_pointers: usize) -> usize {
if is_compression_beneficial() {
num_pointers * 4 } else {
0
}
}
#[derive(Debug, Clone)]
pub struct CompressionStats {
pub heap_base: *const u8,
pub max_offset: u32,
pub compression_ratio: f64,
pub memory_saved_bytes: usize,
pub pointers_compressed: usize,
}
impl CompressionStats {
pub fn new() -> Self {
Self {
heap_base: get_heap_base(),
max_offset: 0,
compression_ratio: 0.5, memory_saved_bytes: 0,
pointers_compressed: 0,
}
}
pub fn record_compression(&mut self, _ptr: *const u8) {
self.pointers_compressed += 1;
self.memory_saved_bytes = calculate_compression_savings(self.pointers_compressed);
}
pub fn is_effective(&self) -> bool {
!self.heap_base.is_null() && self.memory_saved_bytes > 0
}
pub fn efficiency_percentage(&self) -> f64 {
if self.pointers_compressed == 0 {
0.0
} else {
self.compression_ratio * 100.0
}
}
}
impl Default for CompressionStats {
fn default() -> Self {
Self::new()
}
}
pub struct CompressedPtrArray {
pointers: Vec<CompressedPtr>,
stats: CompressionStats,
}
impl CompressedPtrArray {
pub fn new() -> Self {
Self {
pointers: Vec::new(),
stats: CompressionStats::new(),
}
}
pub fn with_capacity(capacity: usize) -> Self {
Self {
pointers: Vec::with_capacity(capacity),
stats: CompressionStats::new(),
}
}
pub fn push(&mut self, ptr: *const u8) {
let compressed = CompressedPtr::compress(ptr);
self.pointers.push(compressed);
self.stats.record_compression(ptr);
}
pub fn get(&self, index: usize) -> Option<*const u8> {
self.pointers.get(index).map(|cp| cp.decompress())
}
pub fn get_compressed(&self, index: usize) -> Option<CompressedPtr> {
self.pointers.get(index).copied()
}
pub fn len(&self) -> usize {
self.pointers.len()
}
pub fn is_empty(&self) -> bool {
self.pointers.is_empty()
}
pub fn clear(&mut self) {
self.pointers.clear();
self.stats = CompressionStats::new();
}
pub fn stats(&self) -> &CompressionStats {
&self.stats
}
pub fn memory_usage(&self) -> (usize, usize) {
let compressed_size = self.pointers.len() * std::mem::size_of::<CompressedPtr>();
let uncompressed_size = self.pointers.len() * std::mem::size_of::<*const u8>();
(compressed_size, uncompressed_size)
}
}
impl Default for CompressedPtrArray {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Mutex, OnceLock};
fn compression_test_guard() -> std::sync::MutexGuard<'static, ()> {
static LOCK: OnceLock<Mutex<()>> = OnceLock::new();
LOCK.get_or_init(|| Mutex::new(()))
.lock()
.unwrap_or_else(|poison| poison.into_inner())
}
#[test]
fn test_compressed_ptr_null() {
let _guard = compression_test_guard();
let null_ptr = CompressedPtr::null();
assert!(null_ptr.is_null());
assert_eq!(null_ptr.decompress(), ptr::null());
}
#[test]
fn test_compressed_ptr_basic() {
let _guard = compression_test_guard();
reset_heap_base_for_test();
let test_data = [1, 2, 3, 4];
let ptr = test_data.as_ptr();
let compressed = CompressedPtr::compress(ptr);
assert!(!compressed.is_null());
let decompressed = compressed.decompress();
assert_eq!(ptr, decompressed);
}
#[test]
fn test_compressed_ptr_offset() {
let _guard = compression_test_guard();
reset_heap_base_for_test();
let test_data = [1u8; 1000];
let base_ptr = test_data.as_ptr();
let offset_ptr = unsafe { base_ptr.add(100) };
set_heap_base_for_test(base_ptr as *mut u8);
let compressed = CompressedPtr::compress(offset_ptr);
assert!(!compressed.is_null());
assert_eq!(compressed.offset(), 100);
let decompressed = compressed.decompress();
assert_eq!(offset_ptr, decompressed);
}
#[test]
fn test_compressed_ptr_array() {
let _guard = compression_test_guard();
reset_heap_base_for_test();
let mut array = CompressedPtrArray::new();
assert!(array.is_empty());
let test_data = [1u8, 2, 3, 4, 5];
for (i, item) in test_data.iter().enumerate() {
array.push(item as *const u8);
assert_eq!(array.len(), i + 1);
}
for (i, expected) in test_data.iter().enumerate() {
let retrieved = array.get(i).expect("should have pointer");
assert_eq!(retrieved, expected as *const u8);
}
let stats = array.stats();
assert_eq!(stats.pointers_compressed, 5);
let (compressed_size, uncompressed_size) = array.memory_usage();
assert!(compressed_size <= uncompressed_size);
}
#[test]
fn test_compression_savings() {
let _guard = compression_test_guard();
let savings = calculate_compression_savings(1000);
if is_compression_beneficial() {
assert_eq!(savings, 4000); } else {
assert_eq!(savings, 0);
}
}
#[test]
fn test_compressed_ptr_add_offset() {
let _guard = compression_test_guard();
let ptr = CompressedPtr { offset: 100 };
let new_ptr = ptr.add_offset(50).expect("should add offset");
assert_eq!(new_ptr.offset(), 149);
let large_ptr = CompressedPtr {
offset: u32::MAX - 10,
};
assert!(large_ptr.add_offset(20).is_none());
}
#[test]
fn test_compressed_ptr_offset_from() {
let _guard = compression_test_guard();
let ptr1 = CompressedPtr { offset: 100 };
let ptr2 = CompressedPtr { offset: 150 };
let diff = ptr2.offset_from(&ptr1).expect("should calculate offset");
assert_eq!(diff, 50);
let diff_rev = ptr1.offset_from(&ptr2).expect("should calculate offset");
assert_eq!(diff_rev, -50);
}
#[test]
fn test_compression_stats() {
let _guard = compression_test_guard();
let mut stats = CompressionStats::new();
assert_eq!(stats.pointers_compressed, 0);
assert_eq!(stats.memory_saved_bytes, 0);
let dummy_ptr = 0x1000 as *const u8;
stats.record_compression(dummy_ptr);
assert_eq!(stats.pointers_compressed, 1);
if is_compression_beneficial() {
assert_eq!(stats.memory_saved_bytes, 4);
}
}
}