fix km_set_grow()

1 files changed, 21 insertions, 12 deletions

diff --git a/km-set.c b/km-set.c
index 63ec22a..58792c0 100644
--- a/km-set.c
+++ b/km-set.c
@@ -2,6 +2,7 @@
 #include <stdint.h> // size_t
 #include <string.h> // memcpy()
 #include <stdlib.h> // rand()
+#include <errno.h> // errno
 #include "util.h"
 #include "km.h"
 
@@ -13,20 +14,28 @@ km_set_grow(
   km_set_t * const set,
   const size_t capacity
 ) {
-  // alloc floats
-  const size_t num_floats = set->shape.num_floats * capacity;
-  float * const floats = malloc(sizeof(float) * num_floats);
-  if (!floats) {
-    // return failure
-    return false;
+  float *floats = NULL;
+  const size_t floats_size = sizeof(float) * set->shape.num_floats * capacity;
+  // fprintf(stderr, "floats_size = %zu\n", floats_size);
+  if (floats_size > 0) {
+    // alloc floats
+    floats = realloc(set->floats, floats_size);
+    if (!floats) {
+      // return failure
+      return false;
+    }
   }
 
-  // alloc ints
-  const size_t num_ints = set->shape.num_ints * capacity;
-  int * const ints = malloc(sizeof(int) * num_ints);
-  if (!ints) {
-    // return failure
-    return false;
+  int *ints = NULL;
+  const size_t ints_size = sizeof(int) * set->shape.num_ints * capacity;
+  // fprintf(stderr, "ints_size = %zu\n", ints_size);
+  if (ints_size > 0) {
+    // alloc ints
+    ints = realloc(set->ints, ints_size);
+    if (!ints) {
+      // return failure
+      return false;
+    }
   }
 
   // update set