3 files changed, 110 insertions, 33 deletions

diff --git a/include/conts/map.h b/include/conts/map.h
index db1e386..e3f0b2f 100644
--- a/include/conts/map.h
+++ b/include/conts/map.h
@@ -18,6 +18,22 @@
 #error "Need map name"
 #endif
 
+#ifndef MAP_MALLOC
+#define MAP_MALLOC malloc
+#endif
+
+#ifndef MAP_CALLOC
+#define MAP_CALLOC calloc
+#endif
+
+#ifndef MAP_REALLOC
+#define MAP_REALLOC realloc
+#endif
+
+#ifndef MAP_FREE
+#define MAP_FREE free
+#endif
+
 #include <stddef.h>
 #include <stdlib.h>
 #include <stdbool.h>
@@ -101,9 +117,9 @@ static inline struct MAP_ROOT MAP(create)(size_t count)
 static inline void MAP(destroy)(struct MAP_ROOT *root)
 {
 	for (size_t i = 0; i < root->count; ++i)
-		free(root->buckets[i]);
+		MAP_FREE(root->buckets[i]);
 
-	free(root->buckets);
+	MAP_FREE(root->buckets);
 }
 
 static inline MAP_TYPE *MAP(insert)(struct MAP_ROOT *root, MAP_KEY key, MAP_TYPE data)
@@ -133,11 +149,20 @@ static inline MAP_TYPE *MAP(insert)(struct MAP_ROOT *root, MAP_KEY key, MAP_TYPE
 	/* no bucket available, create new one */
 	size_t size = root->pow2 << root->count;
 	size_t bytes = sizeof(struct MAP_BUCKET) + sizeof(struct MAP_NODE) * size;
-	struct MAP_BUCKET *bucket = calloc(1, bytes);
+	struct MAP_BUCKET *bucket = MAP_CALLOC(1, bytes);
+	if (!bucket)
+		return NULL;
+
 	bucket->size = size;
 
 	size_t buckets_bytes = sizeof(struct MAP_BUCKET *) * (root->count + 1);
-	root->buckets = realloc(root->buckets, buckets_bytes);
+	struct MAP_BUCKET **new_buckets = MAP_REALLOC(root->buckets, buckets_bytes);
+	if (!new_buckets) {
+		free(bucket);
+		return NULL;
+	}
+
+	root->buckets = new_buckets;
 	if (root->count != 0)
 		root->buckets[root->count - 1]->next = bucket;
 
@@ -248,3 +273,8 @@ static inline size_t MAP(len)(struct MAP_ROOT *root)
 #undef MAP_CMP
 #undef MAP_HASH
 #undef MAP_NAME
+
+#undef MAP_MALLOC
+#undef MAP_CALLOC
+#undef MAP_REALLOC
+#undef MAP_FREE
diff --git a/include/conts/sptree.h b/include/conts/sptree.h
index 8c74604..1464d0f 100644
--- a/include/conts/sptree.h
+++ b/include/conts/sptree.h
@@ -16,6 +16,22 @@
 #error "Need sptree name"
 #endif
 
+#ifndef SPTREE_MALLOC
+#define SPTREE_MALLOC malloc
+#endif
+
+#ifndef SPTREE_CALLOC
+#define SPTREE_CALLOC calloc
+#endif
+
+#ifndef SPTREE_REALLOC
+#define SPTREE_REALLOC realloc
+#endif
+
+#ifndef SPTREE_FREE
+#define SPTREE_FREE free
+#endif
+
 #include "conts.h"
 
 #define SPTREE(a) CONTS_JOIN(SPTREE_NAME, a)
@@ -220,7 +236,7 @@ static inline SPTREE_TYPE *SPTREE(insert)(struct SPROOT *s, SPTREE_TYPE data)
 {
 	if (!s->root) {
 		assert(s->n == 0);
-		struct SPNODE *new = malloc(sizeof(struct SPNODE));
+		struct SPNODE *new = SPTREE_MALLOC(sizeof(struct SPNODE));
 		if (!new)
 			return NULL;
 
@@ -255,7 +271,7 @@ static inline SPTREE_TYPE *SPTREE(insert)(struct SPROOT *s, SPTREE_TYPE data)
 		return &n->data;
 	}
 
-	struct SPNODE *new = malloc(sizeof(struct SPNODE));
+	struct SPNODE *new = SPTREE_MALLOC(sizeof(struct SPNODE));
 	if (!new)
 		return NULL;
 
@@ -403,7 +419,7 @@ static inline void SPTREE(free_found)(struct SPROOT *s, SPTREE_TYPE *found)
 {
 	(void)s; /* unused */
 	struct SPNODE *del = CONTAINER_OF(found, struct SPNODE, data);
-	free(del);
+	SPTREE_FREE(del);
 }
 
 static inline void SPTREE(remove)(struct SPROOT *s, SPTREE_TYPE data)
@@ -426,9 +442,11 @@ static inline void SPTREE(destroy)(struct SPROOT *s)
 	}
 }
 
-#undef SPTREE
-#undef SPNODE
-#undef SPROOT
 #undef SPTREE_TYPE
 #undef SPTREE_NAME
 #undef SPTREE_CMP
