From 9193105c12a7be064d45228d44a3c73660a25473 Mon Sep 17 00:00:00 2001
From: ChinYikMing <yikming2222@gmail.com>
Date: Sun, 24 Dec 2023 17:10:32 +0800
Subject: [PATCH] Consolidate memory allocation

---
 src/cache.c   | 12 ++++++++++++
 src/elf.c     |  3 +++
 src/emulate.c |  7 ++++++-
 src/mpool.c   |  2 ++
 src/riscv.c   |  2 ++
 5 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/cache.c b/src/cache.c
index 328c740d..f363ef46 100644
--- a/src/cache.c
+++ b/src/cache.c
@@ -195,11 +195,20 @@ cache_t *cache_create(int size_bits)
     cache_size = 1 << size_bits;
     for (int i = 0; i < THRESHOLD; i++) {
         cache->lists[i] = malloc(sizeof(struct list_head));
+        if (!cache->lists[i]) {
+            for (int j = 0; j < i; j++) {
+                free(cache->lists[j]);
+            }
+            return NULL;
+        }
         INIT_LIST_HEAD(cache->lists[i]);
     }
 
     cache->map = malloc(sizeof(hashtable_t));
     if (!cache->map) {
+        for (int i = 0; i < THRESHOLD; i++) {
+            free(cache->lists[i]);
+        }
         free(cache->lists);
         free(cache);
         return NULL;
@@ -207,6 +216,9 @@ cache_t *cache_create(int size_bits)
     cache->map->ht_list_head = malloc(cache_size * sizeof(struct hlist_head));
     if (!cache->map->ht_list_head) {
         free(cache->map);
+        for (int i = 0; i < THRESHOLD; i++) {
+            free(cache->lists[i]);
+        }
         free(cache->lists);
         free(cache);
         return NULL;
diff --git a/src/elf.c b/src/elf.c
index 48b5d7ec..7d441402 100644
--- a/src/elf.c
+++ b/src/elf.c
@@ -3,6 +3,7 @@
  * "LICENSE" for information on usage and redistribution of this file.
  */
 
+#include <assert.h>
 #include <stdlib.h>
 #include <string.h>
 
@@ -70,6 +71,7 @@ struct elf_internal {
 elf_t *elf_new(void)
 {
     elf_t *e = malloc(sizeof(elf_t));
+    assert(e);
     e->hdr = NULL;
     e->raw_size = 0;
     e->symbols = map_init(int, char *, map_cmp_uint);
@@ -342,6 +344,7 @@ bool elf_open(elf_t *e, const char *input)
     /* allocate memory */
     free(e->raw_data);
     e->raw_data = malloc(e->raw_size);
+    assert(e->raw_data);
 
     /* read data into memory */
     const size_t r = fread(e->raw_data, 1, e->raw_size, f);
diff --git a/src/emulate.c b/src/emulate.c
index 9caba2f0..5a999d03 100644
--- a/src/emulate.c
+++ b/src/emulate.c
@@ -632,8 +632,10 @@ static void block_translate(riscv_t *rv, block_t *block)
 #if RV32_HAS(EXT_C)
                 || ir->opcode == rv_insn_cjalr || ir->opcode == rv_insn_cjr
 #endif
-            )
+            ) {
                 ir->branch_table = calloc(1, sizeof(branch_history_table_t));
+                assert(ir->branch_table);
+            }
             break;
         }
 
@@ -659,6 +661,7 @@ static void block_translate(riscv_t *rv, block_t *block)
     if (count > 1) {                                              \
         ir->opcode = IIF(RW)(rv_insn_fuse4, rv_insn_fuse3);       \
         ir->fuse = malloc(count * sizeof(opcode_fuse_t));         \
+        assert(ir->fuse);                                         \
         ir->imm2 = count;                                         \
         memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));              \
         ir->impl = dispatch_table[ir->opcode];                    \
@@ -849,6 +852,7 @@ static void match_pattern(riscv_t *rv, block_t *block)
                 if (count > 1) {
                     ir->opcode = rv_insn_fuse1;
                     ir->fuse = malloc(count * sizeof(opcode_fuse_t));
+                    assert(ir->fuse);
                     ir->imm2 = count;
                     memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
                     ir->impl = dispatch_table[ir->opcode];
@@ -887,6 +891,7 @@ static void match_pattern(riscv_t *rv, block_t *block)
             }
             if (count > 1) {
                 ir->fuse = malloc(count * sizeof(opcode_fuse_t));
+                assert(ir->fuse);
                 memcpy(ir->fuse, ir, sizeof(opcode_fuse_t));
                 ir->opcode = rv_insn_fuse7;
                 ir->imm2 = count;
diff --git a/src/mpool.c b/src/mpool.c
index 0815b4bd..0f8e1a27 100644
--- a/src/mpool.c
+++ b/src/mpool.c
@@ -82,6 +82,8 @@ static void *mpool_extend(mpool_t *mp)
     if (!p)
         return NULL;
     area_t *new_area = malloc(sizeof(area_t));
+    if (!new_area)
+        return NULL;
     new_area->mapped = p;
     new_area->next = NULL;
     size_t chunk_count = pool_size / (sizeof(memchunk_t) + mp->chunk_size);
diff --git a/src/riscv.c b/src/riscv.c
index c5c14e4b..a0ae232c 100644
--- a/src/riscv.c
+++ b/src/riscv.c
@@ -27,6 +27,7 @@ static void block_map_init(block_map_t *map, const uint8_t bits)
     map->block_capacity = 1 << bits;
     map->size = 0;
     map->map = calloc(map->block_capacity, sizeof(struct block *));
+    assert(map->map);
 }
 
 /* clear all block in the block map */
@@ -113,6 +114,7 @@ riscv_t *rv_create(const riscv_io_t *io,
     assert(io);
 
     riscv_t *rv = calloc(1, sizeof(riscv_t));
+    assert(rv);
 
     /* copy over the IO interface */
     memcpy(&rv->io, io, sizeof(riscv_io_t));