From 2497eb2d28b275b740a8de8a557924cb8d08799f Mon Sep 17 00:00:00 2001
From: Juhana Lankinen <juhana.lankinen@csc.fi>
Date: Tue, 27 Aug 2024 14:06:07 +0300
Subject: [PATCH] Fix integer overflows; Add tests

---
 src/tools.hpp                    | 43 +++++++++++++++++++++-----------
 tests/unit_tests/tools_tests.cpp | 42 +++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/src/tools.hpp b/src/tools.hpp
index ca4336b..b642669 100644
--- a/src/tools.hpp
+++ b/src/tools.hpp
@@ -61,21 +61,36 @@ namespace FsGridTools {
    //! Helper function: given a global cellID, calculate the global cell coordinate from it.
    // This is then used do determine the task responsible for this cell, and the
    // local cell index in it.
-   static std::array<FsIndex_t, 3> globalIDtoCellCoord(GlobalID id, std::array<FsSize_t, 3>& globalSize) {
-
-      // Transform globalID to global cell coordinate
-      std::array<FsIndex_t, 3> cell;
-
-      assert(id >= 0);
-      assert(id < globalSize[0] * globalSize[1] * globalSize[2]);
-
-      int stride = 1;
-      for (int i = 0; i < 3; i++) {
-         cell[i] = (id / stride) % globalSize[i];
-         stride *= globalSize[i];
+   static std::array<FsIndex_t, 3> globalIDtoCellCoord(GlobalID id, const std::array<FsSize_t, 3>& globalSize) {
+      // We're returning FsIndex_t, which is int32_t. globalSizes are FsSize_t, which are uint32_t.
+      // Need to check that the globalSizes aren't larger than maximum of int32_t
+      const bool globalSizeOverflows = globalSize[0] > std::numeric_limits<FsIndex_t>::max() ||
+                                       globalSize[1] > std::numeric_limits<FsIndex_t>::max() ||
+                                       globalSize[2] > std::numeric_limits<FsIndex_t>::max();
+      const bool idNegative = id < 0;
+
+      // To avoid overflow, do this instead of checking for id < product of globalSize
+      // The number of divisions stays the same anyway
+      const GlobalID idPerGs0 = id / globalSize[0];
+      const GlobalID idPerGs0PerGs1 = idPerGs0 / globalSize[1];
+      const bool idTooLarge = idPerGs0PerGs1 >= globalSize[2];
+
+      const bool badInput = idTooLarge || idNegative || globalSizeOverflows;
+
+      if (badInput) {
+         // For bad input, return bad output
+         return {
+             std::numeric_limits<FsIndex_t>::min(),
+             std::numeric_limits<FsIndex_t>::min(),
+             std::numeric_limits<FsIndex_t>::min(),
+         };
+      } else {
+         return {
+             (FsIndex_t)(id % globalSize[0]),
+             (FsIndex_t)(idPerGs0 % globalSize[1]),
+             (FsIndex_t)(idPerGs0PerGs1 % globalSize[2]),
+         };
       }
-
-      return cell;
    }
 
    //! Helper function to optimize decomposition of this grid over the given number of tasks
diff --git a/tests/unit_tests/tools_tests.cpp b/tests/unit_tests/tools_tests.cpp
index 72a79cb..6c0e4a1 100644
--- a/tests/unit_tests/tools_tests.cpp
+++ b/tests/unit_tests/tools_tests.cpp
@@ -1,6 +1,7 @@
 #include "tools.hpp"
 
 #include <gtest/gtest.h>
+#include <limits>
 
 TEST(FsGridToolsTests, calcLocalStart1) {
     constexpr FsGridTools::FsSize_t numGlobalCells = 1024u;
@@ -45,3 +46,44 @@ TEST(FsGridToolsTests, calcLocalSize2) {
       ASSERT_EQ(FsGridTools::calcLocalSize(numGlobalCells, numTasks, i), 10);
    }
 }
+
+TEST(FsGridToolsTests, globalIDtoCellCoord1) {
+   constexpr std::array<FsGridTools::FsSize_t, 3> globalSize = {3, 7, 45};
+
+   for (FsGridTools::FsSize_t i = 0; i < globalSize[0] * globalSize[1] * globalSize[2]; i++) {
+      const std::array<FsGridTools::FsIndex_t, 3> result = {
+          (FsGridTools::FsIndex_t)(i % globalSize[0]),
+          (FsGridTools::FsIndex_t)((i / globalSize[0]) % globalSize[1]),
+          (FsGridTools::FsIndex_t)((i / (globalSize[0] * globalSize[1])) % globalSize[2]),
+      };
+      ASSERT_EQ(FsGridTools::globalIDtoCellCoord(i, globalSize), result);
+   }
+}
+
+TEST(FsGridToolsTests, globalIDtoCellCoord_globalSize_would_overflow) {
+   constexpr uint32_t maxUint = std::numeric_limits<uint32_t>::max();
+   constexpr std::array<FsGridTools::FsSize_t, 3> globalSize = {maxUint, 1, 1};
+   const std::array<FsGridTools::FsIndex_t, 3> result = {
+       -2147483648,
+       -2147483648,
+       -2147483648,
+   };
+   ASSERT_EQ(FsGridTools::globalIDtoCellCoord(globalSize[0] - 1, globalSize), result);
+}
+
+TEST(FsGridToolsTests, globalIDtoCellCoord_globalSize0_is_maximum_int) {
+   constexpr int32_t maxInt = std::numeric_limits<int32_t>::max();
+   constexpr std::array<FsGridTools::FsSize_t, 3> globalSize = {maxInt, maxInt, maxInt};
+
+   std::array<FsGridTools::FsIndex_t, 3> result = {maxInt - 1, 0, 0};
+   ASSERT_EQ(FsGridTools::globalIDtoCellCoord(globalSize[0] - 1, globalSize), result);
+
+   result = {0, 1, 0};
+   ASSERT_EQ(FsGridTools::globalIDtoCellCoord(globalSize[0], globalSize), result);
+
+   result = {0, 0, 1};
+   ASSERT_EQ(FsGridTools::globalIDtoCellCoord((int64_t)globalSize[0] * globalSize[0], globalSize), result);
+
+   result = {maxInt - 1, maxInt - 1, 0};
+   ASSERT_EQ(FsGridTools::globalIDtoCellCoord((int64_t)globalSize[0] * globalSize[0] - 1, globalSize), result);
+}