From 1a6917171d441224549ac64f748d21bbb6cfd87c Mon Sep 17 00:00:00 2001
From: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com>
Date: Sat, 10 Oct 2026 05:52:01 +0530
Subject: [PATCH] fix: prevent integer overflow in Psi offsets

---
 source/source_psi/psi.cpp           | 27 ++++++-----
 source/source_psi/test/psi_test.cpp | 69 +++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 12 deletions(-)

diff --git a/source/source_psi/psi.cpp b/source/source_psi/psi.cpp
index 8f64b0b8ea8..80a03005a75 100644
--- a/source/source_psi/psi.cpp
+++ b/source/source_psi/psi.cpp
@@ -278,7 +278,7 @@ T* Psi<T, Device>::get_pointer(const int& ikb) const
 {
     assert(ikb >= 0);
     assert(this->k_first ? ikb < this->nbands : ikb < this->nk);
-    return this->psi_current + ikb * this->nbasis;
+    return this->psi_current + static_cast<std::size_t>(ikb) * this->nbasis;
 }
 
 template <typename T, typename Device>
@@ -377,7 +377,7 @@ void Psi<T, Device>::fix_k(const int ik) const
     {
         this->current_b = 0;
     }
-    int base = this->current_b * this->nk * this->nbasis;
+    const std::size_t base = static_cast<std::size_t>(this->current_b) * this->nk * this->nbasis;
     if (ik >= this->nk)
     {
         // mem_saver: fix to base
@@ -386,7 +386,8 @@ void Psi<T, Device>::fix_k(const int ik) const
     }
     else
     {
-        this->psi_bias = k_first ? ik * this->nbands * this->nbasis : base + ik * this->nbasis;
+        this->psi_bias = k_first ? static_cast<std::size_t>(ik) * this->nbands * this->nbasis
+                                : base + static_cast<std::size_t>(ik) * this->nbasis;
         this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
     }
 }
@@ -400,7 +401,7 @@ void Psi<T, Device>::fix_b(const int ib) const
     {
         this->current_k = 0;
     }
-    int base = this->current_k * this->nbands * this->nbasis;
+    const std::size_t base = static_cast<std::size_t>(this->current_k) * this->nbands * this->nbasis;
     if (ib >= this->nbands)
     {
         // mem_saver: fix to base
@@ -409,7 +410,8 @@ void Psi<T, Device>::fix_b(const int ib) const
     }
     else
     {
-        this->psi_bias = k_first ? base + ib * this->nbasis : ib * this->nk * this->nbasis;
+        this->psi_bias = k_first ? base + static_cast<std::size_t>(ib) * this->nbasis
+                                : static_cast<std::size_t>(ib) * this->nk * this->nbasis;
         this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
     }
 }
@@ -427,7 +429,8 @@ void Psi<T, Device>::fix_kb(const int ik, const int ib) const
     }
     else
     {
-        this->psi_bias = k_first ? (ik * this->nbands + ib) * this->nbasis : (ib * this->nk + ik) * this->nbasis;
+        this->psi_bias = k_first ? (static_cast<std::size_t>(ik) * this->nbands + ib) * this->nbasis
+                                : (static_cast<std::size_t>(ib) * this->nk + ik) * this->nbasis;
         this->psi_current = const_cast<T*>(&(this->psi[psi_bias]));
     }
 }
@@ -437,8 +440,8 @@ T& Psi<T, Device>::operator()(const int ikb1, const int ikb2, const int ibasis)
 {
     assert(ikb1 >= 0 && ikb2 >= 0 && ibasis >= 0);
     assert(this->k_first ? ikb1 < this->nk && ikb2 < this->nbands : ikb1 < this->nbands && ikb2 < this->nk);
-    return this->k_first ? this->psi[(ikb1 * this->nbands + ikb2) * this->nbasis + ibasis]
-                         : this->psi[(ikb1 * this->nk + ikb2) * this->nbasis + ibasis];
+    return this->k_first ? this->psi[(static_cast<std::size_t>(ikb1) * this->nbands + ikb2) * this->nbasis + ibasis]
+                        : this->psi[(static_cast<std::size_t>(ikb1) * this->nk + ikb2) * this->nbasis + ibasis];
 }
 
 template <typename T, typename Device>
@@ -447,7 +450,7 @@ T& Psi<T, Device>::operator()(const int ikb2, const int ibasis) const
     assert(this->k_first ? this->current_b == 0 : this->current_k == 0);
     assert(this->k_first ? ikb2 >= 0 && ikb2 < this->nbands : ikb2 >= 0 && ikb2 < this->nk);
     assert(ibasis >= 0 && ibasis < this->nbasis);
-    return this->psi_current[ikb2 * this->nbasis + ibasis];
+    return this->psi_current[static_cast<std::size_t>(ikb2) * this->nbasis + ibasis];
 }
 
 template <typename T, typename Device>
