/* **** FILE: malloc.c
----------------------------------------------------------------------
This file implements a replacement for the libc malloc facility.  This
mechanism works in a manner very similar to AMD64, VAX, IBM z series,
etc., multi level paging paging.  This file, therefore, would make
most sense if the reader consulted the AMD64 architecture developer's
manual, or any other architectural manual that explains an
architecture that implements multi level paged virtual memory mapping.
----------------------------------------------------------------------
*/
/*
 * Copyright (C) 2019 Robin Miyagi ...................................
 *
 *    web:     https://www.linuxprogramming.ca/
 *
 *    email:   r.miyagi@linuxprogramming.ca
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation; either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/*
 * Forward declarations...............................................
 */
#include <stddef.h>

void* malloc_blk  (size_t size);
void  free_blk 	  (void* ptr);
void* realloc_blk (void* ptr, size_t size);
/*
 * Project Includes...................................................
 */

/*
 * System includes....................................................
 */
#ifndef _GNU_SOURCE
# define _GNU_SOURCE
#endif
#include <unistd.h>
#include <stddef.h>
#include <errno.h>
#include <sys/mman.h>
#include <assert.h>
#include <string.h>

/*
 * Locally defined interfaces.........................................
 */
/* ===================================================================
 *      		P A G E   T A B L E S                      
 * ===================================================================
 */
/*
 * Type definition for multi level page table like structure of
 * allocation sizes.
 */
typedef struct
{
  void*   top;	/* Pointer to top level page table. */
  int     ps;	/* page size. */
  size_t  lev;	/* number of levels. */                
} t_allocsize;

static _Thread_local t_allocsize  init_sizes;	
static _Thread_local t_allocsize* alloc_sizes = NULL;

/*
 * Type definition for bottom table entries.  Bottom table entries are
 * the last entries to be reached through the pointer chain.
 */
typedef struct
{
  size_t as;	/* The size, as actually allocated.  */
  size_t us;	/* The size, as requested by the user.  */
} t_bottom_entry;

/*
 * Union helpful for parsing a given address into individual level
 * indexes using bit wise shift operations.
 */
union u_intaddr
{
  unsigned long l;	/* integer treatment */
  void*         p;	/* pointer treatment */                     
};

/*
 * bits in (the maximum number of entries per table - 1), recalling
 * that indexes are numbered starting from zero.
 */
#define PAGE_TABLE_SIZE_BITS  	12

/*
 * ==== SEE COMMENT IN PREVIOUS LINE ==== 
 */
#define MAX_ENTRIES_PER_TABLE 	(1 << PAGE_TABLE_SIZE_BITS)

/*
 * Well known fact --- defined for the sake of clarity.
 */
#define BITS_IN_BYTE          	8	 

/*
 * Maximum request size (rounded up to the nearest page boundary) in
 * which the actual allocation will be twice as much.
 */
#define DOUBLE_THRESHOLD 	(1 << 20)

/*
 * Initialize the page table system.
 */
static int page_table_init (void);

/*
 * Level allocator.
 */
static int level_allocator (void* addr, int level, void** p, size_t size, size_t actual);

/*
 * Get the size information for a given address.
 */
static inline t_bottom_entry* get_size_data (void* addr, size_t level, void** p);

/*
 * Return the bit position of the most significant set bit.  Returns
 * (-1) if x is zero.
 */
static inline int max_set_bit (size_t x);

/*
 * Returns the index corresponding to 
 */
static inline long get_level_index (long level, void* ptr);

/*
 * Calculate page aligned size.
 */
static inline size_t page_align_size (size_t s);

/*
 * Implementation.....................................................
 */
