From 45c61930139d5543b4072967d8274628c77b782c Mon Sep 17 00:00:00 2001
From: Loic Gouarin <loic.gouarin@gmail.com>
Date: Thu, 29 Oct 2026 15:45:41 +0100
Subject: [PATCH] Fix out-of-bounds write when reusing a cached xfunction shape

xfunction::broadcast_shape copied its cached shape into the output shape
without checking that both have the same rank. Assigning a rank-2
xfunction whose shape was already computed to an xt::xtensor<T, 1> wrote
past the std::array<std::size_t, 1> (stack-buffer-overflow under ASan).
With a larger dynamic output shape, the cached path returned a different
shape than the uncached one.

The cache is now used only when both ranks match. Otherwise the full
broadcast runs and reports incompatible ranks, as it already did when the
cache was empty.
---
 include/xtensor/core/xfunction.hpp |  4 +++-
 test/test_xfunction.cpp            | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/include/xtensor/core/xfunction.hpp b/include/xtensor/core/xfunction.hpp
index 99144b377..38982e479 100644
--- a/include/xtensor/core/xfunction.hpp
+++ b/include/xtensor/core/xfunction.hpp
@@ -710,7 +710,9 @@ namespace xt
     template <class S>
     inline bool xfunction<F, CT...>::broadcast_shape(S& shape, bool reuse_cache) const
     {
-        if (m_cache.is_initialized && reuse_cache)
+        // The cached shape can only be copied into a shape of the same rank; any other
+        // rank goes through the full broadcast, which also reports incompatible ranks.
+        if (m_cache.is_initialized && reuse_cache && shape.size() == m_cache.shape.size())
         {
             std::copy(m_cache.shape.cbegin(), m_cache.shape.cend(), shape.begin());
             return m_cache.is_trivial;
diff --git a/test/test_xfunction.cpp b/test/test_xfunction.cpp
index ddf08758c..a9f2dccef 100644
--- a/test/test_xfunction.cpp
+++ b/test/test_xfunction.cpp
@@ -110,6 +110,31 @@ namespace xt
         XT_EXPECT_ANY_THROW(xt::xarray<double> res = arr1 * arr2);
     }
 
+    TEST(xfunction, broadcast_shape_cached_other_rank)
+    {
+        xt::xarray<double> a = {{1.0, 2.0, 3.0}, {4.0, 5.0, 6.0}};
+        auto f = a + 1.0;
+        // Fills the shape cache of f
+        EXPECT_EQ(f.shape().size(), 2u);
+
+        SUBCASE("higher rank")
+        {
+            using shape_type = xt::dynamic_shape<std::size_t>;
+            shape_type cached = uninitialized_shape<shape_type>(3);
+            shape_type computed = uninitialized_shape<shape_type>(3);
+            bool cached_trivial = f.broadcast_shape(cached, true);
+            bool computed_trivial = f.broadcast_shape(computed, false);
+            EXPECT_EQ(cached, computed);
+            EXPECT_EQ(cached_trivial, computed_trivial);
+        }
+
+        SUBCASE("lower rank")
+        {
+            xt::xtensor<double, 1> res;
+            XT_EXPECT_ANY_THROW(res = f);
+        }
+    }
+
     TEST(xfunction, shape)
     {
         xfunction_features f;