/* rndheap.c (emx+gcc) */

#ifdef HAVE_ALL
#define HAVE_HEAPCHK
#define HAVE_HEAPMIN
#define HAVE_HEAP_WALK
#define HAVE_UMALLOC
#define HAVE_USTATS
#endif

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <malloc.h>
#include <signal.h>
#include <time.h>
#include <getopt.h>
#include <errno.h>
#include <limits.h>
#if defined (HAVE_UMALLOC) || defined (HAVE_USTATS)
#include <umalloc.h>
#endif
#ifdef HAVE_UMALLOC
#include <stddef.h>
#include <sys/builtin.h>        /* For <sys/fmutex.h> */
#include <sys/fmutex.h>         /* For <sys/rmutex.h> */
#include <sys/rmutex.h>         /* For <emx/umalloc.h> */
#include <emx/umalloc.h>        /* For -T option */
#endif

#define SHOW_INTERVAL   60

struct block
{
  void *mem;
  size_t size;
  int tile;
};


static void *bottom;
static long iter;
static size_t block_size;
static size_t block_count;
static struct block *blocks;
static sig_atomic_t show_flag;
static int free_batch_factor;
static int free_batch;
static int sbrk_count;
static char tile;
static char quiet;
static char verbose;
static char prefer_realloc;
static char prefer_calloc;
static char heap_min;
static char tiled_heap;
static char check_all;
#ifdef __MT__
static int opt_m;
static long *thread_iter;
#endif
#ifdef HAVE_USTATS
static char show_stats;
#endif
#ifdef HAVE_HEAP_WALK
static FILE *dump_file;
static char *dump_fname = "rndheap.dmp";
#endif


static void usage (void)
{
  fputs ("Usage: rndheap [-n <count>] [-i <iterations>] "
         "[-f <free_batch_factor>]\n"
         "               [-w <waste>] [-s <size>]\n"
         "               [-a] [-b] [-c] [-t] [-v] "
#ifdef __MT__
         "[-m] "
#endif
         "[-q] [-r] "
#ifdef HAVE_HEAP_WALK
         "[-D] "
#endif
#ifdef HAVE_HEAPMIN
         "[-M] "
#endif
#ifdef HAVE_USTATS
         "[-S] "
#endif
         "<seed>\n", stderr);
  exit (1);
}


#ifdef HAVE_USTATS
static void do_stats (void)
{
  _HEAPSTATS stats;

  if (_ustats (_RUNTIME_HEAP, &stats) == 0)
    printf ("_provided=%lu, _used=%lu, _max_free=%lu, "
            "_segments=%lu, _crates=%lu\n",
            stats._provided, stats._used, stats._max_free,
            stats._segments, stats._crates);
}
#endif


static void show (void)
{
  long used_size;
  int i, count, tiled;

  used_size = 0; count = 0; tiled = 0;
  for (i = 0; i < block_count; ++i)
    if (blocks[i].mem != NULL)
      {
        ++count;
        if (blocks[i].tile)
          ++tiled;
        used_size += blocks[i].size;
      }
  printf ("Iterations:%10ld, blocks:%4d, tiled:%4d, used:%8ld, size:%8ld\n",
          iter, count, tiled, used_size,
          (size_t)sbrk (0) - (size_t)bottom);
#ifdef HAVE_USTATS
  if (show_stats)
    do_stats ();
#endif
#ifdef __MT__
  if (opt_m != 0 && thread_iter != NULL)
    {
      printf ("Thread iterations:");
      for (i = 0; i < opt_m; ++i)
        printf (" %ld", thread_iter[i]);
      putchar ('\n');
    }
#endif
}


static void check_tile (void *p, size_t s)
{
  char ok;

  if (s == 0)
    ok = 1;
  else if (s >= 0x10000)
    ok = (((unsigned long)p & 0xffff) == 0);
  else
    ok = (((unsigned long)p & ~0xffff) == (((unsigned long)p+s-1) & ~0xffff));
  if (!ok)
    {
      fprintf (stderr, "Tiling error in iteration %ld\n", iter);
      abort ();
    }
}


