windows-nt/Source/XPSP1/NT/sdktools/mtscript/util/memutils.cxx
2020-09-26 16:20:57 +08:00

688 lines
15 KiB
C++

//+------------------------------------------------------------------------
//
// File: memutil.cxx
//
// Contents: Memory utilities
//
// History: Stolen from Trident
//
//-------------------------------------------------------------------------
#include "headers.hxx"
EXTERN_C HANDLE g_hProcessHeap;
#define SMALLBLOCKHEAP 0
void
ClearInterfaceFn(IUnknown **ppUnk)
{
IUnknown * pUnk;
pUnk = *ppUnk;
*ppUnk = NULL;
if (pUnk)
pUnk->Release();
}
//+------------------------------------------------------------------------
// Allocation functions not implemented in this file:
//
// CDUTIL.HXX
// operator new
// operator delete
//
// OLE's OBJBASE.H
// CoTaskMemAlloc, CoTaskMemFree
//
//-------------------------------------------------------------------------
#if SMALLBLOCKHEAP
DeclareTag(tagSmallBlockHeap, "!Memory", "Check small block heap every time")
DeclareTag(tagSmallBlockHeapDisable, "!Memory", "Disable small block heap");
#define _CRTBLD 1
#include "winheap.h"
EXTERN_C CRITICAL_SECTION g_csHeap;
#if DBG == 1
#define CHECKSBH if (IsTagEnabled(tagSmallBlockHeap)) {Assert(CheckSmallBlockHeap() && "Small block heap corrupt");};
BOOL IsSmallBlockHeapEnabled()
{
static int g_fSmallBlockHeap = -1;
if (g_fSmallBlockHeap == -1)
g_fSmallBlockHeap = IsTagEnabled(tagSmallBlockHeapDisable) ? 0 : 1;
return(g_fSmallBlockHeap == 1);
}
BOOL CheckSmallBlockHeap()
{
if (IsSmallBlockHeapEnabled())
{
EnterCriticalSection(&g_csHeap);
BOOL f = __sbh_heap_check() >= 0;
LeaveCriticalSection(&g_csHeap);
return f;
}
return TRUE;
}
#else
#define CHECKSBH
#endif
#else
#if DBG == 1
BOOL CheckSmallBlockHeap()
{
return TRUE;
}
#endif
#endif SMALLBLOCKHEAP
//+------------------------------------------------------------------------
//
// Function: _MemGetSize
//
// Synopsis: Get size of block allocated with MemAlloc/MemRealloc.
//
// Note that MemAlloc/MemRealloc can allocate more than
// the requested number of bytes. Therefore the size returned
// from this function is possibly greater than the size
// passed to MemAlloc/Realloc.
//
// Arguments: [pv] - Return size of this block.
//
// Returns: The size of the block, or zero of pv == NULL.
//
//-------------------------------------------------------------------------
ULONG
_MemGetSize(void *pv)
{
if (pv == NULL)
return 0;
Assert(g_hProcessHeap);
#if SMALLBLOCKHEAP
#if DBG==1
if (IsSmallBlockHeapEnabled())
#endif
{
__sbh_region_t * preg;
__sbh_page_t * ppage;
__map_t * pmap;
EnterCriticalSection(&g_csHeap);
if ((pmap = __sbh_find_block(DbgPreGetSize(pv), &preg, &ppage)) != NULL )
{
size_t s = DbgPostGetSize(((size_t)(*pmap)) << _PARASHIFT);
LeaveCriticalSection(&g_csHeap);
return s;
}
LeaveCriticalSection(&g_csHeap);
}
#endif
return DbgPostGetSize(HeapSize(g_hProcessHeap, 0, DbgPreGetSize(pv)));
}
//+------------------------------------------------------------------------
//
// Function: _MemAlloc
//
// Synopsis: Allocate block of memory.
//
// The contents of the block are undefined. If the requested size
// is zero, this function returns a valid pointer. The returned
// pointer is guaranteed to be suitably aligned for storage of any
// object type.
//
// Arguments: [cb] - Number of bytes to allocate.
//
// Returns: Pointer to the allocated block, or NULL on error.
//
//-------------------------------------------------------------------------
void *
_MemAlloc(ULONG cb)
{
AssertSz (cb, "Requesting zero sized block.");
// The small-block heap will lose its mind if this ever happens, so we
// protect against the possibility.
if (cb == 0)
cb = 1;
Assert(g_hProcessHeap);
#if SMALLBLOCKHEAP
#if DBG==1
if (IsSmallBlockHeapEnabled())
#endif
{
/* round up to the nearest paragraph */
size_t cbr = (DbgPreAlloc(cb) + _PARASIZE - 1) & ~(_PARASIZE - 1);
if (cbr < __sbh_threshold)
{
CHECKSBH;
EnterCriticalSection(&g_csHeap);
void * pv = DbgPostAlloc(__sbh_alloc_block(cbr >> _PARASHIFT));
LeaveCriticalSection(&g_csHeap);
if (pv)
return pv;
}
}
#endif
return DbgPostAlloc(HeapAlloc(g_hProcessHeap, 0, DbgPreAlloc(cb)));
}
//+------------------------------------------------------------------------
// Function: _MemAllocClear
//
// Synopsis: Allocate a zero filled block of memory.
//
// If the requested size is zero, this function returns a valid
// pointer. The returned pointer is guaranteed to be suitably
// aligned for storage of any object type.
//
// Arguments: [cb] - Number of bytes to allocate.
//
// Returns: Pointer to the allocated block, or NULL on error.
//
//-------------------------------------------------------------------------
void *
_MemAllocClear(ULONG cb)
{
AssertSz (cb, "Allocating zero sized block.");
// The small-block heap will lose its mind if this ever happens, so we
// protect against the possibility.
if (cb == 0)
cb = 1;
void * pv;
Assert(g_hProcessHeap);
#if SMALLBLOCKHEAP
#if DBG==1
if (IsSmallBlockHeapEnabled())
#endif
{
/* round up to the nearest paragraph */
size_t cbr = (DbgPreAlloc(cb) + _PARASIZE - 1) & ~(_PARASIZE - 1);
if (cbr < __sbh_threshold)
{
CHECKSBH;
EnterCriticalSection(&g_csHeap);
pv = DbgPostAlloc(__sbh_alloc_block(cbr >> _PARASHIFT));
LeaveCriticalSection(&g_csHeap);
if (pv)
{
memset(pv, 0, cb);
return pv;
}
}
}
#endif
pv = DbgPostAlloc(HeapAlloc(g_hProcessHeap, HEAP_ZERO_MEMORY,
DbgPreAlloc(cb)));
// In debug, DbgPostAlloc set the memory so we need to clear it again.
#if DBG==1
if (pv)
{
memset(pv, 0, cb);
}
#endif
return pv;
}
//+------------------------------------------------------------------------
//
// Function: _MemFree
//
// Synopsis: Free a block of memory allocated with MemAlloc,
// MemAllocFree or MemRealloc.
//
// Arguments: [pv] - Pointer to block to free. A value of zero is
// is ignored.
//
//-------------------------------------------------------------------------
void
_MemFree(void *pv)
{
// The null check is required for HeapFree.
if (pv == NULL)
return;
Assert(g_hProcessHeap);
#if DBG == 1
pv = DbgPreFree(pv);
#endif
#if SMALLBLOCKHEAP
#if DBG==1
if (IsSmallBlockHeapEnabled())
#endif
{
__sbh_region_t *preg;
__sbh_page_t * ppage;
__map_t * pmap;
CHECKSBH;
EnterCriticalSection(&g_csHeap);
if ( (pmap = __sbh_find_block(pv, &preg, &ppage)) != NULL ) {
__sbh_free_block(preg, ppage, pmap);
LeaveCriticalSection(&g_csHeap);
DbgPostFree();
return;
}
LeaveCriticalSection(&g_csHeap);
}
#endif
HeapFree(g_hProcessHeap, 0, pv);
DbgPostFree();
}
//+------------------------------------------------------------------------
// Function: _MemRealloc
//
// Synopsis: Change the size of an existing block of memory, allocate a
// block of memory, or free a block of memory depending on the
// arguments.
//
// If cb is zero, this function always frees the block of memory
// and *ppv is set to zero.
//
// If cb is not zero and *ppv is zero, then this function allocates
// cb bytes.
//
// If cb is not zero and *ppv is non-zero, then this function
// changes the size of the block, possibly by moving it.
//
// On error, *ppv is left unchanged. The block contents remains
// unchanged up to the smaller of the new and old sizes. The
// contents of the block beyond the old size is undefined.
// The returned pointer is guaranteed to be suitably aligned for
// storage of any object type.
//
// The signature of this function is different than thy typical
// realloc-like function to avoid the following common error:
// pv = realloc(pv, cb);
// If realloc fails, then null is returned and the pointer to the
// original block of memory is leaked.
//
// Arguments: [cb] - Requested size in bytes. A value of zero always frees
// the block.
// [ppv] - On input, pointer to existing block pointer or null.
// On output, pointer to new block pointer.
//
// Returns: HRESULT
//
//-------------------------------------------------------------------------
HRESULT
_MemRealloc(void **ppv, ULONG cb)
{
void *pv;
Assert(g_hProcessHeap);
if (cb == 0)
{
_MemFree(*ppv);
*ppv = 0;
}
else if (*ppv == NULL)
{
*ppv = _MemAlloc(cb);
if (*ppv == NULL)
return E_OUTOFMEMORY;
}
else
{
#if DBG == 1
cb = DbgPreRealloc(*ppv, cb, &pv);
#else
pv = *ppv;
#endif
#if SMALLBLOCKHEAP
#if DBG==1
if (IsSmallBlockHeapEnabled())
#endif
{
__sbh_region_t *preg;
__sbh_page_t * ppage;
__map_t * pmap;
ULONG cbr;
void * pvNew;
cbr = (cb + _PARASIZE - 1) & ~(_PARASIZE - 1);
CHECKSBH;
EnterCriticalSection(&g_csHeap);
if ( (pmap = __sbh_find_block(pv, &preg, &ppage)) != NULL )
{
pvNew = NULL;
/*
* If the new size falls below __sbh_threshold, try to
* carry out the reallocation within the small-block
* heap.
*/
if ( cbr < __sbh_threshold ) {
if ( __sbh_resize_block(preg, ppage, pmap, cbr >> _PARASHIFT))
{
pvNew = pv;
}
else if ((pvNew = __sbh_alloc_block(cbr >> _PARASHIFT)) != NULL)
{
ULONG cbOld = ((size_t)(*pmap)) << _PARASHIFT;
memcpy(pvNew, pv, min(cbOld, cb));
__sbh_free_block(preg, ppage, pmap);
}
}
/*
* If the reallocation has not been (successfully)
* performed in the small-block heap, try to allocate a
* new block with HeapAlloc.
*/
if ((pvNew == NULL) && ((pvNew = HeapAlloc(g_hProcessHeap, 0, cb)) != NULL))
{
ULONG cbOld = ((size_t)(*pmap)) << _PARASHIFT;
memcpy(pvNew, pv, min(cbOld, cb));
__sbh_free_block(preg, ppage, pmap);
}
LeaveCriticalSection(&g_csHeap);
*ppv = DbgPostRealloc(pvNew);
if (*ppv)
{
return S_OK;
}
else
{
return E_OUTOFMEMORY;
}
}
else
{
LeaveCriticalSection(&g_csHeap);
}
}
#endif
pv = DbgPostRealloc(HeapReAlloc(g_hProcessHeap, 0, pv, cb));
if (pv == NULL)
return E_OUTOFMEMORY;
*ppv = pv;
}
return S_OK;
}
// MEMGUARD -------------------------------------------------------------------
#if defined(MEMGUARD)
#define MGGUARDDATA 0xF0F0BAAD
struct MGGUARD
{
MGGUARD *pNext;
DWORD dw;
};
MGGUARD * g_pMemList = NULL;
void
_MgMemValidate()
{
EnterCriticalSection(&g_csHeap);
MGGUARD *pg = g_pMemList;
while (pg)
{
if (pg->dw != MGGUARDDATA)
{
DebugBreak();
}
pg = pg->pNext;
}
LeaveCriticalSection(&g_csHeap);
}
void
_MgRemove(MGGUARD *pmg)
{
if (!pmg)
return;
EnterCriticalSection(&g_csHeap);
MGGUARD *pg = g_pMemList;
if (pmg == pg)
{
g_pMemList = pg->pNext;
goto Cleanup;
}
while (pg)
{
if (pg->pNext == pmg)
{
pg->pNext = pg->pNext->pNext;
break;
}
pg = pg->pNext;
}
Cleanup:
LeaveCriticalSection(&g_csHeap);
}
void
_MgAdd(MGGUARD *pmg)
{
EnterCriticalSection(&g_csHeap);
pmg->pNext = g_pMemList;
g_pMemList = pmg;
LeaveCriticalSection(&g_csHeap);
}
void *
_MgMemAlloc(ULONG cb)
{
_MgMemValidate();
MGGUARD * pmg = (MGGUARD *)_MemAlloc(sizeof(MGGUARD) + cb);
if (pmg)
{
pmg->dw = MGGUARDDATA;
_MgAdd(pmg);
return(pmg + 1);
}
else
{
return(NULL);
}
}
void *
_MgMemAllocClear(ULONG cb)
{
_MgMemValidate();
MGGUARD * pmg = (MGGUARD *)_MemAllocClear(sizeof(MGGUARD) + cb);
if (pmg)
{
pmg->dw = MGGUARDDATA;
_MgAdd(pmg);
return(pmg + 1);
}
else
{
return(NULL);
}
}
HRESULT
_MgMemRealloc(void ** ppv, ULONG cb)
{
_MgMemValidate();
if (cb == 0)
{
_MgMemFree(*ppv);
*ppv = 0;
return(S_OK);
}
if (*ppv == NULL)
{
*ppv = _MgMemAlloc(cb);
return(*ppv ? S_OK : E_OUTOFMEMORY);
}
MGGUARD * pmg = (MGGUARD *)*ppv - 1;
_MgRemove(pmg);
HRESULT hr = _MemRealloc((void **)&pmg, sizeof(MGGUARD) + cb);
if (hr == S_OK)
{
pmg->dw = MGGUARDDATA;
_MgAdd(pmg);
*ppv = pmg + 1;
}
return(hr);
}
ULONG
_MgMemGetSize(void * pv)
{
_MgMemValidate();
if (pv == NULL)
return(0);
else
return(_MemGetSize((MGGUARD *)pv - 1) - sizeof(MGGUARD));
}
void
_MgMemFree(void * pv)
{
_MgMemValidate();
if (pv)
{
MGGUARD * pmg = (MGGUARD *)pv - 1;
if (pmg->dw != MGGUARDDATA)
{
// The memory guard DWORD was overwritten! Bogus!
#ifdef _M_IX86
_asm int 3 // To get a proper stacktrace.
#else
DebugBreak();
#endif
}
_MgRemove(pmg);
_MemFree(pmg);
}
}
HRESULT
_MgMemAllocString(LPCTSTR pchSrc, LPTSTR * ppchDst)
{
TCHAR *pch;
size_t cb;
cb = (_tcsclen(pchSrc) + 1) * sizeof(TCHAR);
*ppchDst = pch = (TCHAR *)_MgMemAlloc(cb);
if (!pch)
return E_OUTOFMEMORY;
else
{
memcpy(pch, pchSrc, cb);
return S_OK;
}
}
HRESULT
_MgMemAllocString(ULONG cch, const TCHAR *pchSrc, TCHAR **ppchDest)
{
TCHAR *pch;
size_t cb = cch * sizeof(TCHAR);
*ppchDest = pch = (TCHAR *)_MgMemAlloc(cb + sizeof(TCHAR));
if (!pch)
return E_OUTOFMEMORY;
else
{
memcpy(pch, pchSrc, cb);
pch[cch] = 0;
return S_OK;
}
}
HRESULT
_MgMemReplaceString(const TCHAR *pchSrc, TCHAR **ppchDest)
{
HRESULT hr;
TCHAR *pch;
if (pchSrc)
{
hr = THR(_MgMemAllocString(pchSrc, &pch));
if (hr)
RRETURN(hr);
}
else
{
pch = NULL;
}
_MgMemFreeString(*ppchDest);
*ppchDest = pch;
return S_OK;
}
#endif // MEMGUARD