A memory management module for C pointers

memory.c

/*
 The MIT License (MIT)

 Copyright (c) 2015 Tom Hancocks

 Permission is hereby granted, free of charge, to any person obtaining a copy of
 this software and associated documentation files (the "Software"), to deal in
 the Software without restriction, including without limitation the rights to
 use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 the Software, and to permit persons to whom the Software is furnished to do so,
 subject to the following conditions:

 The above copyright notice and this permission notice shall be included in all
 copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#include <stdlib.h>
#include <assert.h>
#include <stdio.h>

#include "memory.h"

/*
 The following helper macros will mess up the code internal to the memory
 manager. Get rid of them within the scope of this file.
 */
#undef malloc
#undef free

/*
 This is a custom C based memory management module. It is intended to provide
 reference count based memory management. It should be noted that this does not
 improve performance, and should not be used in performance critical
 applications. It is intended to make memory management in less critical
 applications slightly more forgiving.
 */

/*
 Keep the stack responsible for tracking scopes in its own struct.
 */
struct
{
  struct __mm_scope scope;
  struct __mm_scope *top;
} __scope_stack;


/*
 Sets up and configures the root node of the stack for memory management.
 */
void init_mm()
{
  __scope_stack.scope.parent = NULL;
  __scope_stack.scope.first = NULL;
  __scope_stack.scope.last = NULL;
  __scope_stack.scope.size = 0;
  __scope_stack.scope.warn_dirty = NULL;
  __scope_stack.scope.name = "[spawn]";
  __scope_stack.top = &__scope_stack.scope;
};

/*
 Scope related helper functions. These should not be invoked directly from
 outside this file!
 */
void __mm_scope_dirty_warning(struct __mm_scope *scope);
void __mm_scope_do_autorelease(struct __mm_scope *scope);

/*
 Create a new memory management scope.
 */
void push_mm_scope()
{
  static unsigned int scope_count = 0;
  char scope_name[32];
  sprintf(scope_name, "untitled-scope-%d", scope_count);
  push_named_mm_scope((const char *)scope_name);
};

void push_named_mm_scope(const char *name)
{
  struct __mm_scope *scope = malloc(sizeof(*scope));
  scope->parent = __scope_stack.top;
  scope->name = name;
  scope->first = NULL;
  scope->last = NULL;
  scope->warn_dirty = __mm_scope_dirty_warning;
  scope->do_autorelease = __mm_scope_do_autorelease;
  scope->size = 0;
  __scope_stack.top = scope;
};

/*
 Leave the current scope.
 */
void pop_mm_scope()
{
  struct __mm_scope *scope = __scope_stack.top;
  assert(scope->parent);
  assert(scope->warn_dirty);
  assert(scope->do_autorelease);

  scope->do_autorelease(scope);

  if (scope->size > 0) {
    scope->warn_dirty(scope);
  }

  __scope_stack.top = scope->parent;
  free(scope);
};

/*
 Returns a reference to the current scope.
 */
struct __mm_scope *mm_current_scope()
{
  return __scope_stack.top;
}

/*
 Search for an item in the current scope. Returns the reference header if it
 is found. Returns NULL if it is not currently being tracked.
 */
struct __mm_ref_header *__mm_find_ref_header(
  struct __mm_scope *scope,
  void *ptr
) {
  struct __mm_ref_header *ref = scope->first;
  while (ref) {
    if (ref->ref == ptr) {
      return ref;
    }
    ref = ref->next;
  }
  return NULL;
};

/*
 Add a header for a referenced item into the current scope.
 */
struct __mm_ref_header *__mm_add_ref_header(struct __mm_scope *scope, void *ptr)
{
  assert(ptr);
  assert(scope);

  struct __mm_ref_header *ref = __mm_find_ref_header(scope, ptr);
  if (ref) {
    ref->rc++;
    return ref;
  }

  ref = malloc(sizeof(*ref));
  ref->ref = ptr;
  ref->rc = 1;
  ref->dealloc = NULL;
  ref->retain = NULL;
  ref->release = NULL;
  ref->free = NULL;
  ref->autorelease = NULL;
  ref->scope = scope;
  ref->next = NULL;
  ref->prev = scope->last;
  ref->allocated = 1;

  if (scope->last) {
    scope->last->next = ref;
  }

  if (!scope->first) {
    scope->first = ref;
  }

  scope->last = ref;
  scope->size++;

  return ref;
};

/*
 Remove a header for a referenced item from the current stack.
 */
