2 files changed, 46 insertions, 8 deletions

diff --git a/src/til_setup.c b/src/til_setup.c
index aab196b..d819173 100644
--- a/src/til_setup.c
+++ b/src/til_setup.c
@@ -7,7 +7,8 @@
 
 /* Allocate and initialize a new til_setup_t of size bytes.
  * free_func is assigned to til_setup_t.free, and will be used for freeing the
- * instance returned when destroyed.
+ * instance returned when destroyed.  If free_func is NULL, free() will be
+ * used by default.
  *
  * Note this returns void * despite creating a til_setup_t, this is for convenience
  * as the callers are generally using it in place of calloc(), and assign it to a
@@ -18,28 +19,63 @@ void * til_setup_new(size_t size, void (*free_func)(til_setup_t *setup))
 	til_setup_t	*setup;
 
 	assert(size >= sizeof(til_setup_t));
-	assert(free_func);
 
 	setup = calloc(1, size);
 	if (!setup)
 		return NULL;
 
+	setup->refcount = 1;
 	setup->free = free_func;
 
 	return setup;
 }
 
 
-/* Free the setup when non-NULL, using setup->free if non-NULL.
- * Always returns NULL for uses like foo = til_setup_free(foo);
+/* bump refcount on setup */
+void * til_setup_ref(til_setup_t *setup)
+{
+	assert(setup);
+
+	setup->refcount++;
+
+	return setup;
+}
+
+
+/* unref setup, freeing it when refcount reaches zero.
+ * returns NULL if setup is freed (including when NULL was supplied for setup)
+ * resturns setup when setup persists.
+ * the public api is to just use til_setup_free() and discard that information,
+ * but this is kept here as distinct for potential debugging purposes.
  */
-void * til_setup_free(til_setup_t *setup)
+static void * til_setup_unref(til_setup_t *setup)
 {
 	if (!setup)
 		return NULL;
 
-	if (setup->free)
-		setup->free(setup);
+	assert(setup->refcount > 0);
+
+	setup->refcount--;
+	if (!setup->refcount) {
+		if (setup->free)
+			setup->free(setup);
+		else
+			free(setup);
+
+		return NULL;
+	}
+
+	return setup;
+}
+
+
+/* like til_setup_unref() except always returns NULL so you
+ * can't tell if it was actually freed or not, but this is sometimes
+ * a convenient free-style wrapper if you have to NULL-assign a placeholder.
+ */
+void * til_setup_free(til_setup_t *setup)
+{
+	(void) til_setup_unref(setup);
 
 	return NULL;
 }
diff --git a/src/til_setup.h b/src/til_setup.h
index 16380d5..8cc7a88 100644
--- a/src/til_setup.h
+++ b/src/til_setup.h
@@ -4,10 +4,12 @@
 typedef struct til_setup_t til_setup_t;
 
 struct til_setup_t {
-	void (*free)(til_setup_t *setup);
+	unsigned	refcount;
+	void		(*free)(til_setup_t *setup);
 };
 
 void * til_setup_new(size_t size, void (*free_func)(til_setup_t *setup));
+void * til_setup_ref(til_setup_t *setup);
 void * til_setup_free(til_setup_t *setup);
 
 #endif