static void check_zero (char *p, size_t s)
{
  while (s != 0 && *p == 0)
    ++p, --s;
  if (s != 0)
    abort ();
}


static void handler (int sig)
{
  if (show_flag)
    {
      fprintf (stderr, "Infinite loop, aborted in iteration %ld\n", iter);
      abort ();
    }
  show_flag = 1;
  signal (sig, SIG_ACK);
  alarm (SHOW_INTERVAL);
}


static int rand_block (void)
{
  if (block_count > RAND_MAX)
    return (rand () ^ (rand () << 15)) % block_count;
  else
    return rand ()  % block_count;
}


static size_t rand_block_size (void)
{
  if (block_size > RAND_MAX)
    return (size_t)(rand () ^ (rand () << 15)) % block_size;
  else
    return (size_t)rand () % block_size;
}


#ifdef __MT__
static void thread (void *arg)
{
  char *p;
  size_t s;
  int i;
  long *piter;

  piter = (long *)arg;
  for (;;)
    {
      s = rand_block_size ();
      p = malloc (s);
      if (p == NULL)
        abort ();
      if ((i = _heapchk ()) != _HEAPOK)
        abort ();

      s = rand_block_size ();
      if (s == 0) s = 1;
      p = realloc (p, s);
      if (p == NULL)
        abort ();
      if ((i = _heapchk ()) != _HEAPOK)
        abort ();

      free (p);
      if ((i = _heapchk ()) != _HEAPOK)
        abort ();

      *piter += 1;
    }
}
#endif


#ifdef HAVE_HEAP_WALK
static int walk_dump (const void *obj, size_t size, int flag, int status,
                      const char *fname, size_t lineno)
{
  if (status == _HEAPOK)
    fprintf (dump_file, "0x%.8lx %.10lu %s\n",
             (unsigned long)obj, (unsigned long)size,
             flag == _USEDENTRY ? "used" : "free");
  return status;
}
#endif


static void out_of_mem (size_t size)
{
  fprintf (stderr, "Out of memory in iteration %ld (requested size=%lu)\n",
           iter, size);
#ifdef HAVE_USTATS
  {
    _HEAPSTATS stats;

    if (_ustats (_RUNTIME_HEAP, &stats) == 0)
      printf ("ustats(): _provided=%lu, _used=%lu, _max_free=%lu\n",
              stats._provided, stats._used, stats._max_free);
    if (dump_file != NULL)
      {
        _uheap_walk (_RUNTIME_HEAP, walk_dump);
        if (fclose (dump_file) != 0)
          printf ("Error writing `%s'\n", dump_fname);
        else
          printf ("`%s' written\n", dump_fname);
      }
  }
#endif
  abort ();
}


static void check_heap (void)
{
#ifdef HAVE_HEAPCHK
  int i;

  i = _heapchk ();
  if (i != _HEAPOK && i != _HEAPEMPTY)
    {
      fprintf (stderr, "Error %d in iteration %ld\n", i, iter);
      abort ();
    }
#endif
}


static void do_free (void)
{
  void *p;
  int j;

  j = rand_block ();
  p = blocks[j].mem;
  if (p != NULL)
    {
      if (blocks[j].tile)
        _tfree (p);
      else
        free (p);
      blocks[j].mem = NULL;
    }
}


static void do_alloc (int i, int use_calloc)
{
  void *p;
  int j;
  size_t s;

  j = rand_block ();
  p = blocks[j].mem;
  if (p != NULL)
    {
      if (blocks[j].tile)
        _tfree (p);
      else
        free (p);
      if (check_all)
        check_heap ();
    }
  s = rand_block_size ();
  blocks[j].tile = tile && ((i >> 2) & 1);
  blocks[j].size = s;
  if (blocks[j].tile)
    p = use_calloc ? _tcalloc (s, 1) : _tmalloc (s);
  else
    p = use_calloc ? calloc (s, 1) : malloc (s);
  if (p == NULL && s != 0)
    out_of_mem (s);
  blocks[j].mem = p;
  if (p != NULL)
    {
      if (blocks[j].tile || tiled_heap)
        check_tile (p, s);
      if (use_calloc)
        check_zero (p, s);
      memset (p, 0x77, s);
    }
}


