/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 2 -*- */

/* This Source Code Form is subject to the terms of the Mozilla Public

 * License, v. 2.0. If a copy of the MPL was not distributed with this

 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/*

** Thread safe versions of malloc, free, realloc, calloc and cfree.

*/

#include "primpl.h"

#ifdef _PR_ZONE_ALLOCATOR

/*

** The zone allocator code must use native mutexes and cannot

** use PRLocks because PR_NewLock calls PR_Calloc, resulting

** in cyclic dependency of initialization.

*/

#  include <string.h>

union memBlkHdrUn;

typedef struct MemoryZoneStr {

  union memBlkHdrUn* head; /* free list */

  pthread_mutex_t lock;

  size_t blockSize;    /* size of blocks on this free list */

  PRUint32 locked;     /* current state of lock */

  PRUint32 contention; /* counter: had to wait for lock */

  PRUint32 hits;       /* allocated from free list */

  PRUint32 misses;     /* had to call malloc */

  PRUint32 elements;   /* on free list */

} MemoryZone;

typedef union memBlkHdrUn {

  unsigned char filler[48]; /* fix the size of this beast */

  struct memBlkHdrStr {

    union memBlkHdrUn* next;

    MemoryZone* zone;

    size_t blockSize;

    size_t requestedSize;

    PRUint32 magic;

  } s;

} MemBlockHdr;

#  define MEM_ZONES 7

#  define THREAD_POOLS 11 /* prime number for modulus */

#  define ZONE_MAGIC 0x0BADC0DE

static MemoryZone zones[MEM_ZONES][THREAD_POOLS];

static PRBool use_zone_allocator = PR_FALSE;

static void pr_ZoneFree(void* ptr);

void _PR_DestroyZones(void) {

  int i, j;

  if (!use_zone_allocator) {

    return;

  }

  for (j = 0; j < THREAD_POOLS; j++) {

    for (i = 0; i < MEM_ZONES; i++) {

      MemoryZone* mz = &zones[i][j];

      pthread_mutex_destroy(&mz->lock);

      while (mz->head) {

        MemBlockHdr* hdr = mz->head;

        mz->head = hdr->s.next; /* unlink it */

        free(hdr);

        mz->elements--;

      }

    }

  }

  use_zone_allocator = PR_FALSE;

}

/*

** pr_FindSymbolInProg

**

** Find the specified data symbol in the program and return

** its address.

*/

#  ifdef HAVE_DLL

#    if defined(USE_DLFCN) && !defined(NO_DLOPEN_NULL)

#      include <dlfcn.h>

static void* pr_FindSymbolInProg(const char* name) {

  void* h;

  void* sym;

  h = dlopen(0, RTLD_LAZY);

  if (h == NULL) {

    return NULL;

  }

  sym = dlsym(h, name);

  (void)dlclose(h);

  return sym;

}

#    elif defined(USE_HPSHL)

#      include <dl.h>

static void* pr_FindSymbolInProg(const char* name) {

  shl_t h = NULL;

  void* sym;

  if (shl_findsym(&h, name, TYPE_DATA, &sym) == -1) {

    return NULL;

  }

  return sym;

}

#    elif defined(USE_MACH_DYLD) || defined(NO_DLOPEN_NULL)

static void* pr_FindSymbolInProg(const char* name) {

  /* FIXME: not implemented */

  return NULL;

}

#    else

#      error "The zone allocator is not supported on this platform"

#    endif

#  else /* !defined(HAVE_DLL) */

static void* pr_FindSymbolInProg(const char* name) {

  /* can't be implemented */

  return NULL;

}

#  endif /* HAVE_DLL */

