From 587e2f38cebabbcca3d066a781fed35c6a845758 Mon Sep 17 00:00:00 2001
From: Liam Keegan <liam@keegan.ch>
Date: Wed, 29 Jan 2025 11:33:38 +0100
Subject: [PATCH] Fix additional integer overflow issues

- correct the condition for using 32bit int indices
  - previous condition: `n_points < UINT_MAX`
  - however the data array in the C code has size `n_points * n_dim` and is indexed using a 32bit int type (see e.g. the PA macro)
  - this means results were incorrect for the case where `n_points < UINT_MAX` and `n_points * n_dim >= UINT_MAX`
  - new condition: `n_points * n_dim < UINT_MAX`
  - resolves #138
- render mako template
  - the generated _kdtree_core.c did not contain latest changes from _kdtree_core.c.mako (sorry my mistake)
  - this meant results were incorrect for the case `n_points < UINT_MAX` and `n_query * k >= UINT_MAX`
  - resolves #137
- use signed 64bit ints for the OpenMP loop on all platforms
  - to support OpenMP versions < 3 which only allow signed integers as loop counters
---
 pykdtree/_kdtree_core.c      | 40 +++++++++++-------------------------
 pykdtree/_kdtree_core.c.mako |  7 +------
 pykdtree/kdtree.pyx          |  2 +-
 3 files changed, 14 insertions(+), 35 deletions(-)

diff --git a/pykdtree/_kdtree_core.c b/pykdtree/_kdtree_core.c
index bfcc90d..d587820 100644
--- a/pykdtree/_kdtree_core.c
+++ b/pykdtree/_kdtree_core.c
@@ -836,14 +836,10 @@ void search_tree_float_int32_t(Tree_float_int32_t *tree, float *pa, float *point
     int8_t no_dims = tree->no_dims;
     float *bbox = tree->bbox;
     uint32_t *pidx = tree->pidx;
-    uint32_t j = 0;
-#if defined(_MSC_VER) && defined(_OPENMP)
-    int32_t i = 0;
-    int32_t local_num_points = (int32_t) num_points;
-#else
-    uint32_t i;
-    uint32_t local_num_points = num_points;
-#endif
+    /* use 64-bit ints for indexing to avoid overflow, use signed ints to support all Openmp implementations */
+    int64_t i = 0;
+    int64_t j = 0;
+    int64_t local_num_points = (int64_t) num_points;
     Node_float_int32_t *root = (Node_float_int32_t *)tree->root;
 
     /* Queries are OpenMP enabled */
@@ -1412,14 +1408,10 @@ void search_tree_float_int64_t(Tree_float_int64_t *tree, float *pa, float *point
     int8_t no_dims = tree->no_dims;
     float *bbox = tree->bbox;
     uint64_t *pidx = tree->pidx;
-    uint64_t j = 0;
-#if defined(_MSC_VER) && defined(_OPENMP)
+    /* use 64-bit ints for indexing to avoid overflow, use signed ints to support all Openmp implementations */
     int64_t i = 0;
+    int64_t j = 0;
     int64_t local_num_points = (int64_t) num_points;
-#else
-    uint64_t i;
-    uint64_t local_num_points = num_points;
-#endif
     Node_float_int64_t *root = (Node_float_int64_t *)tree->root;
 
     /* Queries are OpenMP enabled */
@@ -2057,14 +2049,10 @@ void search_tree_double_int32_t(Tree_double_int32_t *tree, double *pa, double *p
     int8_t no_dims = tree->no_dims;
     double *bbox = tree->bbox;
     uint32_t *pidx = tree->pidx;
-    uint32_t j = 0;
-#if defined(_MSC_VER) && defined(_OPENMP)
-    int32_t i = 0;
-    int32_t local_num_points = (int32_t) num_points;
-#else
-    uint32_t i;
-    uint32_t local_num_points = num_points;
-#endif
+    /* use 64-bit ints for indexing to avoid overflow, use signed ints to support all Openmp implementations */
+    int64_t i = 0;
+    int64_t j = 0;
+    int64_t local_num_points = (int64_t) num_points;
     Node_double_int32_t *root = (Node_double_int32_t *)tree->root;
 
     /* Queries are OpenMP enabled */
@@ -2633,14 +2621,10 @@ void search_tree_double_int64_t(Tree_double_int64_t *tree, double *pa, double *p
     int8_t no_dims = tree->no_dims;
     double *bbox = tree->bbox;
     uint64_t *pidx = tree->pidx;
-    uint64_t j = 0;
-#if defined(_MSC_VER) && defined(_OPENMP)
+    /* use 64-bit ints for indexing to avoid overflow, use signed ints to support all Openmp implementations */
     int64_t i = 0;
+    int64_t j = 0;
     int64_t local_num_points = (int64_t) num_points;
-#else
-    uint64_t i;
-    uint64_t local_num_points = num_points;
-#endif
     Node_double_int64_t *root = (Node_double_int64_t *)tree->root;
 
     /* Queries are OpenMP enabled */
diff --git a/pykdtree/_kdtree_core.c.mako b/pykdtree/_kdtree_core.c.mako
index fc921b8..94602a9 100644
--- a/pykdtree/_kdtree_core.c.mako
+++ b/pykdtree/_kdtree_core.c.mako
@@ -716,15 +716,10 @@ void search_tree_${DTYPE}_${ITYPE}(Tree_${DTYPE}_${ITYPE} *tree, ${DTYPE} *pa, $
     int8_t no_dims = tree->no_dims;
     ${DTYPE} *bbox = tree->bbox;
     u${ITYPE} *pidx = tree->pidx;
-#if defined(_MSC_VER) && defined(_OPENMP)
+    /* use 64-bit ints for indexing to avoid overflow, use signed ints to support all Openmp implementations */
     int64_t i = 0;
     int64_t j = 0;
     int64_t local_num_points = (int64_t) num_points;
-#else
-    uint64_t i;
-    uint64_t j;
-    uint64_t local_num_points = (uint64_t) num_points;
-#endif
     Node_${DTYPE}_${ITYPE} *root = (Node_${DTYPE}_${ITYPE} *)tree->root;
 
     /* Queries are OpenMP enabled */
diff --git a/pykdtree/kdtree.pyx b/pykdtree/kdtree.pyx
index 476155e..7c2eeca 100644
--- a/pykdtree/kdtree.pyx
+++ b/pykdtree/kdtree.pyx
@@ -162,7 +162,7 @@ cdef class KDTree:
 
         # Get tree info
         self.n = <uint64_t>data_pts.shape[0]
-        self._use_int32_t = self.n < UINT32_MAX
+        self._use_int32_t = self.n * <uint64_t>data_pts.shape[1] < UINT32_MAX
         self.leafsize = <uint32_t>leafsize
         if data_pts.ndim == 1:
             self.ndim = 1