static void do_realloc (void)
{
  void *p0, *p1;
  int j;
  size_t s;

  j = rand_block ();
  p0 = blocks[j].mem;
  s = rand_block_size ();
  blocks[j].size = s;
  if (blocks[j].tile)
    p1 = _trealloc (p0, s);
  else
    p1 = realloc (p0, s);
  if (s != 0 && p1 == NULL)
    out_of_mem (s);
  if (p0 != NULL && s == 0 && p1 != NULL)
    abort ();
  blocks[j].mem = p1;
  if (p1 != NULL)
    {
      if (blocks[j].tile || tiled_heap)
        check_tile (p1, s);
      memset (p1, 0x77, s);
    }
}


static void do_minimize (void)
{
#ifdef HAVE_HEAPMIN
  if (verbose)
    {
      printf ("Minimizing heap...\n");
      show ();
    }
  check_heap ();
  if (_heapmin () != 0)
    abort ();
  check_heap ();
  if (verbose)
    {
      show ();
      printf ("Minimizing heap done\n");
    }
#endif
}


static void do_other (void)
{
  int i, j;

  if (heap_min || sbrk_count != 0 || free_batch_factor != 0)
    {
      j = rand () ^ (rand () << 15);
      if (free_batch_factor != 0 && j % free_batch_factor == 0)
        {
          j = 0;
          for (i = 0; i < block_count; ++i)
            if (blocks[i].mem != NULL && i != j)
              {
                blocks[j++] = blocks[i];
                blocks[i].mem = NULL;
              }
          free_batch = j;
          if (rand () % 3 >= 1)
            free_batch /= 2;
          if (verbose)
            printf ("Freeing batch of blocks (%d of %d)\n",
                    free_batch, j);
        }
      else if (sbrk_count != 0 && j % 2017 == 0)
        {
          if (verbose)
            printf ("Calling sbrk(1)...\n");
          if (sbrk (1) == (void *)-1)
            abort ();
          sbrk_count -= 1;
        }
      else if (heap_min && j % 4096 == 0)
        do_minimize ();
    }
}


#ifdef BREAK_ITER
static void bpt (void)
{
}
#endif


