2 files changed, 13 insertions, 9 deletions

diff --git a/analyzer/structures/hashmap.c b/analyzer/structures/hashmap.c
index 84cc289..5eb27ae 100644
--- a/analyzer/structures/hashmap.c
+++ b/analyzer/structures/hashmap.c
@@ -3,6 +3,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "utils/memory.h"
+
 HashMap *hashmap_new(
     size_t key_size,
     size_t value_size,
@@ -11,9 +13,9 @@ HashMap *hashmap_new(
     void (*key_cleanup)(void *key),
     void (*value_cleanup)(void *key)
 ) {
-    HashMap *map = malloc(sizeof(HashMap));
+    HashMap *map = safe_malloc(sizeof(HashMap));
     // initialize NULLs
-    memset(map->table, 0, sizeof(HashMapNode *) * HASHMAP_TABLE_SIZE);
+    memset(map->table, 0, HASHMAP_TABLE_SIZE * sizeof(HashMapNode *));
     map->key_size = key_size;
     map->value_size = value_size;
     map->hash = hash;
@@ -42,9 +44,9 @@ void hashmap_put(HashMap *map, void *key, void *value) {
     }
 
     // if it doesn't exist - create new node
-    HashMapNode *new_node = malloc(sizeof(HashMapNode));
-    new_node->key = malloc(map->key_size);
-    new_node->value = malloc(map->value_size);
+    HashMapNode *new_node = safe_malloc(sizeof(HashMapNode));
+    new_node->key = safe_malloc(map->key_size);
+    new_node->value = safe_malloc(map->value_size);
     memcpy(new_node->key, key, map->key_size);
     memcpy(new_node->value, value, map->value_size);
 
diff --git a/analyzer/structures/heap.c b/analyzer/structures/heap.c
index 6102940..1267668 100644
--- a/analyzer/structures/heap.c
+++ b/analyzer/structures/heap.c
@@ -3,9 +3,11 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "utils/memory.h"
+
 Heap *heap_new(size_t element_size, int (*compare)(void *element1, void *element2), void (*element_cleanup)(void *element)) {
-    Heap *heap = malloc(sizeof(Heap));
-    heap->data = malloc(INIT_HEAP_CAPACITY * sizeof(void *));
+    Heap *heap = safe_malloc(sizeof(Heap));
+    heap->data = safe_malloc(INIT_HEAP_CAPACITY * sizeof(void *));
     heap->n_elements = 0;
     heap->element_size = element_size;
     heap->capacity = INIT_HEAP_CAPACITY;
@@ -40,7 +42,7 @@ void swap(void **p_el1, void **p_el2) {
 void expand_heap(Heap *heap, size_t el_count) {
     while (heap->capacity < el_count)
         heap->capacity = (heap->capacity << 1) | 1;
-    heap->data = realloc(heap->data, heap->capacity * sizeof(void *));
+    heap->data = safe_realloc(heap->data, heap->capacity * sizeof(void *));
 }
 
 void heap_insert(Heap *heap, void *element) {
@@ -50,7 +52,7 @@ void heap_insert(Heap *heap, void *element) {
 
     // insert into the last position
     size_t i = heap->n_elements;
-    heap->data[i] = malloc(heap->element_size);
+    heap->data[i] = safe_malloc(heap->element_size);
     memcpy(heap->data[i], element, heap->element_size);
     heap->n_elements++;