/**
   * Delete all lines containing
   * 'ALL_IN_ONE' if you split this example
   * into separate files.
   */
  #define ALL_IN_ONE 1


  /****
   * pfxalloc.h
   *
   * Shao Miller, 2011
   * It took several minutes to find
   * a free *alloc. (No pun intended.)
   */
  #ifndef PFXALLOC_H_

  /* For 'offsetof' and 'size_t' */
  #include <stddef.h>

  /*** Macros */
  #define PFXALLOC_H_

  #ifndef alignof

  /* Derived from Mr. Chris M. Thomasson */
  #define alignof(type) \
    (offsetof(struct {char c; type t;}, t))

  #ifdef __alignof_is_defined
  #undef __alignof_is_defined
  #endif /* __alignof_is_defined */

  #define __alignof_is_defined 1
  #endif /* alignof */

  /*** Functions */

  /**
   * Allocate memory for an object and an
   * associated "prefix" object.
   *
   * @v pfx_size         The "prefix" size.
   * @v obj_alignment    The object's alignment
   *                     requirement.
   * @v obj_size         The object's size.
   * @ret void *         Points to the object,
   *                     or is NULL for failure.
   */
  extern void * pfxalloc(size_t, size_t, size_t);

  /**
   * Get an object's associated "prefix" object.
   *
   * @v obj              The object to fetch the
   *                     associated "prefix" for.
   * @ret void *         Points to the "prefix",
   *                     or is NULL for failure.
   */
  extern void * pfxget(void *);

  #endif /* PFXALLOC_H_ */


  /**** pfxalloc.c */
  /* Shao Miller, 2011 */

  /* For 'malloc' */
  #include <stdlib.h>
  #if !ALL_IN_ONE
  /* For 'size_t', 'ptrdiff_t' and 'NULL' */
  #include <stddef.h>
  /* For 'alignof' */
  #include "pfxalloc.h"
  #endif /* !ALL_IN_ONE */

  /*** Constants */
  enum cv {
      cv_ptrdiff_alignment =
        alignof(ptrdiff_t),
      cv_zero = 0
    };

  /*** Functions */

  /**
   * The layout looks like:
   * +-----+-------------+--------+--------+
   * | pfx | opt_padding | offset |        |
   * +-   -+-           -+-      -+ object |
   * |           header           |        |
   * +----------------------------+--------+
   */
  void * pfxalloc(
      size_t pfx_size,
      size_t obj_alignment,
      size_t obj_size
    ) {
      size_t offset_at;
      ptrdiff_t * offset;
      size_t header_size;
      size_t total_size;
      unsigned char * mem;
      unsigned char * obj;

      /* Sanity-check alignment value */
      if (obj_alignment & (obj_alignment - 1))
        return NULL;

      /* Calculate the offset position */
      offset_at = pfx_size;
      offset_at += sizeof *offset;
      /* Check for wrap-around */
      if (offset_at < pfx_size)
        return NULL;
      --offset_at;
      offset_at /= sizeof *offset;

      /* Calculate the header size */
      header_size = (offset_at + 1) * sizeof *offset;
      /* Check for wrap-around */
      if (header_size < pfx_size)
        return NULL;

      /* Calculate padding */
      if (obj_alignment > cv_ptrdiff_alignment) {
          size_t new_hdr_size = header_size;
          new_hdr_size += obj_alignment;
          /* Check for wrap-around */
          if (new_hdr_size < header_size)
            return NULL;
          --new_hdr_size;
          new_hdr_size /= obj_alignment;
          new_hdr_size *= obj_alignment;
          header_size = new_hdr_size;
        }

      /* Allocate storage */
      total_size = header_size + obj_size;
      /* Check for wrap-around */
      if (total_size < pfx_size || total_size < obj_size)
        return NULL;
      mem = malloc(total_size);
      if (!mem)
        return mem;

      /* Point to the object */
      obj = mem + header_size;

      /* Note the offset to the prefix */
      offset = (ptrdiff_t *) obj - 1;
      *offset = obj - mem;

      return obj;
    }

  /* Fetch an asociated prefix for an object */
  void * pfxget(void * obj) {
      ptrdiff_t * offset;
      unsigned char * prefix;

      if (!obj)
        return obj;

      offset = obj;
      --offset;
      prefix = obj;
      prefix -= *offset;

      return prefix;
    }


  /**** ref_cnt.h */
  /* Shao Miller, 2011 */
  #ifndef REF_CNT_H_

  #if !ALL_IN_ONE
  /* For 'size_t' */
  #include <stddef.h>
  #include "align.h"
  #endif /* !ALL_IN_ONE */

  /*** Macros */
  #define REF_CNT_H_

  /*** Function types */
  typedef void f_ref_cnt_free(void *);

  /*** Functions */

  /**
   * Create a new, reference-counted object
   *
   * @v size              The size of the object
   *                      to allocate, in bytes.
   * @v alignment         The alignment requirement
   *                      for the object.
   * @v dealloc           The de-allocation function
   *                      to call, when refs == 0.
   * @ret void *          Points to the newly-
   *                      allocated object, or is
   *                      NULL, upon failure.
   */
  extern void * NewRefCountedObj(
      size_t,
      size_t,
      f_ref_cnt_free *
    );

  /**
   * Increment the reference count for an object.
   *
   * @v obj               The object to reference.
   * @ret void *          Points to the object, or is
   *                      NULL if too many references.
   */
  extern void * GetRef(void *);

  /**
   * Decrement the reference count for an object.
   *
   * @v obj               The object to dereference.
   */
  extern void PutRef(void *);

  #endif /* REF_CNT_H_ */


  /**** ref_cnt.c */
  /* Shao Miller, 2011 */

  #if !ALL_IN_ONE
  /* For 'size_t' */
  #include <stddef.h>
  #include "pfxalloc.h"
  #include "ref_cnt.h"
  #endif /* !ALL_IN_ONE */

  /*** Types */
  struct ref_counted_obj {
      unsigned int ref_count;
      f_ref_cnt_free * dealloc;
    };

  /*** Functions */

  void * NewRefCountedObj(
      size_t size,
      size_t alignment,
      f_ref_cnt_free * dealloc
    ) {
      void * result;
      struct ref_counted_obj * ref;

      /* Allocate the object */
      result = pfxalloc(
          sizeof *ref,
          alignment,
          size
        );
      if (!result)
        return result;

      ref = pfxget(result);
      if (!ref) {
          /* Uh oh; this is unexpected! */
          return ref;
        }

      /* Populate accounting info */
      ref->ref_count = 1;
      ref->dealloc = dealloc;

      return result;
    }

  void * GetRef(void * obj) {
      struct ref_counted_obj * ref;

      ref = pfxget(obj);
      if (!ref)
        return ref;

      /* Increment ref count and check for wrap-around */
      if (!++ref->ref_count) {
          --ref->ref_count;
          return NULL;
        }
      return obj;
    }

  void PutRef(void * obj) {
      struct ref_counted_obj * ref;

      ref = pfxget(obj);
      if (!ref)
        return;
      /* Decrement ref count and check for zero */
      if (!--ref->ref_count) {
          ref->dealloc(obj);
          free(ref);
        }
      return;
    }


  /**** Test program */
  /* Shao Miller, 2011 */

  #include <stdio.h>
  #include <string.h>
  #if !ALL_IN_ONE
  #include <stdlib.h>
  #include "pfxalloc.h"
  #include "ref_cnt.h"
  #endif /* !ALL_IN_ONE */

  /*** Types */
  struct foo {
      double bobble;
      size_t len;
      char bear[1];
    };

  /*** Functions */
  f_ref_cnt_free foo_dealloc;
  void foo_dealloc(void * ptr) {
      struct foo * bar;

      bar = ptr;
      printf(
          "foo_dealloc(%p): {%f, %d, \"%s\"}\n",
          ptr,
          bar->bobble,
          (int)bar->len,
          bar->bear
        );
      return;
    }

  int main(void) {
      struct foo * bar;
      char laughsalot[] = "Laughs-a-lot";

      bar = NewRefCountedObj(
          sizeof *bar + sizeof laughsalot,
          alignof (struct foo),
          foo_dealloc
        );
      if (!bar) {
          puts("NewRefCountedObj() failed!");
          return EXIT_FAILURE;
        }

      bar->bobble = 3.14159;
      bar->len = sizeof laughsalot;
      memcpy(bar->bear, laughsalot, sizeof laughsalot);

      /* Test getting a reference */
      bar = GetRef(bar);
      if (!bar) {
          puts("GetRef() failed!");
          PutRef(bar);
          return EXIT_FAILURE;
        }

      /* We have two references */
      PutRef(bar);
      PutRef(bar);
      return EXIT_SUCCESS;
    }