Move capn_init_* all into capn-malloc.c with a common create and free function

capn-malloc.c

@@ -2,6 +2,7 @@
 #include "capn.h"
 #include <stdlib.h>
 #include <string.h>
+#include <limits.h>
 
 static struct capn_segment *create(void *u, uint32_t id, int sz) {
 	struct capn_segment *s;
@@ -13,6 +14,7 @@ static struct capn_segment *create(void *u, uint32_t id, int sz) {
 	s = (struct capn_segment*) calloc(1, sizeof(*s));
 	s->data = (char*) calloc(1, sz);
 	s->cap = sz;
+	s->user = (void*)(uintptr_t) 1;
 	return s;
 }
 
@@ -21,19 +23,130 @@ void capn_init_malloc(struct capn *c) {
 	c->create = &create;
 }
 
-void capn_free_malloc(struct capn *c) {
+void capn_free(struct capn *c) {
 	struct capn_segment *s = c->seglist;
 	while (s != NULL) {
 		struct capn_segment *n = s->next;
-		free(s->data);
+		if (s->user) {
-		free(s);
+			free(s->data);
+			free(s);
+		}
 		s = n;
 	}
 	s = c->copylist;
 	while (s != NULL) {
 		struct capn_segment *n = s->next;
-		free(s->data);
+		if (s->user) {
-		free(s);
+			free(s->data);
+			free(s);
+		}
 		s = n;
 	}
 }
+
+#define ZBUF_SZ 4096
+
+static int read_fp(void *p, size_t sz, FILE *f, struct capn_stream *z, uint8_t* zbuf, int packed) {
+	if (f && packed) {
+		z->next_out = (uint8_t*) p;
+		z->avail_out = sz;
+
+		while (z->avail_out && capn_inflate(z) == CAPN_NEED_MORE) {
+			int r;
+			memmove(zbuf, z->next_in, z->avail_in);
+			r = fread(zbuf+z->avail_in, 1, ZBUF_SZ - z->avail_in, f);
+			if (r <= 0)
+				return -1;
+			z->avail_in += r;
+		}
+		return 0;
+
+	} else if (f && !packed) {
+		return fread(p, sz, 1, f) != 1;
+
+	} else if (packed) {
+		z->next_out = (uint8_t*) p;
+		z->avail_out = sz;
+		return capn_inflate(z) != 0;
+
+	} else {
+		if (z->avail_in < sz)
+			return -1;
+		memcpy(p, z->next_in, sz);
+		z->next_in += sz;
+		z->avail_in -= sz;
+		return 0;
+	}
+}
+
+static int init_fp(struct capn *c, FILE *f, struct capn_stream *z, int packed) {
+	struct capn_segment *s = NULL;
+	uint32_t i, segnum, total = 0;
+	uint32_t hdr[1024];
+	uint8_t zbuf[ZBUF_SZ];
+	char *data = NULL;
+
+	capn_init_malloc(c);
+
+	if (read_fp(&segnum, 4, f, z, zbuf, packed))
+		goto err;
+
+	segnum = capn_flip32(segnum);
+	if (segnum > 1023)
+		goto err;
+	segnum++;
+
+	s = (struct capn_segment*) calloc(segnum, sizeof(*s));
+	if (!s)
+		goto err;
+
+	if (read_fp(hdr, 8 * (segnum/2) + 4, f, z, zbuf, packed))
+		goto err;
+
+	for (i = 0; i < segnum; i++) {
+		uint32_t n = capn_flip32(hdr[i]);
+		if (n > INT_MAX/8 || n > UINT32_MAX/8 || UINT32_MAX - total < n*8)
+			goto err;
+		s[i].cap = s[i].len = n * 8;
+		total += s[i].len;
+	}
+
+	data = (char*) calloc(1, total);
+	if (!data)
+		goto err;
+
+	if (read_fp(data, total, f, z, zbuf, packed))
+		goto err;
+
+	for (i = 0; i < segnum; i++) {
+		s[i].data = data;
+		data += s[i].len;
+		capn_append_segment(c, &s[i]);
+	}
+
+	/* mark the first segment so capn_free will free the data and
+	 * segment arrays */
+	s[0].user = (void*)(uintptr_t) 1;
+
+	return 0;
+
+err:
+	memset(c, 0, sizeof(*c));
+	free(data);
+	free(s);
+	return -1;
+}
+
+int capn_init_fp(struct capn *c, FILE *f, int packed) {
+	struct capn_stream z;
+	memset(&z, 0, sizeof(z));
+	return init_fp(c, f, &z, packed);
+}
+
+int capn_init_mem(struct capn *c, const uint8_t *p, size_t sz, int packed) {
+	struct capn_stream z;
+	memset(&z, 0, sizeof(z));
+	z.next_in = p;
+	z.avail_in = sz;
+	return init_fp(c, NULL, &z, packed);
+}

capn-stream-test.cpp

@@ -21,7 +21,7 @@ TEST(Stream, ReadEmptyStream_Even) {
   EXPECT_EQ(2, ctx.segnum);
   EXPECT_EQ(0, ctx.seglist->len);
   EXPECT_EQ(0, ctx.seglist->next->len);
-  capn_free_mem(&ctx);
+  capn_free(&ctx);
 }
 
 TEST(Stream, ReadEmptyStream_Odd) {
@@ -40,13 +40,13 @@ TEST(Stream, ReadEmptyStream_Odd) {
   EXPECT_EQ(3, ctx.segnum);
   EXPECT_EQ(0, ctx.seglist->len);
   EXPECT_EQ(0, ctx.seglist->next->len);
-  capn_free_mem(&ctx);
+  capn_free(&ctx);
 
   ASSERT_EQ(0, capn_init_mem(&ctx, data.bytes, 20, 0));
   EXPECT_EQ(3, ctx.segnum);
   EXPECT_EQ(0, ctx.seglist->len);
   EXPECT_EQ(0, ctx.seglist->next->len);
-  capn_free_mem(&ctx);
+  capn_free(&ctx);
 }
 
 TEST(Stream, ReadStream_Even) {
@@ -68,5 +68,5 @@ TEST(Stream, ReadStream_Even) {
   EXPECT_EQ(1, ctx.seglist->data[0]);
   EXPECT_EQ(16, ctx.seglist->next->len);
   EXPECT_EQ(9, ctx.seglist->next->data[0]);
-  capn_free_mem(&ctx);
+  capn_free(&ctx);
 }

capn-stream.c

@@ -1,9 +1,6 @@
 /* vim: set sw=8 ts=8 sts=8 noet: */
 #include "capn.h"
 #include <string.h>
-#include <stdio.h>
-#include <stdlib.h>
-#include <limits.h>
 
 #ifndef min
 static int min(int a, int b) { return (a < b) ? a : b; }
@@ -191,115 +188,3 @@ int capn_inflate(struct capn_stream* s) {
 	return 0;
 }
 
-#define ZBUF_SZ 4096
-
-static int read_fp(void *p, size_t sz, FILE *f, struct capn_stream *z, uint8_t* zbuf, int packed) {
-	if (f && packed) {
-		z->next_out = (uint8_t*) p;
-		z->avail_out = sz;
-
-		while (z->avail_out && capn_inflate(z) == CAPN_NEED_MORE) {
-			int r;
-			memmove(zbuf, z->next_in, z->avail_in);
-			r = fread(zbuf+z->avail_in, 1, ZBUF_SZ - z->avail_in, f);
-			if (r <= 0)
-				return -1;
-			z->avail_in += r;
-		}
-		return 0;
-
-	} else if (f && !packed) {
-		return fread(p, sz, 1, f) != 1;
-
-	} else if (packed) {
-		z->next_out = (uint8_t*) p;
-		z->avail_out = sz;
-		return capn_inflate(z) != 0;
-
-	} else {
-		if (z->avail_in < sz)
-			return -1;
-		memcpy(p, z->next_in, sz);
-		z->next_in += sz;
-		z->avail_in -= sz;
-		return 0;
-	}
-}
-
-static int init_fp(struct capn *c, FILE *f, struct capn_stream *z, int packed) {
-	struct capn_segment *s = NULL;
-	uint32_t i, segnum, total = 0;
-	uint32_t hdr[1024];
-	uint8_t zbuf[ZBUF_SZ];
-	char *data = NULL;
-	memset(c, 0, sizeof(*c));
-
-	if (read_fp(&segnum, 4, f, z, zbuf, packed))
-		goto err;
-
-	segnum = capn_flip32(segnum);
-	if (segnum > 1023)
-		goto err;
-	segnum++;
-
-	s = (struct capn_segment*) calloc(segnum, sizeof(*s));
-	if (!s)
-		goto err;
-
-	if (read_fp(hdr, 8 * (segnum/2) + 4, f, z, zbuf, packed))
-		goto err;
-
-	for (i = 0; i < segnum; i++) {
-		uint32_t n = capn_flip32(hdr[i]);
-		if (n > INT_MAX/8 || n > UINT32_MAX/8 || UINT32_MAX - total < n*8)
-			goto err;
-		s[i].cap = s[i].len = n * 8;
-		total += s[i].len;
-	}
-
-	data = (char*) calloc(1, total);
-	if (!data)
-		goto err;
-
-	if (read_fp(data, total, f, z, zbuf, packed))
-		goto err;
-
-	for (i = 0; i < segnum; i++) {
-		s[i].data = data;
-		data += s[i].len;
-		capn_append_segment(c, &s[i]);
-	}
-
-	return 0;
-
-err:
-	memset(c, 0, sizeof(*c));
-	free(data);
-	free(s);
-	return -1;
-}
-
-int capn_init_fp(struct capn *c, FILE *f, int packed) {
-	struct capn_stream z;
-	memset(&z, 0, sizeof(z));
-	return init_fp(c, f, &z, packed);
-}
-
-int capn_init_mem(struct capn *c, const uint8_t *p, size_t sz, int packed) {
-	struct capn_stream z;
-	memset(&z, 0, sizeof(z));
-	z.next_in = p;
-	z.avail_in = sz;
-	return init_fp(c, NULL, &z, packed);
-}
-
-void capn_free_mem(struct capn *c) {
-	capn_free_fp(c);
-}
-
-void capn_free_fp(struct capn *c) {
-	if (c->seglist) {
-		free(c->seglist->data);
-		free(c->seglist);
-	}
-}

capn-test.cpp

@@ -15,7 +15,7 @@ union AlignedData {
 class Session {
 public:
   Session() {capn_init_malloc(&capn);}
-  ~Session() {capn_free_malloc(&capn);}
+  ~Session() {capn_free(&capn);}
   struct capn capn;
 };
 

capn.h

@@ -73,6 +73,7 @@ struct capn_segment {
 	uint32_t id;
 	char *data;
 	int len, cap;
+	void *user;
 };
 
 enum CAPN_TYPE {
@@ -203,17 +204,18 @@ struct capn_tree *capn_tree_insert(struct capn_tree *root, struct capn_tree *n);
 /* capn_init_malloc inits the capn struct with a create function which
  * allocates segments on the heap using malloc
  *
- * capn_free_all frees all the segment headers and data created by the create
+ * capn_init_(fp|mem) inits by reading segments in from the file/memory buffer
- * function setup by capn_init_malloc
+ * in serialized form (optionally packed). It will then setup the create
+ * function ala capn_init_malloc so that further segments can be created.
+ *
+ * capn_free frees all the segment headers and data created by the create
+ * function setup by capn_init_*
  */
 void capn_init_malloc(struct capn *c);
-void capn_free_malloc(struct capn *c);
-
 int capn_init_fp(struct capn *c, FILE *f, int packed);
-void capn_free_fp(struct capn *c);
-
 int capn_init_mem(struct capn *c, const uint8_t *p, size_t sz, int packed);
-void capn_free_mem(struct capn *c);
+
+void capn_free(struct capn *c);
 
 /* capn_stream encapsulates the needed fields for capn_(deflate|inflate) in a
  * similar manner to z_stream from zlib