void __mm_remove_ref_header(struct __mm_scope *scope, void *ptr)
{
  assert(scope);
  assert(ptr);

  struct __mm_ref_header *ref = __mm_find_ref_header(scope, ptr);
  if (!ref || ref->rc > 0) {
    return;
  }

  struct __mm_ref_header *prev = ref->prev;
  struct __mm_ref_header *next = ref->next;
  scope->size--;

  if (prev) {
    prev->next = next;
  }

  if (next) {
    next->prev = prev;
  }

  !ref->dealloc ?: ref->dealloc(ref);

  free(ref);
};

/*
 Zero out the contents of a pointer.
 */
void __mm_zero_pointer(void *ptr)
{
  *((int *)ptr) = 0;
};

/*
 Memory Management Functions
 */
void *mm_release(void *ptr)
{
  if (!ptr) {
    return NULL;
  }

  struct __mm_ref_header *ref = __mm_find_ref_header(mm_current_scope(), ptr);

  if (!ref) {
    return ptr;
  }

  assert(ref->rc > 0);
  ref->rc--;

  if (ref->rc == 0) {
    !ref->release ?: ref->release(ref);
    char f = ref->allocated;
    __mm_remove_ref_header(mm_current_scope(), ptr);
    if (f) {
      free(ptr);
    }
    __mm_zero_pointer(ptr);
    ptr = NULL;
  }

  return ptr;
};

void *mm_retain(void *ptr)
{
  if (!ptr) {
    return NULL;
  }

  struct __mm_ref_header *ref = __mm_find_ref_header(mm_current_scope(), ptr);

  if (!ref) {
    ref = __mm_add_ref_header(mm_current_scope(), ptr);
    ref->allocated = 0;
    return ptr;
  }

  assert(ref->rc > 0);
  ref->rc++;
  !ref->retain ?: ref->retain(ref);

  return ptr;
};

/*
 Custom malloc function. This is here to allow seemless integration to the
 memory manager.
 */
void *__mm_malloc(size_t n)
{
  void *p = malloc(n);
  __mm_add_ref_header(mm_current_scope(), p);
  return p;
};

/*
 Custom free function. As above.
 */
void *__mm_free(void *ptr)
{
  return mm_release(ptr);
};

/*
 Scope related messages
 */
void __mm_scope_dirty_warning(struct __mm_scope *scope)
{
  fprintf(stderr,
          "Scope [%s] still contains allocated memory!\n",
          scope->name);
};

/*
 Autorelease pool
 */
void __mm_scope_do_autorelease(struct __mm_scope *scope)
{
  
};

memory.h

/*
The MIT License (MIT)

Copyright (c) 2015 Tom Hancocks

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */

#ifndef nebula_test_memory_h
#define nebula_test_memory_h

#include <stdlib.h>
#include <assert.h>
#include <stdio.h>

/*
 This is a custom C based memory management module. It is intended to provide
 reference count based memory management. It should be noted that this does not
 improve performance, and should not be used in performance critical 
 applications. It is intended to make memory management in less critical
 applications slightly more forgiving.
 */

struct __mm_ref_header;
struct __mm_scope;

typedef void(*ref_handler)(struct __mm_ref_header *);
typedef void(*scope_warning)(struct __mm_scope *);
typedef void(*scope_action)(struct __mm_scope *);

struct __mm_ref_header
{
  // Reference
  void *ref;
  unsigned int rc;
  char allocated;

  // Event handling functions
  ref_handler dealloc;
  ref_handler retain;
  ref_handler release;
  ref_handler autorelease;
  ref_handler free;

  // List
  struct __mm_ref_header *next;
  struct __mm_ref_header *prev;

  // Scope
  struct __mm_scope *scope;
};

struct __mm_scope
{
  struct __mm_scope *parent;
  struct __mm_ref_header *first;
  struct __mm_ref_header *last;
  size_t size;
  scope_warning warn_dirty;
  scope_action do_autorelease;
  const char *name;
};

/*
 Sets up and configures the root node of the stack for memory management.
 */
void init_mm();

/*
 Create a new memory management scope.
 */
void push_mm_scope();
void push_named_mm_scope(const char *name);

/*
 Leave the current scope.
 */
void pop_mm_scope();

/*
 Release the memory referenced by the specifified pointer.
 */
void *mm_release(void *ptr);

/*
 Retain the memory referenced by the specified pointer.
 */
void *mm_retain(void *ptr);


/*
 The following macros allow for integration of these functions and features into
 external code.
 */
#if !defined(malloc)
# define malloc(n)  __mm_malloc((n))
#endif
#
#if !defined(free)
# define free(ptr)  __mm_free((ptr))
#endif

#endif