5 files changed, 42 insertions, 1 deletions

diff --git a/Makefile b/Makefile
index aa98152..158373b 100644
--- a/Makefile
+++ b/Makefile
@@ -1,7 +1,7 @@
 all: compile run
 
 compile:
-	gcc -g sirtet.c pieces.c ui.c -lncurses -o sirtet
+	gcc -g sirtet.c pieces.c ui.c util.c -lncurses -o sirtet
 
 run:
 	./sirtet
diff --git a/pieces.c b/pieces.c
index f1ed17c..3b71003 100644
--- a/pieces.c
+++ b/pieces.c
@@ -2,6 +2,7 @@
 #include <stdbool.h>
 #include <string.h>
 #include "pieces.h"
+#include "util.h"
 
 const int N_SHAPES = 17;
 const struct piece SHAPES[] = {
@@ -27,6 +28,9 @@ const struct piece SHAPES[] = {
 void transpose(struct piece* pc) {
     // transpose blocks
     char* old = malloc(strlen(pc->blocks) + 1);
+    if (old == NULL)
+        mallocfail();
+
     strcpy(old, pc->blocks); // before transposition
     for (int r = 0; r < pc->h; r++) {
         for (int c = 0; c < pc->w; c++) {
@@ -44,6 +48,9 @@ void rotate(struct piece* pc) {
     // rotate `pc` ccw by 90 degrees
     // transpose blocks
     char* old = malloc(strlen(pc->blocks) + 1);
+    if (old == NULL)
+        mallocfail();
+
     strcpy(old, pc->blocks); // before transposition
     for (int r = 0; r < pc->h; r++) {
         for (int c = 0; c < pc->w; c++) {
@@ -62,8 +69,14 @@ struct piece* randpiece() {
     // then do random transformations to the copy
     struct piece shape = SHAPES[rand() % N_SHAPES];
     char* blocks = malloc(strlen(shape.blocks) + 1);
+    if (blocks == NULL)
+        mallocfail();
+
     strcpy(blocks, shape.blocks);
     struct piece* pc = malloc(sizeof(struct piece));
+    if (pc == NULL)
+        mallocfail();
+
     pc->h = shape.h;
     pc->w = shape.w;
     pc->points = shape.points;
diff --git a/sirtet.c b/sirtet.c
index 0a8cc8f..61427c1 100644
--- a/sirtet.c
+++ b/sirtet.c
@@ -7,6 +7,7 @@
 #include <ncurses.h>
 #include "pieces.h"
 #include "ui.h"
+#include "util.h"
 
 void clearmap(char** map, int mapH, int mapW) {
     for (int r = 0; r < mapH; r++) {
@@ -23,6 +24,9 @@ int clearfull(char** map, int mapH, int mapW) {
     // returns total number of blocks cleared
     bool* rows_to_clear = malloc(mapH * sizeof(bool));
     bool* cols_to_clear = malloc(mapW * sizeof(bool));
+    if (rows_to_clear == NULL || cols_to_clear == NULL)
+        mallocfail();
+
     int rtc_count = 0;
     int ctc_count = 0;
     // scan all rows and columns
@@ -65,10 +69,18 @@ int clearfull(char** map, int mapH, int mapW) {
 void sirtet(int H, int W, int P) {
     // init memory and game state
     char** map = malloc(H * sizeof(char*));
+    if (map == NULL)
+        mallocfail();
+
     for (int r = 0; r < H; r++) {
         map[r] = malloc(W * sizeof(char));
+        if (map[r] == NULL)
+            mallocfail();
     }
     struct piece** hand = malloc(P * sizeof(struct piece*));
+    if (hand == NULL)
+        mallocfail();
+
     srand(time(NULL));
     clearmap(map, H, W);
     refillpieces(hand, P);
diff --git a/util.c b/util.c
new file mode 100644
index 0000000..0dc757b
--- /dev/null
+++ b/util.c
@@ -0,0 +1,10 @@
+#include <stdlib.h>
+#include <stdio.h>
+#include <ncurses.h>
+
+void mallocfail() {
+    endwin();
+    fprintf(stderr, "Failed to allocate memory\n");
+    exit(1);
+}
+
diff --git a/util.h b/util.h
new file mode 100644
index 0000000..1b5ec50
--- /dev/null
+++ b/util.h
@@ -0,0 +1,6 @@
+#ifndef UTIL_H
+#define UTIL_H
+
+void mallocfail();
+#endif
+