Issue #1621: Fix DataView iterator segfault with axis of size 0 (#2164)

**Context:**
The `DataView` iterators suffer an underflow of unsigned types that can
lead to a segfault when presented an axis of size 0, as discussed in
#1621 and
[here](#1598 (comment)).

**Description of the Change:**
This change removes the `var - 1` type arithmetic that lead to the
underflows, which in turn prevents the segfaults, by adjusting the
indexing done in the for loop.

**Benefits:**
No more underflow errors and segfaults, accurate `view.size()` and
`std::distance(view.begin(), view.end())` results in the case of an axis
of size 0.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Closes #1621 

[sc-88502]

---------

Co-authored-by: Joey Carter <[email protected]>
Co-authored-by: David Ittah <[email protected]>

Author: 3 people

doc/releases/changelog-dev.md

@@ -42,7 +42,12 @@
   [(#1984)](https://github.com/PennyLaneAI/catalyst/pull/1984)
 
 * Split `from_plxpr.py` into two files.
-  [(#2142)](https://github.com/PennyLaneAI/catalyst/pull/2142)
+  [(#2142)](https://github.com/PennyLaneAI/catalyst/pull/2142)	
+
+* Re-work `DataView` to avoid an axis of size 0 possibly triggering a segfault via an underflow
+  error, as discovered in 
+  [this comment](https://github.com/PennyLaneAI/catalyst/pull/1598#issuecomment-2779178046).
+  [(#1621)](https://github.com/PennyLaneAI/catalyst/pull/2164)
 
 <h3>Documentation 📝</h3>
 
@@ -58,5 +63,6 @@ This release contains contributions from (in alphabetical order):
 
 Ali Asadi,
 Christina Lee,
+River McCubbin,
 Roberto Turrado,
 Paul Haochen Wang.

runtime/include/DataView.hpp

@@ -14,6 +14,8 @@
 
 #pragma once
 
+#include <vector>
+
 #include <Exception.hpp>
 
 /**
@@ -56,39 +58,35 @@ template <typename T, size_t R> class DataView {
         iterator &operator++()
         {
             int64_t next_axis = -1;
-            int64_t idx;
-            for (int64_t i = R; i > 0; --i) {
-                idx = i - 1;
-                RT_ASSERT(view.sizes[idx] > 0);
-                if (indices[idx]++ < view.sizes[idx] - 1) {
-                    next_axis = idx;
+            for (int64_t axis = R - 1; axis >= 0; axis--) {
+                if (++indices[axis] < view.sizes[axis]) {
+                    next_axis = axis;
                     break;
                 }
-                indices[idx] = 0;
-                loc -= (view.sizes[idx] - 1) * view.strides[idx];
+
+                indices[axis] = 0;
+
+                loc -= view.sizes[axis] == 0 ? 0 : (view.sizes[axis] - 1) * view.strides[axis];
             }
 
             loc = next_axis == -1 ? -1 : loc + view.strides[next_axis];
             return *this;
         }
         iterator operator++(int)
         {
-            auto tmp = *this;
+            auto cached_iter = *this;
             int64_t next_axis = -1;
-            int64_t idx;
-            for (int64_t i = R; i > 0; --i) {
-                idx = i - 1;
-                RT_ASSERT(view.sizes[idx] > 0);
-                if (indices[idx]++ < view.sizes[idx] - 1) {
-                    next_axis = idx;
+            for (int64_t axis = R - 1; axis >= 0; axis--) {
+                if (++indices[axis] < view.sizes[axis]) {
+                    next_axis = axis;
                     break;
                 }
-                indices[idx] = 0;
-                loc -= (view.sizes[idx] - 1) * view.strides[idx];
+                indices[axis] = 0;
+                loc -= view.sizes[axis] == 0 ? 0 : (view.sizes[axis] - 1) * view.strides[axis];
             }
 
             loc = next_axis == -1 ? -1 : loc + view.strides[next_axis];
-            return tmp;
+            return cached_iter;
         }
         bool operator==(const iterator &other) const
         {
@@ -138,13 +136,15 @@ template <typename T, size_t R> class DataView {
 
         size_t loc = offset;
         for (size_t axis = 0; axis < R; axis++) {
-            RT_ASSERT(indices[axis] < sizes[axis]);
             loc += indices[axis] * strides[axis];
         }
         return data_aligned[loc];
     }
 
-    iterator begin() { return iterator{*this, static_cast<int64_t>(offset)}; }
+    iterator begin()
+    {
+        return iterator{*this, (*this).size() == 0 ? -1 : static_cast<int64_t>(offset)};
+    }
 
     iterator end() { return iterator{*this, -1}; }
 };

runtime/tests/CMakeLists.txt

@@ -54,6 +54,7 @@ endif()
 
 add_executable(runner_tests_qir_runtime)
 target_sources(runner_tests_qir_runtime PRIVATE
+    Test_DataView.cpp
     Test_NullQubit.cpp
     Test_ResourceTracker.cpp
 )

runtime/tests/Test_DataView.cpp

@@ -0,0 +1,192 @@
+// Copyright 2025 Xanadu Quantum Technologies Inc.
+
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+
+//     http://www.apache.org/licenses/LICENSE-2.0
+
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <cstdio>
+
+#include <catch2/catch_approx.hpp>
+#include <catch2/catch_test_macros.hpp>
+#include <catch2/matchers/catch_matchers_string.hpp>
+
+#include "DataView.hpp"
+
+using namespace Catch::Matchers;
+
+using namespace Catalyst::Runtime;
+
+TEST_CASE("Test DataView Pre-Increment Iterator - double, 1", "[DataView]")
+{
+    double data_aligned[3] = {1.0, 1.1, 1.2};
+    size_t offset = 0U;
+    size_t sizes[1] = {3};
+    size_t strides[1] = {1};
+
+    DataView<double, 1> view(data_aligned, offset, sizes, strides);
+
+    auto view_iter = view.begin();
+
+    CHECK(*view_iter == Catch::Approx(1.0).epsilon(1e-5));
+    CHECK(*++view_iter == Catch::Approx(1.1).epsilon(1e-5));
+    CHECK(*++view_iter == Catch::Approx(1.2).epsilon(1e-5));
+}
+
+TEST_CASE("Test DataView Pre-Increment Iterator - int, 2", "[DataView]")
+{
+    int data_aligned[3][3] = {{0, 1, 2}, {3, 4, 5}};
+    size_t offset = 0U;
+    size_t sizes[2] = {3, 3};
+    size_t strides[2] = {3, 1};
+
+    DataView<int, 2> view(*data_aligned, offset, sizes, strides);
+
+    auto view_iter = view.begin();
+
+    CHECK(*view_iter == 0);
+    for (int i = 1; i < 6; i++) {
+        CHECK(*++view_iter == i);
+    }
+}
+
+TEST_CASE("Test DataView Pre-Increment Iterator - int, 3", "[DataView]")
+{
+    int data_aligned[18] = {0};
+    for (int i = 0; i < 18; i++) {
+        data_aligned[i] = i;
+    }
+
+    size_t offset = 0U;
+    size_t sizes[3] = {3, 2, 3};
+    size_t strides[3] = {6, 3, 1};
+
+    DataView<int, 3> view(data_aligned, offset, sizes, strides);
+
+    auto view_iter = view.begin();
+
+    CHECK(*view_iter == 0);
+    for (int i = 1; i < 18; i++) {
+        CHECK(*++view_iter == i);
+    }
+}
+
+TEST_CASE("Test DataView Post-Increment Iterator - double, 1", "[DataView]")
+{
+    double data_aligned[3] = {3.2, 4.1, 1.6};
+    size_t offset = 0;
+    size_t sizes[1] = {3};
+    size_t strides[1] = {1};
+
+    DataView<double, 1> view(data_aligned, offset, sizes, strides);
+
+    auto view_iter = view.begin();
+
+    CHECK(*view_iter++ == Catch::Approx(3.2).epsilon(1e-5));
+    CHECK(*view_iter++ == Catch::Approx(4.1).epsilon(1e-5));
+    CHECK(*view_iter == Catch::Approx(1.6).epsilon(1e-5));
+}
+
+TEST_CASE("Test DataView Post-Increment Iterator - int, 2", "[DataView]")
+{
+    int data_aligned[3][3] = {{0, 1, 2}, {3, 4, 5}};
+    size_t offset = 0U;
+    size_t sizes[2] = {3, 3};
+    size_t strides[2] = {3, 1};
+
+    DataView<int, 2> view(*data_aligned, offset, sizes, strides);
+
+    auto view_iter = view.begin();
+
+    for (int i = 0; i < 6; i++) {
+        CHECK(*view_iter++ == i);
+    }
+}
+
+TEST_CASE("Test DataView Post-Increment Iterator - int, 3", "[DataView]")
+{
+    int data_aligned[18] = {0};
+    for (int i = 0; i < 18; i++) {
+        data_aligned[i] = i;
+    }
+
+    size_t offset = 0U;
+    size_t sizes[3] = {3, 2, 3};
+    size_t strides[3] = {6, 3, 1};
+
+    DataView<int, 3> view(data_aligned, offset, sizes, strides);
+
+    auto view_iter = view.begin();
+
+    for (int i = 0; i < 18; i++) {
+        CHECK(*view_iter++ == i);
+    }
+}
+
+TEST_CASE("DataView Iterator Distance 0 - 0 first axis", "[DataView]")
+{
+    int *data_aligned = nullptr;
+    size_t offset = 0;
+    size_t sizes[2] = {0, 10};
+    size_t strides[2] = {0, 0};
+
+    DataView<int, 2> view(data_aligned, offset, sizes, strides);
+
+    CHECK(std::distance(view.begin(), view.end()) == 0);
+}
+
+TEST_CASE("DataView Iterator Distance 0 - 0 second axis", "[DataView]")
+{
+    int *data_aligned = nullptr;
+    size_t offset = 0;
+    size_t sizes[2] = {10, 0};
+    size_t strides[2] = {0, 0};
+
+    DataView<int, 2> view(data_aligned, offset, sizes, strides);
+
+    CHECK(std::distance(view.begin(), view.end()) == 0);
+}
+
+TEST_CASE("DataView Iterator Distance 4 - int, 2", "[DataView]")
+{
+    int data_aligned[2][2] = {{1, 2}, {3, 4}};
+    size_t offset = 0;
+    size_t sizes[2] = {2, 2};
+    size_t strides[2] = {2, 1};
+
+    DataView<int, 2> view(*data_aligned, offset, sizes, strides);
+
+    CHECK(std::distance(view.begin(), view.end()) == 4);
+}
+
+TEST_CASE("DataView Iterator Distance 12 - double, 3", "[DataView]")
+{
+    double data_aligned[2][2][3] = {{{3.1, 2.6, 9.5}, {5.4, 2.3, 8.1}},
+                                    {{9.8, 8.2, 7.2}, {0.7, 9.6, 6.6}}};
+    size_t offset = 0;
+    size_t sizes[3] = {2, 2, 3};
+    size_t strides[3] = {6, 3, 1};
+
+    DataView<double, 3> view(**data_aligned, offset, sizes, strides);
+
+    CHECK(std::distance(view.begin(), view.end()) == 12);
+}
+
+TEST_CASE("DataView Size - 0 first axis", "[DataView]")
+{
+    int *data_aligned = nullptr;
+    size_t offset = 0;
+    size_t sizes[2] = {0, 10};
+    size_t strides[2] = {0, 0};
+
+    DataView<int, 2> view(data_aligned, offset, sizes, strides);
+
+    CHECK(view.size() == 0);
+}