From 4b3dfd134f39b702e617c0816981a990e89b2b0f Mon Sep 17 00:00:00 2001
From: Tom de Geus <tom@geus.me>
Date: Wed, 19 May 2021 16:01:50 +0200
Subject: [PATCH] Adding size assertion `.flat(i)` + adding a few tests on size
 assertions

---
 include/xtensor/xcontainer.hpp |  2 ++
 test/test_xarray.cpp           | 21 ++++++++++++++++++++-
 test/test_xtensor.cpp          | 19 +++++++++++++++++++
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/include/xtensor/xcontainer.hpp b/include/xtensor/xcontainer.hpp
index cf5e52cdf..933801cb0 100644
--- a/include/xtensor/xcontainer.hpp
+++ b/include/xtensor/xcontainer.hpp
@@ -651,12 +651,14 @@ namespace xt
     template <class D>
     inline auto xcontainer<D>::flat(size_type i) -> reference
     {
+        XTENSOR_ASSERT(i < size());
         return storage()[i];
     }
 
     template <class D>
     inline auto xcontainer<D>::flat(size_type i) const -> const_reference
     {
+        XTENSOR_ASSERT(i < size());
         return storage()[i];
     }
 
diff --git a/test/test_xarray.cpp b/test/test_xarray.cpp
index 81724d940..962f68f19 100644
--- a/test/test_xarray.cpp
+++ b/test/test_xarray.cpp
@@ -314,6 +314,19 @@ namespace xt
         EXPECT_EQ(a.layout(), c.layout());
     }
 
+    TEST(xarray, operator_brace)
+    {
+#ifdef XTENSOR_ENABLE_ASSERT
+        xt::xarray<size_t> a = {{0,1,2}, {3,4,5}};
+        EXPECT_THROW(a(2, 0), std::runtime_error);
+        EXPECT_THROW(a(0, 3), std::runtime_error);
+
+        xt::xarray<size_t> b = {{0,1,2}};
+        EXPECT_THROW(a(0, 3), std::runtime_error);
+        EXPECT_THROW(a(1, 3), std::runtime_error);
+#endif
+    }
+
     TEST(xarray, periodic)
     {
         xt::xarray<size_t> a = {{0,1,2}, {3,4,5}};
@@ -336,7 +349,7 @@ namespace xt
         EXPECT_EQ(a.back(), 6);
     }
 
-    TEST(xarray, flat)
+TEST(xarray, flat)
     {
         {
             xt::xarray<size_t, xt::layout_type::row_major> a = {{0,1,2}, {3,4,5}};
@@ -345,6 +358,9 @@ namespace xt
             a.flat(4) = 40;
             a.flat(5) = 50;
             EXPECT_EQ(a, b);
+#ifdef XTENSOR_ENABLE_ASSERT
+            EXPECT_THROW(a.flat(6), std::runtime_error);
+#endif
         }
         {
             xt::xarray<size_t, xt::layout_type::column_major> a = {{0,1,2}, {3,4,5}};
@@ -353,6 +369,9 @@ namespace xt
             a.flat(3) = 40;
             a.flat(5) = 50;
             EXPECT_EQ(a, b);
+#ifdef XTENSOR_ENABLE_ASSERT
+            EXPECT_THROW(a.flat(6), std::runtime_error);
+#endif
         }
     }
 
diff --git a/test/test_xtensor.cpp b/test/test_xtensor.cpp
index 859b061f7..482cd0939 100644
--- a/test/test_xtensor.cpp
+++ b/test/test_xtensor.cpp
@@ -370,6 +370,19 @@ namespace xt
         EXPECT_TRUE(idx==jdx);
     }
 
+    TEST(xarray, operator_brace)
+    {
+#ifdef XTENSOR_ENABLE_ASSERT
+        xt::xtensor<size_t, 2> a = {{0,1,2}, {3,4,5}};
+        EXPECT_THROW(a(2, 0), std::runtime_error);
+        EXPECT_THROW(a(0, 3), std::runtime_error);
+
+        xt::xtensor<size_t, 2> b = {{0,1,2}};
+        EXPECT_THROW(a(0, 3), std::runtime_error);
+        EXPECT_THROW(a(1, 3), std::runtime_error);
+#endif
+    }
+
     TEST(xtensor, periodic)
     {
         xt::xtensor<size_t,2> a = {{0,1,2}, {3,4,5}};
@@ -389,6 +402,9 @@ namespace xt
             a.flat(4) = 40;
             a.flat(5) = 50;
             EXPECT_EQ(a, b);
+#ifdef XTENSOR_ENABLE_ASSERT
+            EXPECT_THROW(a.flat(6), std::runtime_error);
+#endif
         }
         {
             xt::xtensor<size_t, 2, xt::layout_type::column_major> a = {{0,1,2}, {3,4,5}};
@@ -397,6 +413,9 @@ namespace xt
             a.flat(3) = 40;
             a.flat(5) = 50;
             EXPECT_EQ(a, b);
+#ifdef XTENSOR_ENABLE_ASSERT
+            EXPECT_THROW(a.flat(6), std::runtime_error);
+#endif
         }
     }