Friday, April 27, 2012

Simple Memory Pool in C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>


#define MEMORY_POOL_ALLOCED_CONST   ((void *) 0xFFFFFFFFu)

typedef struct memory_pool_s
{
  void *pool;
  void *empty_blocks;
  size_t block_size;
  size_t count;
} __attribute__ ((__aligned__)) memory_pool_t;

memory_pool_t * memory_pool_create(size_t bs, size_t c);
void memory_pool_destroy(memory_pool_t *mp);
void memory_pool_clear(memory_pool_t *mp);
void memory_pool_dump(memory_pool_t *mp, void (* print_func) (void *value));
void * memory_pool_alloc(memory_pool_t *mp);
void memory_pool_free(memory_pool_t *mp, void *p);


memory_pool_t * memory_pool_create(size_t bs, size_t c)
{
  memory_pool_t *mp = malloc(sizeof(memory_pool_t));
  if (!mp)
    return NULL;

  mp->block_size = bs;
  mp->count = c;

  mp->pool = malloc((mp->block_size + sizeof(void *)) * mp->count);

  memory_pool_clear(mp);

  mp->empty_blocks = mp->pool;

  return mp;
}

void memory_pool_destroy(memory_pool_t *mp)
{
  if (!mp)
    return;

  memory_pool_clear(mp);
  free(mp->pool);
  free(mp);
}


void memory_pool_clear(memory_pool_t *mp)
{
  if (!mp)
    return;

  size_t i;
  void **p;

  for (i = 0; i < mp->count - 1; i++)
  {
    p = (void **) ((uint8_t *) mp->pool + (mp->block_size * (i + 1) +
                   sizeof(void *) * i));
    *p = (uint8_t *) mp->pool + (mp->block_size + sizeof(void *)) * (i + 1);
  }

  p = (void **) ((uint8_t *) mp->pool + (mp->block_size * mp->count +
                 sizeof(void *) * (mp->count - 1)));
  *p = NULL;

  mp->empty_blocks = mp->pool;
}

void memory_pool_dump(memory_pool_t *mp, void (* print_func) (void *value))
{
  printf("start: %p, size: %d, count: %d\n", mp->pool,
         (mp->block_size + sizeof(void *)) * mp->count, mp->count);

  void *block;
  void **next;
  size_t i;

  for (i = 0; i < mp->count; i++)
  {
    block = (void *) ((uint8_t *) mp->pool + (mp->block_size * i) +
                      sizeof(void *) * i);
    next = (void **) ((uint8_t *) block + sizeof(void *));

    printf("block #%i(%p):", i, block);

    if (*next == MEMORY_POOL_ALLOCED_CONST)
    {
      printf(" allocated");

      if (print_func)
      {
        printf(", value: ");
        print_func(block);
      }

      printf("\n");
    } else
    {
      printf(" free, next address %p\n", *next);
    }
  }
}

void * memory_pool_alloc(memory_pool_t *mp)
{
  void *p;

  if (mp->empty_blocks)
  {
    p = mp->empty_blocks;
    mp->empty_blocks = * (void **) ((uint8_t *) mp->empty_blocks +
                                    mp->block_size);
    *(void **) ((uint8_t *) p + mp->block_size) = MEMORY_POOL_ALLOCED_CONST;
    return p;
  } else
  {
    return NULL;
  }
}

void memory_pool_free(memory_pool_t *mp, void *p)
{
  if (p && (p >= mp->pool) && (p <= (void *) ((uint8_t *) mp->pool +
      (mp->block_size + sizeof(void *)) * mp->count)))
  {
    *(void **) ((uint8_t *) p + mp->block_size) = mp->empty_blocks;
    mp->empty_blocks = p;
  }
}


void print_intptr(void *p)
{
  printf("%d", *(int *) p);
}

int main(int argc, char *argv[])
{
  int *value;
  int c = 0;

  memory_pool_t *mp = memory_pool_create(sizeof(int), 128);

  memory_pool_dump(mp, print_intptr);

  while ((value = memory_pool_alloc(mp)))
  {
    *value = c++;
  }

  memory_pool_dump(mp, print_intptr);

  memory_pool_destroy(mp);

  return EXIT_SUCCESS;
}