void _PR_InitZones(void) {

  int i, j;

  char* envp;

  PRBool* sym;

  if ((sym = (PRBool*)pr_FindSymbolInProg("nspr_use_zone_allocator")) != NULL) {

    use_zone_allocator = *sym;

  } else if ((envp = getenv("NSPR_USE_ZONE_ALLOCATOR")) != NULL) {

    use_zone_allocator = (atoi(envp) == 1);

  }

  if (!use_zone_allocator) {

    return;

  }

  for (j = 0; j < THREAD_POOLS; j++) {

    for (i = 0; i < MEM_ZONES; i++) {

      MemoryZone* mz = &zones[i][j];

      int rv = pthread_mutex_init(&mz->lock, NULL);

      PR_ASSERT(0 == rv);

      if (rv != 0) {

        goto loser;

      }

      mz->blockSize = 16 << (2 * i);

    }

  }

  return;

loser:

  _PR_DestroyZones();

  return;

}

PR_IMPLEMENT(void)

PR_FPrintZoneStats(PRFileDesc* debug_out) {

  int i, j;

  for (j = 0; j < THREAD_POOLS; j++) {

    for (i = 0; i < MEM_ZONES; i++) {

      MemoryZone* mz = &zones[i][j];

      MemoryZone zone = *mz;

      if (zone.elements || zone.misses || zone.hits) {

        PR_fprintf(debug_out,

                   "pool: %d, zone: %d, size: %d, free: %d, hit: %d, miss: %d, "

                   "contend: %d\n",

                   j, i, zone.blockSize, zone.elements, zone.hits, zone.misses,

                   zone.contention);

      }

    }

  }

}

static void* pr_ZoneMalloc(PRUint32 size) {

  void* rv;

  unsigned int zone;

  size_t blockSize;

  MemBlockHdr *mb, *mt;

  MemoryZone* mz;

  /* Always allocate a non-zero amount of bytes */

  if (size < 1) {

    size = 1;

  }

  for (zone = 0, blockSize = 16; zone < MEM_ZONES; ++zone, blockSize <<= 2) {

    if (size <= blockSize) {

      break;

    }

  }

  if (zone < MEM_ZONES) {

    pthread_t me = pthread_self();

    unsigned int pool = (PRUptrdiff)me % THREAD_POOLS;

    PRUint32 wasLocked;

    mz = &zones[zone][pool];

    wasLocked = mz->locked;

    pthread_mutex_lock(&mz->lock);

    mz->locked = 1;

    if (wasLocked) {

      mz->contention++;

    }

    if (mz->head) {

      mb = mz->head;

      PR_ASSERT(mb->s.magic == ZONE_MAGIC);

      PR_ASSERT(mb->s.zone == mz);

      PR_ASSERT(mb->s.blockSize == blockSize);

      PR_ASSERT(mz->blockSize == blockSize);

      mt = (MemBlockHdr*)(((char*)(mb + 1)) + blockSize);

      PR_ASSERT(mt->s.magic == ZONE_MAGIC);

      PR_ASSERT(mt->s.zone == mz);

      PR_ASSERT(mt->s.blockSize == blockSize);

      mz->hits++;

      mz->elements--;

      mz->head = mb->s.next; /* take off free list */

      mz->locked = 0;

      pthread_mutex_unlock(&mz->lock);

      mt->s.next = mb->s.next = NULL;

      mt->s.requestedSize = mb->s.requestedSize = size;

      rv = (void*)(mb + 1);

      return rv;

    }

    mz->misses++;

    mz->locked = 0;

    pthread_mutex_unlock(&mz->lock);

    mb = (MemBlockHdr*)malloc(blockSize + 2 * (sizeof *mb));

    if (!mb) {

      PR_SetError(PR_OUT_OF_MEMORY_ERROR, 0);

      return NULL;

    }

    mb->s.next = NULL;

    mb->s.zone = mz;

    mb->s.magic = ZONE_MAGIC;

    mb->s.blockSize = blockSize;

    mb->s.requestedSize = size;

    mt = (MemBlockHdr*)(((char*)(mb + 1)) + blockSize);

    memcpy(mt, mb, sizeof *mb);

    rv = (void*)(mb + 1);

    return rv;

  }

  /* size was too big.  Create a block with no zone */

  blockSize = (size & 15) ? size + 16 - (size & 15) : size;

  mb = (MemBlockHdr*)malloc(blockSize + 2 * (sizeof *mb));

  if (!mb) {

    PR_SetError(PR_OUT_OF_MEMORY_ERROR, 0);

    return NULL;

  }

  mb->s.next = NULL;

  mb->s.zone = NULL;

  mb->s.magic = ZONE_MAGIC;

  mb->s.blockSize = blockSize;

  mb->s.requestedSize = size;

  mt = (MemBlockHdr*)(((char*)(mb + 1)) + blockSize);

  memcpy(mt, mb, sizeof *mb);

  rv = (void*)(mb + 1);

  return rv;

}