+
+#undef SPTREE_MALLOC
+#undef SPTREE_CALLOC
+#undef SPTREE_REALLOC
+#undef SPTREE_FREE
diff --git a/include/conts/vec.h b/include/conts/vec.h
index c08fd49..b8dfde9 100644
--- a/include/conts/vec.h
+++ b/include/conts/vec.h
@@ -6,7 +6,24 @@
 #error "Need vector name"
 #endif
 
+#ifndef VEC_MALLOC
+#define VEC_MALLOC malloc
+#endif
+
+#ifndef VEC_CALLOC
+#define VEC_CALLOC calloc
+#endif
+
+#ifndef VEC_REALLOC
+#define VEC_REALLOC realloc
+#endif
+
+#ifndef VEC_FREE
+#define VEC_FREE free
+#endif
+
 #include <stdbool.h>
+#include <stdint.h>
 #include <string.h>
 #include <stdlib.h>
 #include <assert.h>
@@ -26,21 +43,15 @@ typedef VEC_TYPE *VEC(iter);
 
 static inline struct VEC_STRUCT VEC(create)(size_t reserve)
 {
-	if (reserve == 0)
-		return (struct VEC_STRUCT) {.n = 0, .s = 0, .buf = NULL};
-
+	/* first alloc doubles size of s, so divide by two to get wanted
+	 * reservation size. Add 1 for odd cases. */
 	return (struct VEC_STRUCT) {
 		       .n = 0,
-		       .s = reserve,
-		       .buf = malloc(reserve * sizeof(VEC_TYPE)),
+		       .s = (reserve + 1) / 2,
+		       .buf = NULL,
 	};
 }
 
-static inline bool VEC(uninit)(struct VEC_STRUCT *v)
-{
-	return v->buf == NULL;
-}
-
 static inline size_t VEC(len)(struct VEC_STRUCT *v)
 {
 	return v->n;
@@ -65,15 +76,22 @@ static inline VEC_TYPE *VEC(pop)(struct VEC_STRUCT *v)
 	return &v->buf[v->n];
 }
 
-static inline void VEC(append)(struct VEC_STRUCT *v, VEC_TYPE n)
+static inline VEC_TYPE *VEC(append)(struct VEC_STRUCT *v, VEC_TYPE n)
 {
 	v->n++;
-	if (v->n >= v->s) {
+	if (v->n >= v->s || !v->buf) {
 		v->s = v->s == 0 ? 1 : 2 * v->s;
-		v->buf = realloc(v->buf, v->s * sizeof(VEC_TYPE));
+		VEC_TYPE *new_buf = VEC_REALLOC(v->buf, v->s * sizeof(VEC_TYPE));
+		if (!new_buf) {
+			v->n--;
+			return NULL;
+		}
+
+		v->buf = new_buf;
 	}
 
 	v->buf[v->n - 1] = n;
+	return &v->buf[v->n - 1];
 }
 
 static inline void VEC(reset)(struct VEC_STRUCT *v)
@@ -82,7 +100,7 @@ static inline void VEC(reset)(struct VEC_STRUCT *v)
 }
 
 static inline void VEC(destroy)(struct VEC_STRUCT *v) {
-	free(v->buf);
+	VEC_FREE(v->buf);
 }
 
 typedef int (*VEC(comp_t))(VEC_TYPE *a, VEC_TYPE *b);
@@ -91,19 +109,25 @@ static inline void VEC(sort)(struct VEC_STRUCT *v, VEC(comp_t) comp)
 	qsort(v->buf, v->n, sizeof(VEC_TYPE), (__compar_fn_t)comp);
 }
 
-static inline void VEC(reserve)(struct VEC_STRUCT *v, size_t n)
+static inline VEC_TYPE *VEC(reserve)(struct VEC_STRUCT *v, size_t n)
 {
-	if (v->n >= n)
-		return;
+	if (v->s >= n) {
+		v->n = n;
+		return v->buf;
+	}
 
-	v->n = n;
-	if (v->s >= v->n)
-		return;
+	size_t s = v->s;
+	while (s <= n)
+		s = s == 0 ? 1 : 2 * s;
 
-	while (v->s < v->n)
-		v->s = v->s == 0 ? 1 : 2 * v->s;
+	VEC_TYPE *new_buf = VEC_REALLOC(v->buf, s * sizeof(VEC_TYPE));
+	if (!new_buf)
+		return NULL;
 
-	v->buf = realloc(v->buf, v->s * sizeof(VEC_TYPE));
+	v->n = n;
+	v->s = s;
+	v->buf = new_buf;
+	return v->buf;
 }
 
 static inline void VEC(shrink)(struct VEC_STRUCT *v, size_t n)
@@ -139,3 +163,8 @@ static inline VEC_TYPE *VEC(next)(VEC_TYPE *i)
 #undef VEC_TYPE
 #undef VEC_NAME
 #undef VEC_STRUCT
+
+#undef VEC_MALLOC
+#undef VEC_CALLOC
+#undef VEC_REALLOC
+#undef VEC_FREE