int main (int argc, char *argv[])
{
  int c, i, k;
  long n, iter_limit;
  char *q;

  bottom = sbrk (0);
  setvbuf (stdout, NULL, _IOLBF, BUFSIZ);

  iter_limit = LONG_MAX; block_size = 0x8000; block_count = 256;
  while ((c = getopt (argc, argv, "abcf:i:mn:qrs:tvw:DMST")) != -1)
    switch (c)
      {
      case 'a':
        check_all = 1;
        break;
      case 'b':
        ++sbrk_count;
        break;
      case 'c':
        prefer_calloc = 1;
        break;
      case 'f':
        errno = 0;
        free_batch_factor = strtol (optarg, &q, 0);
        if (errno != 0 || free_batch_factor < 0 || *q != 0)
          usage ();
        break;
      case 'i':
        errno = 0;
        iter_limit = strtol (optarg, &q, 0);
        if (errno != 0 || iter_limit < 1 || *q != 0)
          usage ();
        break;
#ifdef __MT__
      case 'm':
        opt_m += 1;
        break;
#endif
      case 'n':
        errno = 0;
        block_count = strtol (optarg, &q, 0);
        if (errno != 0 || block_count < 2 || *q != 0)
          usage ();
        break;
      case 'q':
        quiet = 1; verbose = 0;
        break;
      case 'r':
        prefer_realloc = 1;
        break;
      case 's':
        errno = 0;
        block_size = strtol (optarg, &q, 0);
        if (errno != 0 || block_size < 2 || *q != 0)
          usage ();
        break;
      case 't':
        tile = 1; 
        break;
      case 'v':
        verbose = 1; quiet = 0;
        break;
      case 'w':
        errno = 0;
        n = strtol (optarg, &q, 0);
        if (errno != 0 || n < 0 || *q != 0)
          usage ();
        if (sbrk ((int)n) == (void *)-1)
          printf ("sbrk(%ld) failed\n", n);
        break;
#ifdef HAVE_HEAP_WALK
      case 'D':
        if (dump_file != NULL)
          usage ();
        dump_file = fopen (dump_fname, "w");
        if (dump_file == NULL)
          {
            printf ("Cannot create `%s'\n", dump_fname);
            exit (1);
          }
        /* Avoid deadlock */
        setvbuf (dump_file, NULL, _IOFBF, BUFSIZ);
        break;
#endif
#ifdef HAVE_HEAPMIN
      case 'M':
        heap_min = 1;
        break;
#endif
#ifdef HAVE_USTATS
      case 'S':
        show_stats = 1;
        break;
#endif
#ifdef HAVE_UMALLOC
      case 'T':
        tiled_heap = 1;
        break;
#endif
      default:
        usage ();
      }
  argv += optind;
  argc -= optind;
  if (argc != 1)
    usage ();

  errno = 0;
  n = strtol (argv[0], &q, 0);
  if (errno != 0 || *q != 0)
    usage ();

  srand ((unsigned)n);

  blocks = alloca (block_count * sizeof (*blocks));

  for (i = 0; i < block_count; ++i)
    {
      blocks[i].mem = NULL;
      blocks[i].size = 0;
      blocks[i].tile = 0;
    }

  signal (SIGALRM, handler);

  if (iter_limit != LONG_MAX && !quiet)
    printf ("Start time: %ld\n", (long)clock ());

#ifdef HAVE_UMALLOC
  if (tiled_heap)
    {
      _tmalloc (0);             /* Create the heap */
      _RUNTIME_HEAP->type = _HEAP_TILED;
      if (_heapchk () != _HEAPOK)
        abort ();
      tile = 0;
    }
#endif

  if (!quiet)
    {
      show ();
      alarm (SHOW_INTERVAL);
    }

#ifdef __MT__
  if (opt_m != 0)
    {
      thread_iter = alloca (opt_m * sizeof (*thread_iter));
      for (i = 0; i < opt_m; ++i)
        thread_iter[i] = 0;
      for (i = 0; i < opt_m; ++i)
        if (_beginthread (thread, NULL, 0x8000, thread_iter + i) == -1)
          abort ();
    }
#endif

  k = 0; free_batch = 0;
  for (iter = 0; iter < iter_limit; ++iter)
    {
#ifdef BREAK_ITER
      if (iter == BREAK_ITER)
        {
          alarm (0);
          bpt ();
        }
#endif
      if (show_flag)
        {
          show_flag = 0;
          show ();
        }
      if (k == 97 || check_all)
        check_heap ();
      if (free_batch > 0)
        {
          free_batch -= 1;
          if (blocks[free_batch].mem == NULL)
            abort ();
          if (blocks[free_batch].tile)
            _tfree (blocks[free_batch].mem);
          else
            free (blocks[free_batch].mem);
          blocks[free_batch].mem = NULL;
          if (free_batch == 0 && verbose)
            printf ("Freeing batch of blocks done\n");
          continue;
        }
      i = rand () / 253;
      if ((prefer_realloc || prefer_calloc) && i % 32 >= 27)
        {
          if (!prefer_calloc)
            do_realloc ();
          else if (!prefer_realloc)
            do_alloc (i, 1);
          else if (i & 32)
            do_realloc ();
          else
            do_alloc (i, 1);
        }
      else
        switch (i % 4)
          {
          case 0:
            do_free ();
            break;
          case 1:
            do_alloc (i, i & 32);
            break;
          case 2:
            do_realloc ();
            break;
          case 3:
            do_other ();
            break;
          }
    }

  if (!quiet)
    show ();

  check_heap ();
  for (i = 0; i < block_count; ++i)
    {
      if (blocks[i].mem != NULL)
        {
          free (blocks[i].mem);
          blocks[i].mem = NULL;
          check_heap ();
#ifdef HAVE_HEAPMIN
          if (heap_min)
            {
              do_minimize ();
              check_heap ();
            }
#endif
        }
    }
#ifdef HAVE_USTATS
  if (show_stats)
    do_stats ();
#endif

  if (!quiet)
    printf ("End time:   %ld\n", (long)clock ());
  return 0;
}