static void* pr_ZoneCalloc(PRUint32 nelem, PRUint32 elsize) {

  PRUint32 size = nelem * elsize;

  void* p = pr_ZoneMalloc(size);

  if (p) {

    memset(p, 0, size);

  }

  return p;

}

static void* pr_ZoneRealloc(void* oldptr, PRUint32 bytes) {

  void* rv;

  MemBlockHdr* mb;

  int ours;

  MemBlockHdr phony;

  if (!oldptr) {

    return pr_ZoneMalloc(bytes);

  }

  mb = (MemBlockHdr*)((char*)oldptr - (sizeof *mb));

  if (mb->s.magic != ZONE_MAGIC) {

    /* Maybe this just came from ordinary malloc */

#  ifdef DEBUG

    fprintf(stderr,

            "Warning: reallocing memory block %p from ordinary malloc\n",

            oldptr);

#  endif

    /*

     * We are going to realloc oldptr.  If realloc succeeds, the

     * original value of oldptr will point to freed memory.  So this

     * function must not fail after a successfull realloc call.  We

     * must perform any operation that may fail before the realloc

     * call.

     */

    rv = pr_ZoneMalloc(bytes); /* this may fail */

    if (!rv) {

      return rv;

    }

    /* We don't know how big it is.  But we can fix that. */

    oldptr = realloc(oldptr, bytes);

    /*

     * If realloc returns NULL, this function loses the original

     * value of oldptr.  This isn't a leak because the caller of

     * this function still has the original value of oldptr.

     */

    if (!oldptr) {

      if (bytes) {

        PR_SetError(PR_OUT_OF_MEMORY_ERROR, 0);

        pr_ZoneFree(rv);

        return oldptr;

      }

    }

    phony.s.requestedSize = bytes;

    mb = &phony;

    ours = 0;

  } else {

    size_t blockSize = mb->s.blockSize;

    MemBlockHdr* mt = (MemBlockHdr*)(((char*)(mb + 1)) + blockSize);

    PR_ASSERT(mt->s.magic == ZONE_MAGIC);

    PR_ASSERT(mt->s.zone == mb->s.zone);

    PR_ASSERT(mt->s.blockSize == blockSize);

    if (bytes <= blockSize) {

      /* The block is already big enough. */

      mt->s.requestedSize = mb->s.requestedSize = bytes;

      return oldptr;

    }

    ours = 1;

    rv = pr_ZoneMalloc(bytes);

    if (!rv) {

      return rv;

    }

  }

  if (oldptr && mb->s.requestedSize) {

    memcpy(rv, oldptr, mb->s.requestedSize);

  }

  if (ours) {

    pr_ZoneFree(oldptr);

  } else if (oldptr) {

    free(oldptr);

  }

  return rv;

}

static void pr_ZoneFree(void* ptr) {

  MemBlockHdr *mb, *mt;

  MemoryZone* mz;

  size_t blockSize;

  PRUint32 wasLocked;

  if (!ptr) {

    return;

  }

  mb = (MemBlockHdr*)((char*)ptr - (sizeof *mb));

  if (mb->s.magic != ZONE_MAGIC) {

    /* maybe this came from ordinary malloc */

#  ifdef DEBUG

    fprintf(stderr, "Warning: freeing memory block %p from ordinary malloc\n",

            ptr);

#  endif

    free(ptr);

    return;

  }

  blockSize = mb->s.blockSize;

  mz = mb->s.zone;

  mt = (MemBlockHdr*)(((char*)(mb + 1)) + blockSize);

  PR_ASSERT(mt->s.magic == ZONE_MAGIC);

  PR_ASSERT(mt->s.zone == mz);

  PR_ASSERT(mt->s.blockSize == blockSize);

  if (!mz) {

    PR_ASSERT(blockSize > 65536);

    /* This block was not in any zone.  Just free it. */

    free(mb);

    return;

  }

  PR_ASSERT(mz->blockSize == blockSize);

  wasLocked = mz->locked;

  pthread_mutex_lock(&mz->lock);

  mz->locked = 1;

  if (wasLocked) {

    mz->contention++;

  }

  mt->s.next = mb->s.next = mz->head; /* put on head of list */

  mz->head = mb;

  mz->elements++;

  mz->locked = 0;

  pthread_mutex_unlock(&mz->lock);

}

