Allow using the rust global allocator for memory allocations

Author: bjorn3

Cargo.toml

@@ -40,8 +40,11 @@ features = ["xxh64", "xxh32"]
 
 
 [features]
+default = ["rust-allocator"]
 export-symbols = [] # whether the zlib api symbols are publicly exported
 semver-prefix = ["export-symbols"] # prefix all symbols in a semver-compatible way
 no-prefetch = [] # do not prefetch
 trace = []
 nightly = []
+c-allocator = []
+rust-allocator = []

lib/common/allocations.rs

@@ -1,12 +1,10 @@
 use std::ptr;
 
-use libc::{calloc, free, malloc, size_t};
-
 use crate::lib::zstd::ZSTD_customMem;
 
 #[inline]
 pub(crate) unsafe fn ZSTD_customMalloc(
-    size: size_t,
+    size: usize,
     customMem: ZSTD_customMem,
 ) -> *mut core::ffi::c_void {
     if customMem.customAlloc.is_some() ^ customMem.customFree.is_some() {
@@ -17,11 +15,15 @@ pub(crate) unsafe fn ZSTD_customMalloc(
         return f(customMem.opaque, size);
     }
 
-    malloc(size)
+    #[cfg(feature = "rust-allocator")]
+    return std::alloc::alloc(core::alloc::Layout::from_size_align_unchecked(size, 16)).cast();
+
+    #[cfg(feature = "c-allocator")]
+    return libc::malloc(size);
 }
 #[inline]
 pub(crate) unsafe fn ZSTD_customCalloc(
-    size: size_t,
+    size: usize,
     customMem: ZSTD_customMem,
 ) -> *mut core::ffi::c_void {
     if customMem.customAlloc.is_some() ^ customMem.customFree.is_some() {
@@ -33,15 +35,32 @@ pub(crate) unsafe fn ZSTD_customCalloc(
         ptr::write_bytes(ptr, 0, size);
         return ptr;
     }
-    calloc(1, size)
+
+    #[cfg(feature = "rust-allocator")]
+    return std::alloc::alloc_zeroed(core::alloc::Layout::from_size_align_unchecked(size, 16))
+        .cast();
+
+    #[cfg(feature = "c-allocator")]
+    return libc::calloc(1, size);
 }
 #[inline]
-pub(crate) unsafe fn ZSTD_customFree(ptr: *mut core::ffi::c_void, customMem: ZSTD_customMem) {
+pub(crate) unsafe fn ZSTD_customFree(
+    ptr: *mut core::ffi::c_void,
+    _size: usize,
+    customMem: ZSTD_customMem,
+) {
     if !ptr.is_null() {
         if let Some(f) = customMem.customFree {
             f(customMem.opaque, ptr);
         } else {
-            free(ptr);
+            #[cfg(feature = "rust-allocator")]
+            return std::alloc::dealloc(
+                ptr.cast(),
+                core::alloc::Layout::from_size_align_unchecked(_size, 16),
+            );
+
+            #[cfg(feature = "c-allocator")]
+            return libc::free(ptr);
         }
     }
 }

lib/common/pool.rs

@@ -136,9 +136,23 @@ pub unsafe fn POOL_free(ctx: *mut POOL_ctx) {
     ptr::drop_in_place(ptr::addr_of_mut!((*ctx).queueMutex));
     ptr::drop_in_place(ptr::addr_of_mut!((*ctx).queuePushCond));
     ptr::drop_in_place(ptr::addr_of_mut!((*ctx).queuePopCond));
-    ZSTD_customFree((*ctx).queue as *mut core::ffi::c_void, (*ctx).customMem);
-    ZSTD_customFree((*ctx).threads as *mut core::ffi::c_void, (*ctx).customMem);
-    ZSTD_customFree(ctx as *mut core::ffi::c_void, (*ctx).customMem);
+    ZSTD_customFree(
+        (*ctx).queue as *mut core::ffi::c_void,
+        ((*ctx).queueSize).wrapping_mul(::core::mem::size_of::<POOL_job>()),
+        (*ctx).customMem,
+    );
+    ZSTD_customFree(
+        (*ctx).threads as *mut core::ffi::c_void,
+        (*ctx)
+            .threadCapacity
+            .wrapping_mul(::core::mem::size_of::<JoinHandle<()>>()),
+        (*ctx).customMem,
+    );
+    ZSTD_customFree(
+        ctx as *mut core::ffi::c_void,
+        ::core::mem::size_of::<POOL_ctx>(),
+        (*ctx).customMem,
+    );
 }
 pub unsafe fn POOL_joinJobs(ctx: *mut POOL_ctx) {
     let mut guard = (*ctx).queueMutex.lock().unwrap();
@@ -181,7 +195,13 @@ unsafe fn POOL_resize_internal(ctx: *mut POOL_ctx, numThreads: size_t) -> core::
         (*ctx).threads as *const core::ffi::c_void,
         ((*ctx).threadCapacity).wrapping_mul(::core::mem::size_of::<JoinHandle<()>>()),
     );
-    ZSTD_customFree((*ctx).threads as *mut core::ffi::c_void, (*ctx).customMem);
+    ZSTD_customFree(
+        (*ctx).threads as *mut core::ffi::c_void,
+        (*ctx)
+            .threadCapacity
+            .wrapping_mul(::core::mem::size_of::<JoinHandle<()>>()),
+        (*ctx).customMem,
+    );
     (*ctx).threads = threadPool;
     for threadId in (*ctx).threadCapacity..numThreads {
         let ctx = SendPoolCtxPtr(ctx);

lib/compress/zstd_compress.rs

@@ -1251,8 +1251,11 @@ unsafe fn ZSTD_cwksp_create(
 #[inline]
 unsafe fn ZSTD_cwksp_free(ws: *mut ZSTD_cwksp, customMem: ZSTD_customMem) {
     let ptr = (*ws).workspace;
+    let size = (*ws)
+        .workspaceEnd
+        .byte_offset_from_unsigned((*ws).workspace);
     ptr::write_bytes(ws as *mut u8, 0, ::core::mem::size_of::<ZSTD_cwksp>());
-    ZSTD_customFree(ptr, customMem);
+    ZSTD_customFree(ptr, size, customMem);
 }
 #[inline]
 unsafe fn ZSTD_cwksp_move(dst: *mut ZSTD_cwksp, src: *mut ZSTD_cwksp) {
@@ -1705,7 +1708,11 @@ pub unsafe extern "C" fn ZSTD_initStaticCCtx(
     cctx
 }
 unsafe fn ZSTD_clearAllDicts(cctx: *mut ZSTD_CCtx) {
-    ZSTD_customFree((*cctx).localDict.dictBuffer, (*cctx).customMem);
+    ZSTD_customFree(
+        (*cctx).localDict.dictBuffer,
+        (*cctx).localDict.dictSize,
+        (*cctx).customMem,
+    );
     ZSTD_freeCDict((*cctx).localDict.cdict);
     ptr::write_bytes(
         &mut (*cctx).localDict as *mut ZSTD_localDict as *mut u8,
@@ -1746,7 +1753,11 @@ pub unsafe extern "C" fn ZSTD_freeCCtx(cctx: *mut ZSTD_CCtx) -> size_t {
         ZSTD_cwksp_owns_buffer(&(*cctx).workspace, cctx as *const core::ffi::c_void);
     ZSTD_freeCCtxContent(cctx);
     if cctxInWorkspace == 0 {
-        ZSTD_customFree(cctx as *mut core::ffi::c_void, (*cctx).customMem);
+        ZSTD_customFree(
+            cctx as *mut core::ffi::c_void,
+            ::core::mem::size_of::<ZSTD_CCtx>(),
+            (*cctx).customMem,
+        );
     }
     0
 }
@@ -1965,7 +1976,11 @@ pub unsafe extern "C" fn ZSTD_freeCCtxParams(params: *mut ZSTD_CCtx_params) -> s
     if params.is_null() {
         return 0;
     }
-    ZSTD_customFree(params as *mut core::ffi::c_void, (*params).customMem);
+    ZSTD_customFree(
+        params as *mut core::ffi::c_void,
+        ::core::mem::size_of::<ZSTD_CCtx_params>(),
+        (*params).customMem,
+    );
     0
 }
 #[cfg_attr(feature = "export-symbols", export_name = crate::prefix!(ZSTD_CCtxParams_reset))]
@@ -5564,7 +5579,7 @@ pub unsafe extern "C" fn ZSTD_generateSequences(
     seqCollector.maxSequences = outSeqsSize;
     (*zc).seqCollector = seqCollector;
     let ret = ZSTD_compress2(zc, dst, dstCapacity, src, srcSize);
-    ZSTD_customFree(dst, ZSTD_customMem::default());
+    ZSTD_customFree(dst, dstCapacity, ZSTD_customMem::default());
     let err_code_1 = ret;
     if ERR_isError(err_code_1) {
         return err_code_1;
@@ -8657,7 +8672,6 @@ unsafe fn ZSTD_createCDict_advanced_internal(
         isStatic: ZSTD_cwksp_dynamic_alloc,
     };
     if workspace.is_null() {
-        ZSTD_customFree(workspace, customMem);
         return core::ptr::null_mut();
     }
     ZSTD_cwksp_init(&mut ws, workspace, workspaceSize, ZSTD_cwksp_dynamic_alloc);
@@ -8765,7 +8779,6 @@ pub unsafe extern "C" fn ZSTD_createCDict_advanced2(
         targetLength: 0,
         strategy: 0,
     };
-    let mut cdict = core::ptr::null_mut::<ZSTD_CDict>();
     if cctxParams.enableDedicatedDictSearch != 0 {
         cParams = ZSTD_dedicatedDictSearch_getCParams(cctxParams.compressionLevel, dictSize);
         ZSTD_overrideCParams(&mut cParams, &cctxParams.cParams);
@@ -8789,7 +8802,7 @@ pub unsafe extern "C" fn ZSTD_createCDict_advanced2(
     cctxParams.cParams = cParams;
     cctxParams.useRowMatchFinder =
         ZSTD_resolveRowMatchFinderMode(cctxParams.useRowMatchFinder, &cParams);
-    cdict = ZSTD_createCDict_advanced_internal(
+    let cdict = ZSTD_createCDict_advanced_internal(
         dictSize,
         dictLoadMethod,
         cctxParams.cParams,
@@ -8880,7 +8893,11 @@ pub unsafe extern "C" fn ZSTD_freeCDict(cdict: *mut ZSTD_CDict) -> size_t {
         ZSTD_cwksp_owns_buffer(&(*cdict).workspace, cdict as *const core::ffi::c_void);
     ZSTD_cwksp_free(&mut (*cdict).workspace, cMem);
     if cdictInWorkspace == 0 {
-        ZSTD_customFree(cdict as *mut core::ffi::c_void, cMem);
+        ZSTD_customFree(
+            cdict as *mut core::ffi::c_void,
+            (*cdict).dictContentSize,
+            cMem,
+        );
     }
     0
 }
@@ -8898,18 +8915,14 @@ pub unsafe extern "C" fn ZSTD_initStaticCDict(
         ZSTD_resolveRowMatchFinderMode(ZSTD_ParamSwitch_e::ZSTD_ps_auto, &cParams);
     let matchStateSize = ZSTD_sizeof_matchState(&cParams, useRowMatchFinder, 1, 0);
     let neededSize = (ZSTD_cwksp_alloc_size(::core::mem::size_of::<ZSTD_CDict>()))
-        .wrapping_add(
-            if dictLoadMethod as core::ffi::c_uint
-                == ZSTD_dlm_byRef as core::ffi::c_int as core::ffi::c_uint
-            {
-                0
-            } else {
-                ZSTD_cwksp_alloc_size(ZSTD_cwksp_align(
-                    dictSize,
-                    ::core::mem::size_of::<*mut core::ffi::c_void>(),
-                ))
-            },
-        )
+        .wrapping_add(if dictLoadMethod == ZSTD_dlm_byRef {
+            0
+        } else {
+            ZSTD_cwksp_alloc_size(ZSTD_cwksp_align(
+                dictSize,
+                ::core::mem::size_of::<*mut core::ffi::c_void>(),
+            ))
+        })
         .wrapping_add(ZSTD_cwksp_alloc_size(HUF_WORKSPACE_SIZE as size_t))
         .wrapping_add(matchStateSize);
     let mut cdict = core::ptr::null_mut::<ZSTD_CDict>();

lib/compress/zstdmt_compress.rs

@@ -294,17 +294,23 @@ unsafe fn ZSTDMT_freeBufferPool(bufPool: *mut ZSTDMT_bufferPool) {
         while u < (*bufPool).totalBuffers {
             ZSTD_customFree(
                 (*((*bufPool).buffers).offset(u as isize)).start,
+                (*((*bufPool).buffers).offset(u as isize)).capacity,
                 (*bufPool).cMem,
             );
             u = u.wrapping_add(1);
         }
         ZSTD_customFree(
             (*bufPool).buffers as *mut core::ffi::c_void,
+            (*bufPool).bufferSize,
             (*bufPool).cMem,
         );
     }
     core::ptr::drop_in_place(core::ptr::addr_of_mut!((*bufPool).poolMutex));
-    ZSTD_customFree(bufPool as *mut core::ffi::c_void, (*bufPool).cMem);
+    ZSTD_customFree(
+        bufPool as *mut core::ffi::c_void,
+        size_of::<ZSTDMT_bufferPool>(),
+        (*bufPool).cMem,
+    );
 }
 unsafe fn ZSTDMT_createBufferPool(
     maxNbBuffers: core::ffi::c_uint,
@@ -389,7 +395,7 @@ unsafe fn ZSTDMT_getBuffer(bufPool: *mut ZSTDMT_bufferPool) -> Buffer {
         {
             return buf;
         }
-        ZSTD_customFree(buf.start, (*bufPool).cMem);
+        ZSTD_customFree(buf.start, buf.capacity, (*bufPool).cMem);
     }
     drop(guard);
     let mut buffer = buffer_s {
@@ -414,7 +420,7 @@ unsafe fn ZSTDMT_releaseBuffer(bufPool: *mut ZSTDMT_bufferPool, buf: Buffer) {
         return;
     }
     drop(guard);
-    ZSTD_customFree(buf.start, (*bufPool).cMem);
+    ZSTD_customFree(buf.start, buf.capacity, (*bufPool).cMem);
 }
 unsafe fn ZSTDMT_sizeof_seqPool(seqPool: *mut ZSTDMT_seqPool) -> size_t {
     ZSTDMT_sizeof_bufferPool(seqPool)
@@ -478,9 +484,17 @@ unsafe fn ZSTDMT_freeCCtxPool(pool: *mut ZSTDMT_CCtxPool) {
             ZSTD_freeCCtx(*((*pool).cctxs).offset(cid as isize));
             cid += 1;
         }
-        ZSTD_customFree((*pool).cctxs as *mut core::ffi::c_void, (*pool).cMem);
+        ZSTD_customFree(
+            (*pool).cctxs as *mut core::ffi::c_void,
+            ((*pool).totalCCtx as usize).wrapping_mul(::core::mem::size_of::<*mut ZSTD_CCtx>()),
+            (*pool).cMem,
+        );
     }
-    ZSTD_customFree(pool as *mut core::ffi::c_void, (*pool).cMem);
+    ZSTD_customFree(
+        pool as *mut core::ffi::c_void,
+        size_of::<ZSTDMT_CCtxPool>(),
+        (*pool).cMem,
+    );
 }
 unsafe fn ZSTDMT_createCCtxPool(
     nbWorkers: core::ffi::c_int,
@@ -605,6 +619,7 @@ unsafe fn ZSTDMT_serialState_reset(
         {
             ZSTD_customFree(
                 (*serialState).ldmState.hashTable as *mut core::ffi::c_void,
+                hashSize,
                 cMem,
             );
             (*serialState).ldmState.hashTable =
@@ -613,6 +628,7 @@ unsafe fn ZSTDMT_serialState_reset(
         if ((*serialState).ldmState.bucketOffsets).is_null() || prevBucketLog < bucketLog {
             ZSTD_customFree(
                 (*serialState).ldmState.bucketOffsets as *mut core::ffi::c_void,
+                1 << prevBucketLog,
                 cMem,
             );
             (*serialState).ldmState.bucketOffsets = ZSTD_customMalloc(numBuckets, cMem) as *mut u8;
@@ -676,12 +692,19 @@ unsafe fn ZSTDMT_serialState_free(serialState: *mut SerialState) {
     core::ptr::drop_in_place(core::ptr::addr_of_mut!((*serialState).cond));
     core::ptr::drop_in_place(core::ptr::addr_of_mut!((*serialState).ldmWindowMutex));
     core::ptr::drop_in_place(core::ptr::addr_of_mut!((*serialState).ldmWindowCond));
+    let hashLog = (*serialState).params.ldmParams.hashLog;
+    let hashSize = ((1 as size_t) << hashLog).wrapping_mul(::core::mem::size_of::<ldmEntry_t>());
+    let bucketLog = ((*serialState).params.ldmParams.hashLog)
+        .wrapping_sub((*serialState).params.ldmParams.bucketSizeLog);
+    let numBuckets = 1usize << bucketLog;
     ZSTD_customFree(
         (*serialState).ldmState.hashTable as *mut core::ffi::c_void,
+        hashSize,
         cMem,
     );
     ZSTD_customFree(
         (*serialState).ldmState.bucketOffsets as *mut core::ffi::c_void,
+        numBuckets,
         cMem,
     );
 }
@@ -1066,7 +1089,11 @@ unsafe fn ZSTDMT_freeJobsTable(
         ));
         jobNb = jobNb.wrapping_add(1);
     }
-    ZSTD_customFree(jobTable as *mut core::ffi::c_void, cMem);
+    ZSTD_customFree(
+        jobTable as *mut core::ffi::c_void,
+        (nbJobs as usize).wrapping_mul(::core::mem::size_of::<ZSTDMT_jobDescription>()),
+        cMem,
+    );
 }
 unsafe fn ZSTDMT_createJobsTable(
     nbJobsPtr: *mut u32,
@@ -1268,10 +1295,15 @@ pub unsafe fn ZSTDMT_freeCCtx(mtctx: *mut ZSTDMT_CCtx) -> size_t {
     if !((*mtctx).roundBuff.buffer).is_null() {
         ZSTD_customFree(
             (*mtctx).roundBuff.buffer as *mut core::ffi::c_void,
+            (*mtctx).roundBuff.capacity,
             (*mtctx).cMem,
         );
     }
-    ZSTD_customFree(mtctx as *mut core::ffi::c_void, (*mtctx).cMem);
+    ZSTD_customFree(
+        mtctx as *mut core::ffi::c_void,
+        ::core::mem::size_of::<ZSTDMT_CCtx>(),
+        (*mtctx).cMem,
+    );
     0
 }
 pub unsafe fn ZSTDMT_sizeof_CCtx(mtctx: *mut ZSTDMT_CCtx) -> size_t {
@@ -1560,6 +1592,7 @@ pub unsafe fn ZSTDMT_initCStream_internal(
         if !((*mtctx).roundBuff.buffer).is_null() {
             ZSTD_customFree(
                 (*mtctx).roundBuff.buffer as *mut core::ffi::c_void,
+                (*mtctx).roundBuff.capacity,
                 (*mtctx).cMem,
             );
         }

lib/decompress/zstd_ddict.rs

@@ -297,8 +297,12 @@ pub unsafe extern "C" fn ZSTD_freeDDict(ddict: *mut ZSTD_DDict) -> size_t {
         return 0;
     }
     let cMem = (*ddict).cMem;
-    ZSTD_customFree((*ddict).dictBuffer, cMem);
-    ZSTD_customFree(ddict as *mut core::ffi::c_void, cMem);
+    ZSTD_customFree((*ddict).dictBuffer, (*ddict).dictSize, cMem);
+    ZSTD_customFree(
+        ddict as *mut core::ffi::c_void,
+        size_of::<ZSTD_DDict>(),
+        cMem,
+    );
     0
 }