Fixed Memory Leak

5 files changed, 84 insertions, 25 deletions

diff --git a/include/lib/mem/galloc.h b/include/lib/mem/galloc.h
index 8c773cd..763b21a 100644
--- a/include/lib/mem/galloc.h
+++ b/include/lib/mem/galloc.h
@@ -3,9 +3,14 @@
 
 #include <sys/types.h>
 
-void *g_malloc(size_t size);
-void *g_calloc(size_t n_memb, size_t size);
-void g_free(void *);
+#define g_malloc(size) _g_malloc(size, __FILE__, __LINE__)
+#define g_calloc(n_memb, size) _g_calloc(n_memb, size, __FILE__, __LINE__)
+#define g_realloc(ptr, size) _g_realloc(ptr, size, __FILE__, __LINE__)
+#define g_free(ptr) _g_free(ptr, __FILE__, __LINE__)
+void *_g_malloc(size_t size, char* file, unsigned int line);
+void *_g_calloc(size_t n_memb, size_t size, char* file, unsigned int line);
+void *_g_realloc(void *ptr, size_t size, char* file, unsigned int line);
+void _g_free(void *ptr, char* file, unsigned int line);
 ssize_t g_outstanding_allocations();
 
 #endif
diff --git a/lib/data/image_types.c b/lib/data/image_types.c
index f5c9115..6e8f6b4 100644
--- a/lib/data/image_types.c
+++ b/lib/data/image_types.c
@@ -68,29 +68,45 @@ Mask *create_image_mask(size_t width, size_t height) {
 }
 
 void free_image(Image *image) {
-  if (image->image[0] != NULL) {
-    g_free(image->image[0]);
-    image->image[0] = NULL;
+  if (image == NULL) {
+    return;
   }
-  if (image->image != NULL) {
-    g_free(image->image);
-    image->image = NULL;
+  if (image->image == NULL) {
+    g_free(image);
+    return;
   }
-  if (image != NULL) {
+  if (image->image[0] != NULL) {
+    if (image->image[0][0] != NULL) {
+      g_free(image->image[0][0]);
+    }
+  } else {
+    g_free(image->image);
     g_free(image);
+    return;
+  }
+  for (size_t y = 0; y < image->height; y++) {
+    if (image->image[y] != NULL) {
+      g_free(image->image[y]);
+    }
   }
+  g_free(image->image);
+  image->image = NULL;
+  g_free(image);
 }
 
 void free_image_mask(Mask *image_mask) {
+  if (image_mask == NULL) {
+    return;
+  }
+  if (image_mask->image == NULL) {
+    g_free(image_mask);
+    return;
+  }
   if (image_mask->image[0] != NULL) {
     g_free(image_mask->image[0]);
     image_mask->image[0] = NULL;
   }
-  if (image_mask->image != NULL) {
-    g_free(image_mask->image);
-    image_mask->image = NULL;
-  }
-  if (image_mask != NULL) {
-    g_free(image_mask);
-  }
+  g_free(image_mask->image);
+  image_mask->image = NULL;
+  g_free(image_mask);
 }
diff --git a/lib/dir.c b/lib/dir.c
index 361775d..f0776ca 100644
--- a/lib/dir.c
+++ b/lib/dir.c
@@ -22,7 +22,7 @@ char **list_directory(char *dirname) {
   struct dirent *dir;
   d = opendir(dirname);
   char **file_names = (char **)g_malloc(sizeof(char *));
-  if (!file_names) {
+  if (file_names == NULL) {
     return NULL;
   }
   file_names[0] = NULL;
@@ -32,10 +32,24 @@ char **list_directory(char *dirname) {
       if (dir->d_type == DT_REG) {
         // When a regular file is reached
         /// Create space for it in the list
-        file_names = realloc(file_names, (file_count + 1 + 1) * sizeof(char *));
+        char **temp =
+            g_realloc(file_names, (file_count + 1 + 1) * sizeof(char *));
+        if (temp == NULL) {
+          for (size_t file_idx = 0; file_idx < file_count; file_idx++) {
+	    g_free(file_names[file_idx]);
+          }
+	  return NULL;
+        }
+        file_names = temp;
         /// Create space for the name
         file_names[file_count] =
             g_calloc(strlen(dir->d_name) + 1, sizeof(char));
+        if (file_names[file_count] == NULL) {
+          for (size_t file_idx = 0; file_idx < file_count; file_idx++) {
+	    g_free(file_names[file_idx]);
+          }
+	  return NULL;
+        }
         /// Copy the name
         strcpy(file_names[file_count], dir->d_name);
         /// Mark the end of the file list with a NULL entry
@@ -53,6 +67,9 @@ char *full_path(char *dir, char *file) {
   size_t dir_len = strlen(dir);
   size_t file_len = strlen(file);
   fpath = (char *)g_calloc(dir_len + file_len + 2, sizeof(char));
+  if (fpath == NULL) {
+    return NULL;
+  }
   strcpy(fpath, dir);
   strcpy(fpath + dir_len + 1, file);
   fpath[dir_len] = '/';
diff --git a/lib/mem/galloc.c b/lib/mem/galloc.c
index 642e010..ea4ea16 100644
--- a/lib/mem/galloc.c
+++ b/lib/mem/galloc.c
@@ -1,29 +1,46 @@
-#include <lib/mem/galloc.h>
+ #include <lib/mem/galloc.h>
 #include <stdlib.h>
+#include <stdio.h>
 
 static ssize_t standing_allocations = 0;
 
-void *g_malloc(size_t size) {
+void *_g_malloc(size_t size, char* file, unsigned int line) {
   void *ptr = malloc(size);
   if (ptr == NULL) {
     return ptr;
   }
+  fprintf(stderr, "+%p :a: %s:%d\n", ptr, file, line);
   standing_allocations++;
   return ptr;
 }
 
-void *g_calloc(size_t n_memb, size_t size) {
+void *_g_calloc(size_t n_memb, size_t size, char* file, unsigned int line) {
   void *ptr = calloc(n_memb, size);
   if (ptr == NULL) {
     return ptr;
   }
+  fprintf(stderr, "+%p :c: %s:%d\n", ptr, file, line);
   standing_allocations++;
   return ptr;
 }
 
-void g_free(void *ptr) {
-  free(ptr);
-  standing_allocations--;
+void *_g_realloc(void *ptr, size_t size, char* file, unsigned int line) {
+  fprintf(stderr, "-%p :r: %s:%d\n", ptr, file, line);
+  void* temp = realloc(ptr, size);
+  if (temp == NULL) {
+    fprintf(stderr, "+%p :r: %s:%d\n", ptr, file, line);
+  } else {
+    fprintf(stderr, "+%p :r: %s:%d\n", temp, file, line);
+  }
+  return temp;
+}
+
+void _g_free(void *ptr, char* file, unsigned int line) {
+  if (ptr != NULL) {
+    fprintf(stderr, "-%p :f: %s:%d\n", ptr, file, line);
+    free(ptr);
+    standing_allocations--;
+  }
 }
 
 ssize_t g_outstanding_allocations() { return standing_allocations; }
diff --git a/src/prog.c b/src/prog.c
index e08bb72..6999a32 100644
--- a/src/prog.c
+++ b/src/prog.c
@@ -97,6 +97,7 @@ int main(int arg_count, char **arg_value) {
         for (size_t index = 0; file_list[index] != NULL; index++) {
           char *fname = file_list[index];
           if (is_tif_ext(fname) == FALSE) {
+	    fprintf(stderr, "Allocation error?\n");
             g_free(file_list[index]);
             continue;
           }
@@ -107,6 +108,7 @@ int main(int arg_count, char **arg_value) {
           //  4. Free up allocations made in this process
           char *fpath = full_path(process_directory, fname);
           if (fpath == NULL) {
+	    fprintf(stderr, "Allocation error?\n");
             g_free(file_list[index]);
             continue;
           }
@@ -118,6 +120,7 @@ int main(int arg_count, char **arg_value) {
           //-----------------------------------------------
           Mask *file_im = tif_to_labels(fpath, &starting_label);
           if (file_im == NULL) {
+	    fprintf(stderr, "Allocation error?\n");
             g_free(fpath);
             g_free(file_list[index]);
             continue;
@@ -206,6 +209,7 @@ int main(int arg_count, char **arg_value) {
     Bitmap *bitmap = image_mask_data_to_bitmap(masks_im);
     if (bitmap != NULL) {
       save_png(bitmap, png_output_file_fullpath);
+      g_free(bitmap->image_buffer);
       g_free(bitmap);
     }
     write_array(bin_output_file_fullpath, masks_im->image[0],