Merge pull request #37 from detly/write_size

Implemented capn_size() for calculating buffer size

lib/capn-malloc.c

@@ -407,3 +407,29 @@ int capn_write_fd(struct capn *c, ssize_t (*write_fd)(int fd, const void *p, siz
 
 	return datasz;
 }
+
+int capn_size(struct capn *c)
+{
+	size_t headersz, datasz = 0;
+	struct capn_ptr root;
+	struct capn_segment *seg;
+	int i;
+
+	if (c->segnum == 0)
+		return -1;
+
+	root = capn_root(c);
+	seg = root.seg;
+
+    headersz = 8 * ((2 + c->segnum) / 2);
+
+	for (i = 0; i < c->segnum; i++, seg = seg->next) {
+		if (0 == seg)
+			return -1;
+		datasz += seg->len;
+	}
+	if (0 != seg)
+		return -1;
+
+	return (int) headersz+datasz;
+}

lib/capnp_c.h

@@ -276,6 +276,14 @@ void capn_init_malloc(struct capn *c);
 int capn_init_fp(struct capn *c, FILE *f, int packed);
 int capn_init_mem(struct capn *c, const uint8_t *p, size_t sz, int packed);
 
+/* capn_size() calculates the amount of memory required to serialise the given
+ * Cap'n Proto structure in the unpacked format. It does NOT apply to packed
+ * serialisation, as that may (in rare cases) actually become bigger than the
+ * input. A buffer of this size can then be passed to capn_write_mem() without
+ * fear of truncation (again, only in the unpacked case).
+ */
+int capn_size(struct capn *c);
+
 /* capn_write_(fp|mem) writes segments to the file/memory buffer in
  * serialized form and returns the number of bytes written.
  */

tests/capn-stream-test.cpp

@@ -80,6 +80,71 @@ TEST(Stream, ReadStream_Even) {
   capn_free(&ctx);
 }
 
+static struct capn_segment *CreateSmallSegment(void *u, uint32_t id, int sz) {
+  struct capn_segment *s = (struct capn_segment*) calloc(1, sizeof(*s));
+  s->data = (char*) calloc(1, sz);
+  s->cap = sz;
+  return s;
+}
+
+TEST(Stream, SizeEmptyStream) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  struct capn_ptr root = capn_root(&ctx);
+  ASSERT_EQ(CAPN_PTR_LIST, root.type);
+  EXPECT_EQ(2*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
+TEST(Stream, SizeOneSegment) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  struct capn_ptr root = capn_root(&ctx);
+  struct capn_ptr ptr = capn_new_struct(root.seg, 8, 0);
+  EXPECT_EQ(0, capn_setp(root, 0, ptr));
+  EXPECT_EQ(0, capn_write64(ptr, 0, UINT64_C(0x1011121314151617)));
+  EXPECT_EQ(3*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
+TEST(Stream, SizeTwoSegments) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  ctx.create = &CreateSmallSegment;
+  struct capn_ptr root = capn_root(&ctx);
+  struct capn_ptr ptr1 = capn_new_struct(root.seg, 8, 0);
+  EXPECT_EQ(0, capn_setp(root, 0, ptr1));
+  EXPECT_EQ(0, capn_write64(ptr1, 0, UINT64_C(0xfffefdfcfbfaf9f8)));
+  EXPECT_EQ(2, ctx.segnum);
+
+  /* 2 words: header
+   * 1 word: segment 1
+   * 2 words: segment 2
+   */
+  EXPECT_EQ(5*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
+TEST(Stream, SizeThreeSegments) {
+  struct capn ctx;
+  capn_init_malloc(&ctx);
+  ctx.create = &CreateSmallSegment;
+  struct capn_ptr root = capn_root(&ctx);
+  struct capn_ptr ptr1 = capn_new_struct(root.seg, 0, 1);
+  EXPECT_EQ(0, capn_setp(root, 0, ptr1));
+  struct capn_ptr ptr2 = capn_new_struct(ptr1.seg, 4, 0);
+  EXPECT_EQ(0, capn_setp(ptr1, 0, ptr2));
+  EXPECT_EQ(0, capn_write32(ptr2, 0, 0x12345678));
+  EXPECT_EQ(3, ctx.segnum);
+
+  EXPECT_EQ(7*8, capn_size(&ctx));
+
+  capn_free(&ctx);
+}
+
 TEST(Stream, WriteEmptyStream) {
   uint8_t buf[2048];
 
@@ -164,13 +229,6 @@ TEST(Stream, WriteOneSegmentPacked) {
   capn_free(&ctx2);
 }
 
-static struct capn_segment *CreateSmallSegment(void *u, uint32_t id, int sz) {
-  struct capn_segment *s = (struct capn_segment*) calloc(1, sizeof(*s));
-  s->data = (char*) calloc(1, sz);
-  s->cap = sz;
-  return s;
-}
-
 TEST(Stream, WriteTwoSegments) {
   struct capn ctx1, ctx2;
   uint8_t buf[5*8];