From 5b1a2a5f5ecf29fe0b9ab02314ef3fefc3539a46 Mon Sep 17 00:00:00 2001
From: ShubhamDesai <42180509+ShubhamDesai@users.noreply.github.com>
Date: Thu, 15 Aug 2024 02:21:12 -0400
Subject: [PATCH] r.terraflow: Handle memory reallocation errors gracefully
 (#3970)

Co-authored-by: Nicklas Larsson <n_larsson@yahoo.com>
---
 raster/r.terraflow/unionFind.h | 48 ++++++++++++++++++++++++----------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/raster/r.terraflow/unionFind.h b/raster/r.terraflow/unionFind.h
index 698dcfc0c97..920421ca5d6 100644
--- a/raster/r.terraflow/unionFind.h
+++ b/raster/r.terraflow/unionFind.h
@@ -19,11 +19,15 @@
 #ifndef __UNION_FIND
 #define __UNION_FIND
 
-#include <assert.h>
-#include <stdlib.h>
-#include <string.h>
+#include <cassert>
+#include <cstdlib>
+#include <cstring>
 #include <iostream>
 
+extern "C" {
+#include <grass/glocale.h>
+}
+
 /* initial range guesstimate */
 #define UNION_INITIAL_SIZE 2000
 
@@ -68,11 +72,19 @@ unionFind<T>::unionFind()
 {
     maxsize = UNION_INITIAL_SIZE;
     /*  parent = new (long)[maxsize]; */
-    parent = (T *)calloc(maxsize, sizeof(T));
-    assert(parent);
+    if (void *new_parent = std::calloc(maxsize, sizeof(T))) {
+        parent = static_cast<T *>(new_parent);
+    }
+    else {
+        G_fatal_error(_("Not enough memory for %s"), "parent");
+    }
     /*  rank = new (long)[maxsize]; */
-    rank = (T *)calloc(maxsize, sizeof(T));
-    assert(rank);
+    if (void *new_rank = std::calloc(maxsize, sizeof(T))) {
+        rank = static_cast<T *>(new_rank);
+    }
+    else {
+        G_fatal_error(_("Not enough memory for %s"), "rank");
+    }
 }
 
 /************************************************************/
@@ -126,13 +138,21 @@ inline void unionFind<T>::makeSet(T x)
     if (x >= maxsize) {
         /* reallocate parent */
         cout << "UnionFind::makeSet: reallocate double " << maxsize << "\n";
-        parent = (T *)realloc(parent, 2 * maxsize * sizeof(T));
-        assert(parent);
-        memset(parent + maxsize, 0, maxsize * sizeof(T));
-        /*reallocate rank */
-        rank = (T *)realloc(rank, 2 * maxsize * sizeof(T));
-        assert(rank);
-        memset(rank + maxsize, 0, maxsize * sizeof(T));
+        if (void *new_parent = std::realloc(parent, 2 * maxsize * sizeof(T))) {
+            parent = static_cast<T *>(new_parent);
+            std::memset(parent + maxsize, 0, maxsize * sizeof(T));
+        }
+        else {
+            G_fatal_error(_("Not enough memory for %s"), "parent");
+        }
+        /* reallocate rank */
+        if (void *new_rank = std::realloc(rank, 2 * maxsize * sizeof(T))) {
+            rank = static_cast<T *>(new_rank);
+            std::memset(rank + maxsize, 0, maxsize * sizeof(T));
+        }
+        else {
+            G_fatal_error(_("Not enough memory for %s"), "rank");
+        }
         /*update maxsize */
         maxsize *= 2;
     }