void*
malloc_blk
(size_t size)
{
  void*  ret;
  int    res;
  size_t actual;
 
  if ((alloc_sizes == NULL) && (page_table_init () == 1))
    goto err;
  else
    {
      /*
       * allocate requested memory with mmap ().
       */
      actual = page_align_size (size);
      if (actual <= DOUBLE_THRESHOLD)
	actual <<= 1;
      ret = mmap (NULL, actual, PROT_READ | PROT_WRITE,
		  MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
      if (ret == MAP_FAILED)
	goto err;
      else
	{
	  if (level_allocator (ret, 0, alloc_sizes->top, size, actual) == 1)
	    goto err1;
	}
    }
 quit:
  return ret;
 err:
  ret = NULL;
  goto quit;
 err1:
  res = munmap (ret, size);
  assert (res > 0);
  goto err;
}
void*
realloc_blk
(void* ptr, size_t size)
{
  void*           ret;
  t_bottom_entry* be;
  size_t          actual;
  
  be = get_size_data (ptr, 0, alloc_sizes->top);
  actual = be->as;
  
  if (size < actual)
    {
      if (size <= (actual >> 2))
	{
	  actual >>= 1;
	  ret = mremap (ptr, be->as, actual, MREMAP_MAYMOVE);
	  if (ret == MAP_FAILED)
	    goto err;
	  if (ret == ptr)
	    {
	      be->as = actual;
	      be->us = size;
	    }
	  else
	    {
	      bzero (be, sizeof (t_bottom_entry));
	      if (level_allocator (ret, 0, alloc_sizes->top, size, actual) == 1)
		goto err1;
	    }
	}
      else
	{
	  be->us = size;
	  ret = ptr;
	}
    }
  else 
    {
      actual = page_align_size (size);
      if (actual <= DOUBLE_THRESHOLD)
	actual <<= 1;
      ret = mremap (ptr, be->as, actual, MREMAP_MAYMOVE);
      if (ret == MAP_FAILED)
	goto err;
      if (ret == ptr)
	{
	  be->as = actual;
	  be->us = size;
	}
      else
	{
	  bzero (be, sizeof (t_bottom_entry));
	  if (level_allocator (ret, 0, alloc_sizes->top, size, actual) == 1)
	    goto err1;
	}
    }
 quit:
  return ret;
 err:
  ret = NULL;
  goto quit;
 err1:
  munmap (ret, actual);
  goto err;
}
void
free_blk
(void* ptr)
{
  t_bottom_entry* sd;		/* size data */
 
  sd = get_size_data (ptr, 0, alloc_sizes->top);

  assert (munmap (ptr, sd->as) != -1);

  /*
   * Reset size data to zero.
   */
  bzero (sd, sizeof (t_bottom_entry));
}
static
int
page_table_init
(void)
{ /*----------------------------------------------------------------*/
  int 	  ret     = 0;	/* return value */
  void*   resptr;	/* pointer for accepting return value of
			   function.  */
  int 	  ps;		/* page size.  */
  int 	  bips;		/* bits in page size */
  int 	  bias;		/* bits in address size */
  int 	  biti;		/* bits in table index, this contains the
	  		   number of bits in the index taking the
	  		   indices of all levels as one index, since
	  		   all of these bits are consecutive.  */
  int 	  levels;	/* number of paging levels */
  int 	  remainder;	/* remainder variable */
  ssize_t top_entries;	/* number of entries in the top level page
			   table.  */
  size_t  table_size;	/* The size of the top level page level */
  /* ---------------------------------------------------------------*/
  /*
   * Set up the alloc_sizes pointer.
   */
  alloc_sizes = &init_sizes;

  /*
   * Get the page size reported by the operating system.
   */
  ps = sysconf (_SC_PAGE_SIZE);
  if (ps == -1)
    goto err;
  else
    {
      bips = max_set_bit (ps);
      assert (bips >= PAGE_TABLE_SIZE_BITS);	
      if (bips < PAGE_TABLE_SIZE_BITS)
	/*
         * AN EXTREMELY UNLIKELY CONDITION IN A WELL FUNCTIONING UNIX
         * ENVIRONMENT.  This code is only reached if the macro NDEBUG
         * is defined.
	 */
	{
	  errno = 0;                  
	  goto err;
	}
      else
	{
	  bias      = BITS_IN_BYTE * sizeof (void*);
	  biti      = bias - bips;
	  levels    = biti / PAGE_TABLE_SIZE_BITS;
	  remainder = biti % PAGE_TABLE_SIZE_BITS;
	  if (remainder != 0)
	    {
	      /*
	       * The top level has fewer entries than
	       * MAX_ENTRIES_PER_TABLE, and these remaining entries
	       * add one incomplete level.
	       */
	      ++levels;
	      top_entries = (1 << remainder);
	    }
	  else
	    top_entries = MAX_ENTRIES_PER_TABLE;
	    
	  /*
	   * Allocate entries for top level page table.
	   */
	  table_size = top_entries * sizeof (void*);
	  resptr = sbrk ((ptrdiff_t) table_size);
	  if (resptr == (void *) -1)
	    goto err;
	  else
	    {
	      /*
	       * Finally, save the result data in module wide storage,
	       * and then call the level allocation function.
	       */
	      alloc_sizes->top = resptr;
	      alloc_sizes->ps  = ps;
	      alloc_sizes->lev = levels;
	    }
	}
    }
 quit:
  return ret;
 err:
  ret         = 1;
  alloc_sizes = NULL;
  goto quit;
}
static
int
level_allocator
( /* PARAMETERS: ---------------------------------------------------*/
 void*  addr,	/* address added to page tables */
 int    level,	/* level passed from the previous recursion */
 void** p,	/* Pointer pass from the previous recursion */
 size_t size,	/* the length requested */
 size_t actual	/* the actual length allocated. */
 /* ----------------------------------------------------------------*/
  )
{
  int  	 nlev;	/* number or levels */
  void*  table;	/* pointer to new table */
  size_t index;	/* index */
  /*
   * Get the number of levels from module wide storage.
   */
  nlev = alloc_sizes->lev;
  index = get_level_index (level, addr);
  if (level == nlev - 1)
    {
      /*
       * we have completed processing all levels.
       */
      t_bottom_entry* bptr;		/* bottom pointer */
      
      bptr           = (t_bottom_entry*) p;
      bptr[index].us = size;
      bptr[index].as = actual;
      return 0;
    }
  else
    /*
     * We have not completed processing all levels.
     */
    {
      size_t s;
      void** ptr; 		/* pointer variable */
      
      ptr   = p + index;
      if (*ptr == NULL)
	{
	  s     = ((level < (nlev - 2))
		   ? sizeof (void*) : sizeof (t_bottom_entry));
	  table = sbrk ((intptr_t) s * MAX_ENTRIES_PER_TABLE);
	  if (table == (void*) -1)
	    return 1;
	  *ptr  = table;
	}
      return level_allocator (addr, level + 1, *ptr, size, actual);
    }
}
static inline
long
get_level_index
(long level, void* ptr)
{ /*----------------------------------------------------------------*/
  long            ret; /* return value */
  int             ps;  /* page size */
  int             lvl; /* number of levels */
  union u_intaddr ia;  /* union for integer/pointer treatment */
  /* ---------------------------------------------------------------*/
  /*
   * Get page size and levels from module wide storage.
   */
  ps  = alloc_sizes->ps;
  lvl = alloc_sizes->lev;
  
  /*
   * Shift out the offset portion of the address.
   */
  ia.p =   ptr;
  ret  =   ia.l;
  ret  >>= max_set_bit (ps);

  /*
   * Determine how many bit positions to shift, and shen shift that
   * many spaces.  Recall that we count levels starting from zero.
   */
  ret  >>= (lvl - 1 - level) * PAGE_TABLE_SIZE_BITS;

  /*
   * Mask out bits that cannot be part of a level index.
   */
  ret  &=  (1 << PAGE_TABLE_SIZE_BITS) - 1;

  /*
   * Return the calculated index to the calling function.
   */
  return ret;
}
static inline
t_bottom_entry*
get_size_data
(void* addr, size_t level, void** p)
{
  void**  ptr;
  t_bottom_entry* be;

  if (level == alloc_sizes->lev - 1)
    {
      be = (t_bottom_entry*) p;
      be += get_level_index (level, addr);
      return be;
    }
  else
    {
      ptr = p + get_level_index (level, addr);
      return get_size_data (addr, level + 1, *ptr);
    }
}
static inline
int
max_set_bit
(size_t x)
{
  int c;

  for (c = 0; x != 0; ++c)
    x >>= 1;
  --c;
  return c;	/* c == -1 means that x was zero. */
}
static inline
size_t
page_align_size
(size_t s)
{
  size_t ret;
  
  ret =  s / alloc_sizes->ps;
  if ((s % alloc_sizes->ps) != 0)
    ++ret;
  ret *= alloc_sizes->ps;
  return ret;
}
                                                                                                    