@@ -508,13 +511,13 @@ std::tuple<const T*, int> Psi<T, Device>::to_range(const Range& range) const
     }
     else if (i1 < 0) // [r1, r2] is the range of index1 with length m
     {
-        const T* p = &this->psi[r1 * (k_first ? this->nbands : this->nk) * this->nbasis];
+        const T* p = &this->psi[static_cast<std::size_t>(r1) * (k_first ? this->nbands : this->nk) * this->nbasis];
         int m = (r2 - r1 + 1) * this->get_npol();
         return std::tuple<const T*, int>(p, m);
     }
     else // [r1, r2] is the range of index2 with length m
     {
-        const T* p = &this->psi[(i1 * (k_first ? this->nbands : this->nk) + r1) * this->nbasis];
+        const T* p = &this->psi[(static_cast<std::size_t>(i1) * (k_first ? this->nbands : this->nk) + r1) * this->nbasis];
         int m = (r2 - r1 + 1) * this->get_npol();
         return std::tuple<const T*, int>(p, m);
     }
@@ -553,4 +556,4 @@ template Psi<std::complex<float>, base_device::DEVICE_GPU>::Psi(
 template Psi<std::complex<double>, base_device::DEVICE_GPU>::Psi(
     const Psi<std::complex<float>, base_device::DEVICE_GPU>&);
 #endif
-} // namespace psi
\ No newline at end of file
+} // namespace psi
diff --git a/source/source_psi/test/psi_test.cpp b/source/source_psi/test/psi_test.cpp
index 532d308cf24..74f105de8ed 100644
--- a/source/source_psi/test/psi_test.cpp
+++ b/source/source_psi/test/psi_test.cpp
@@ -1,6 +1,75 @@
 #include "source_psi/psi.h"
 
 #include <gtest/gtest.h>
+#include <memory>
+
+TEST(PsiLargeOffsets, ThreeDimensionalAccess)
+{
+    // No zero-initialization: the test touches only one page of this 12 GiB buffer.
+    std::unique_ptr<float[]> storage(new float[3221225472ULL]);
+    float* const data = storage.get();
+    data[2684354567ULL] = 42.0f;
+    for (const bool k_first : {true, false})
+    {
+        const int nk = k_first ? 3 : 2;
+        const int nbands = k_first ? 2 : 3;
+        const int ik = k_first ? 2 : 1;
+        const int ib = k_first ? 1 : 2;
+        psi::Psi<float> wave(data, nk, nbands, 536870912, 536870912, k_first);
+
+        if (k_first)
+        {
+            wave.fix_k(ik);
+        }
+        else
+        {
+            wave.fix_b(ib);
+        }
+        ASSERT_EQ(wave.get_psi_bias(), 2147483648ULL);
+        ASSERT_EQ(wave.get_pointer(), data + 2147483648ULL);
+        ASSERT_EQ(wave.get_pointer(1), data + 2684354560ULL);
+
+        if (k_first)
+        {
+            wave.fix_b(ib);
+        }
+        else
+        {
+            wave.fix_k(ik);
+        }
+        EXPECT_EQ(wave.get_pointer(), data + 2684354560ULL);
+
+        wave.fix_kb(ik, ib);
+        ASSERT_EQ(wave.get_psi_bias(), 2684354560ULL);
+        EXPECT_EQ(wave.get_pointer(), data + 2684354560ULL);
+        ASSERT_EQ(&wave(2, 1, 7), data + 2684354567ULL);
+        EXPECT_FLOAT_EQ(wave(2, 1, 7), 42.0f);
+        EXPECT_EQ(&wave(1, 1, 536870911), data + 2147483647ULL);
+
+        const psi::Range range(k_first, 2, 1, 1);
+        const auto result = wave.to_range(range);
+        ASSERT_EQ(std::get<0>(result), data + 2684354560ULL);
+        const psi::Range outer_range(k_first, -1, 2, 2);
+        const auto outer_result = wave.to_range(outer_range);
+        EXPECT_EQ(std::get<0>(outer_result), data + 2147483648ULL);
+    }
+}
+
+TEST(PsiLargeOffsets, TwoDimensionalAccess)
+{
+    std::unique_ptr<float[]> storage(new float[3221225472ULL]);
+    float* const data = storage.get();
+    data[2147483655ULL] = 24.0f;
+    for (const bool k_first : {true, false})
+    {
+        const int nk = k_first ? 1 : 3;
+        const int nbands = k_first ? 3 : 1;
+        psi::Psi<float> wave(data, nk, nbands, 1073741824, 1073741824, k_first);
+        ASSERT_EQ(wave.get_pointer(2), data + 2147483648ULL);
+        ASSERT_EQ(&wave(2, 7), data + 2147483655ULL);
+        EXPECT_FLOAT_EQ(wave(2, 7), 24.0f);
+    }
+}
 
 class TestPsi : public ::testing::Test
 {