PR_IMPLEMENT(void*) PR_Malloc(PRUint32 size) {

  if (!_pr_initialized) {

    _PR_ImplicitInitialization();

  }

  return use_zone_allocator ? pr_ZoneMalloc(size) : malloc(size);

}

PR_IMPLEMENT(void*) PR_Calloc(PRUint32 nelem, PRUint32 elsize) {

  if (!_pr_initialized) {

    _PR_ImplicitInitialization();

  }

  return use_zone_allocator ? pr_ZoneCalloc(nelem, elsize)

                            : calloc(nelem, elsize);

}

PR_IMPLEMENT(void*) PR_Realloc(void* ptr, PRUint32 size) {

  if (!_pr_initialized) {

    _PR_ImplicitInitialization();

  }

  return use_zone_allocator ? pr_ZoneRealloc(ptr, size) : realloc(ptr, size);

}

PR_IMPLEMENT(void) PR_Free(void* ptr) {

  if (use_zone_allocator) {

    pr_ZoneFree(ptr);

  } else {

    free(ptr);

  }

}

#else /* !defined(_PR_ZONE_ALLOCATOR) */

/*

** The PR_Malloc, PR_Calloc, PR_Realloc, and PR_Free functions simply

** call their libc equivalents now.  This may seem redundant, but it

** ensures that we are calling into the same runtime library.  On

** Win32, it is possible to have multiple runtime libraries (e.g.,

** objects compiled with /MD and /MDd) in the same process, and

** they maintain separate heaps, which cannot be mixed.

*/

PR_IMPLEMENT(void*) PR_Malloc(PRUint32 size) { return malloc(size); }

PR_IMPLEMENT(void*) PR_Calloc(PRUint32 nelem, PRUint32 elsize) {

  return calloc(nelem, elsize);

}

PR_IMPLEMENT(void*) PR_Realloc(void* ptr, PRUint32 size) {

  return realloc(ptr, size);

}

PR_IMPLEMENT(void) PR_Free(void* ptr) { free(ptr); }

#endif /* _PR_ZONE_ALLOCATOR */

/*

** Complexity alert!

**

** If malloc/calloc/free (etc.) were implemented to use pr lock's then

** the entry points could block when called if some other thread had the

** lock.

**

** Most of the time this isn't a problem. However, in the case that we

** are using the thread safe malloc code after PR_Init but before

** PR_AttachThread has been called (on a native thread that nspr has yet

** to be told about) we could get royally screwed if the lock was busy

** and we tried to context switch the thread away. In this scenario

**  PR_CURRENT_THREAD() == NULL

**

** To avoid this unfortunate case, we use the low level locking

** facilities for malloc protection instead of the slightly higher level

** locking. This makes malloc somewhat faster so maybe it's a good thing

** anyway.

*/

#ifdef _PR_OVERRIDE_MALLOC

/* Imports */

extern void* _PR_UnlockedMalloc(size_t size);

extern void* _PR_UnlockedMemalign(size_t alignment, size_t size);

extern void _PR_UnlockedFree(void* ptr);

extern void* _PR_UnlockedRealloc(void* ptr, size_t size);

extern void* _PR_UnlockedCalloc(size_t n, size_t elsize);

static PRBool _PR_malloc_initialised = PR_FALSE;

