From 4026ec1d89a5c2bd721368cd9a51e781763e0bb6 Mon Sep 17 00:00:00 2001 From: Loic Gouarin Date: Thu, 29 Oct 2026 18:55:27 +0100 Subject: [PATCH] Fix concatenate and stack access with a different number of indices Containers accept fewer indices than dimensions (the missing leading ones are 0) or more (the extra leading ones are ignored). concatenate_access computed the offset into the index range as distance(first, last) - dimension() in size_t. With fewer indices this wrapped around, so concatenate and hstack read before the start of the index array: c(1) returned garbage for a 2-D concatenation and hstack crashed. stack_access ignored the index count entirely, so stack and 1-D vstack returned wrong values with both fewer and more indices. Both access methods now go through detail::trailing_index, which aligns the indices on the last dimensions like containers do. It also converts the index values explicitly: flat_expression_adaptor passes the signed indices returned by unravel_index, which are never negative. This removes the -Wsign-conversion warnings in xbuilder.hpp. The new test fails on master and passes with ASan and UBSan. --- include/xtensor/generators/xbuilder.hpp | 88 ++++++++++++++----------- test/test_xbuilder.cpp | 33 ++++++++++ 2 files changed, 83 insertions(+), 38 deletions(-) diff --git a/include/xtensor/generators/xbuilder.hpp b/include/xtensor/generators/xbuilder.hpp index 1026b21c7..a96894542 100644 --- a/include/xtensor/generators/xbuilder.hpp +++ b/include/xtensor/generators/xbuilder.hpp @@ -19,6 +19,7 @@ #include #include #include +#include #include #include @@ -485,6 +486,28 @@ namespace xt namespace detail { + /** + * Returns the index along dimension @p d of a @p dim-dimensional expression, + * read from the indices in [first, last). + * + * As for containers, the indices are aligned on the last dimensions: extra + * leading indices are ignored and missing leading indices are 0. The indices + * may have a signed type (flat_expression_adaptor passes the result of + * unravel_index) but they are never negative. + */ + template + inline std::size_t trailing_index(It first, It last, std::size_t dim, std::size_t d) + { + using difference_type = typename std::iterator_traits::difference_type; + // d < dim, so the distance from the end is in [1, dim]. + const auto from_end = static_cast(dim - d); + if (std::distance(first, last) < from_end) + { + return 0; + } + return static_cast(*(last - from_end)); + } + template class concatenate_access { @@ -497,9 +520,8 @@ namespace xt template inline value_type access(const tuple_type& t, size_type axis, It first, It last) const { - // trim off extra indices if provided to match behavior of containers - auto dim_offset = std::distance(first, last) - std::get<0>(t).dimension(); - std::size_t axis_dim = *(first + axis + dim_offset); + const size_type dim = std::get<0>(t).dimension(); + size_type axis_dim = trailing_index(first, last, dim, axis); auto match = [&](auto& arr) { if (axis_dim >= arr.shape()[axis]) @@ -512,28 +534,21 @@ namespace xt auto get = [&](auto& arr) { - std::size_t offset = 0; - const std::size_t end = arr.dimension(); - for (std::size_t i = 0; i < end; i++) + size_type offset = 0; + const auto& shape = arr.shape(); + for (size_type i = 0; i < dim; i++) { - const auto& shape = arr.shape(); - const std::size_t stride = std::accumulate( - shape.begin() + i + 1, + const size_type stride = std::accumulate( + shape.begin() + static_cast(i) + 1, shape.end(), - size_t(1), - std::multiplies() + size_type(1), + std::multiplies() ); - if (i == axis) - { - offset += axis_dim * stride; - } - else - { - const auto len = (*(first + i + dim_offset)); - offset += len * stride; - } + const size_type index = i == axis ? axis_dim : trailing_index(first, last, dim, i); + offset += index * stride; } - const auto element = arr.begin() + offset; + // offset < arr.size(), so it fits in a ptrdiff_t. + const auto element = arr.begin() + static_cast(offset); return *element; }; @@ -559,34 +574,31 @@ namespace xt using value_type = xtl::promote_type_t::value_type...>; template - inline value_type access(const tuple_type& t, size_type axis, It first, It) const + inline value_type access(const tuple_type& t, size_type axis, It first, It last) const { + // The stacked expression has one more dimension than its operands. + const size_type dim = std::get<0>(t).dimension() + 1; auto get_item = [&](auto& arr) { - std::size_t offset = 0; - const std::size_t end = arr.dimension(); - std::size_t after_axis = 0; - for (std::size_t i = 0; i < end; i++) + size_type offset = 0; + const auto& shape = arr.shape(); + for (size_type i = 0; i < arr.dimension(); i++) { - if (i == axis) - { - after_axis = 1; - } - const auto& shape = arr.shape(); - const std::size_t stride = std::accumulate( + const size_type stride = std::accumulate( shape.begin() + static_cast(i) + 1, shape.end(), - std::size_t(1), - std::multiplies() + size_type(1), + std::multiplies() ); - const auto len = (*(first + static_cast(i + after_axis))); - offset += len * stride; + // Dimensions of arr after the stacking axis are shifted by one. + const size_type index = trailing_index(first, last, dim, i < axis ? i : i + 1); + offset += index * stride; } + // offset < arr.size(), so it fits in a ptrdiff_t. const auto element = arr.begin() + static_cast(offset); return *element; }; - size_type i = *(first + static_cast(axis)); - return apply(i, get_item, t); + return apply(trailing_index(first, last, dim, axis), get_item, t); } }; diff --git a/test/test_xbuilder.cpp b/test/test_xbuilder.cpp index d35847a8e..b4a12a735 100644 --- a/test/test_xbuilder.cpp +++ b/test/test_xbuilder.cpp @@ -532,6 +532,39 @@ namespace xt EXPECT_EQ(c2, e2); } + // Like containers, builders accept fewer indices than dimensions (missing + // leading ones are 0) or more (extra leading ones are ignored). + TEST(xbuilder, concatenate_stack_index_count) + { + xarray a = {{1, 2, 3}, {4, 5, 6}}; + xarray b = {{7, 8, 9}}; + + auto c = concatenate(xtuple(a, b), 0); + xarray ec = c; + EXPECT_EQ(c(1), ec(1)); + EXPECT_EQ(c(5, 2, 1), ec(5, 2, 1)); + + auto h = hstack(xtuple(a, a)); + xarray eh = h; + EXPECT_EQ(h(4), eh(4)); + EXPECT_EQ(h(5, 1, 4), eh(5, 1, 4)); + + auto s = stack(xtuple(a, a), 1); + xarray es = s; + EXPECT_EQ(s(2), es(2)); + EXPECT_EQ(s(1, 2), es(1, 2)); + EXPECT_EQ(s(5, 1, 1, 2), es(5, 1, 1, 2)); + + auto v = vstack(xtuple(xarray{1, 2, 3}, xarray{4, 5, 6})); + xarray ev = v; + EXPECT_EQ(v(2), ev(2)); + EXPECT_EQ(v(5, 1, 2), ev(5, 1, 2)); + + // Signed indices, as passed by flat_expression_adaptor + const std::array idx = {2, 1}; + EXPECT_EQ(c.element(idx.cbegin(), idx.cend()), ec(2, 1)); + } + TEST(xbuilder, meshgrid) { auto mesh = meshgrid(linspace(0.0, 1.0, 3), linspace(0.0, 1.0, 2));