/* -*- 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 */