#  ifdef _PR_PTHREADS

static pthread_mutex_t _PR_MD_malloc_crustylock;

#    define _PR_Lock_Malloc()                                 \

      {                                                       \

        if (PR_TRUE == _PR_malloc_initialised) {              \

          PRStatus rv;                                        \

          rv = pthread_mutex_lock(&_PR_MD_malloc_crustylock); \

          PR_ASSERT(0 == rv);                                 \

        }

#    define _PR_Unlock_Malloc()                               \

      if (PR_TRUE == _PR_malloc_initialised) {                \

        PRStatus rv;                                          \

        rv = pthread_mutex_unlock(&_PR_MD_malloc_crustylock); \

        PR_ASSERT(0 == rv);                                   \

      }                                                       \

      }

#  else /* _PR_PTHREADS */

static _MDLock _PR_MD_malloc_crustylock;

#    define _PR_Lock_Malloc()                                 \

      {                                                       \

        PRIntn _is;                                           \

        if (PR_TRUE == _PR_malloc_initialised) {              \

          if (_PR_MD_CURRENT_THREAD() &&                      \

              !_PR_IS_NATIVE_THREAD(_PR_MD_CURRENT_THREAD())) \

            _PR_INTSOFF(_is);                                 \

          _PR_MD_LOCK(&_PR_MD_malloc_crustylock);             \

        }

#    define _PR_Unlock_Malloc()                             \

      if (PR_TRUE == _PR_malloc_initialised) {              \

        _PR_MD_UNLOCK(&_PR_MD_malloc_crustylock);           \

        if (_PR_MD_CURRENT_THREAD() &&                      \

            !_PR_IS_NATIVE_THREAD(_PR_MD_CURRENT_THREAD())) \

          _PR_INTSON(_is);                                  \

      }                                                     \

      }

#  endif /* _PR_PTHREADS */

PR_IMPLEMENT(PRStatus) _PR_MallocInit(void) {

  PRStatus rv = PR_SUCCESS;

  if (PR_TRUE == _PR_malloc_initialised) {

    return PR_SUCCESS;

  }

#  ifdef _PR_PTHREADS

  {

    int status;

    pthread_mutexattr_t mattr;

    status = _PT_PTHREAD_MUTEXATTR_INIT(&mattr);

    PR_ASSERT(0 == status);

    status = _PT_PTHREAD_MUTEX_INIT(_PR_MD_malloc_crustylock, mattr);

    PR_ASSERT(0 == status);

    status = _PT_PTHREAD_MUTEXATTR_DESTROY(&mattr);

    PR_ASSERT(0 == status);

  }

#  else  /* _PR_PTHREADS */

  _MD_NEW_LOCK(&_PR_MD_malloc_crustylock);

#  endif /* _PR_PTHREADS */

  if (PR_SUCCESS == rv) {

    _PR_malloc_initialised = PR_TRUE;

  }

  return rv;

}

void* malloc(size_t size) {

  void* p;

  _PR_Lock_Malloc();

  p = _PR_UnlockedMalloc(size);

  _PR_Unlock_Malloc();

  return p;

}

void free(void* ptr) {

  _PR_Lock_Malloc();

  _PR_UnlockedFree(ptr);

  _PR_Unlock_Malloc();

}

void* realloc(void* ptr, size_t size) {

  void* p;

  _PR_Lock_Malloc();

  p = _PR_UnlockedRealloc(ptr, size);

  _PR_Unlock_Malloc();

  return p;

}

void* calloc(size_t n, size_t elsize) {

  void* p;

  _PR_Lock_Malloc();

  p = _PR_UnlockedCalloc(n, elsize);

  _PR_Unlock_Malloc();

  return p;

}

void cfree(void* p) {

  _PR_Lock_Malloc();

  _PR_UnlockedFree(p);

  _PR_Unlock_Malloc();

}

void _PR_InitMem(void) {

  PRStatus rv;

  rv = _PR_MallocInit();

  PR_ASSERT(PR_SUCCESS == rv);

}

#endif /* _PR_OVERRIDE_MALLOC */