1// -*- C++ -*-
2//===----------------------------------------------------------------------===//
3//
4// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5// See https://llvm.org/LICENSE.txt for license information.
6// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7//
8//                        Kokkos v. 4.0
9//       Copyright (2022) National Technology & Engineering
10//               Solutions of Sandia, LLC (NTESS).
11//
12// Under the terms of Contract DE-NA0003525 with NTESS,
13// the U.S. Government retains certain rights in this software.
14//
15//===---------------------------------------------------------------------===//
16
17#ifndef _LIBCPP___MDSPAN_MDSPAN_H
18#define _LIBCPP___MDSPAN_MDSPAN_H
19
20#include <__assert>
21#include <__config>
22#include <__fwd/mdspan.h>
23#include <__mdspan/default_accessor.h>
24#include <__mdspan/extents.h>
25#include <__type_traits/extent.h>
26#include <__type_traits/is_abstract.h>
27#include <__type_traits/is_array.h>
28#include <__type_traits/is_constructible.h>
29#include <__type_traits/is_convertible.h>
30#include <__type_traits/is_default_constructible.h>
31#include <__type_traits/is_nothrow_constructible.h>
32#include <__type_traits/is_pointer.h>
33#include <__type_traits/is_same.h>
34#include <__type_traits/rank.h>
35#include <__type_traits/remove_all_extents.h>
36#include <__type_traits/remove_cv.h>
37#include <__type_traits/remove_pointer.h>
38#include <__type_traits/remove_reference.h>
39#include <__utility/integer_sequence.h>
40#include <array>
41#include <cinttypes>
42#include <cstddef>
43#include <limits>
44#include <span>
45
46#if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
47#  pragma GCC system_header
48#endif
49
50_LIBCPP_PUSH_MACROS
51#include <__undef_macros>
52
53_LIBCPP_BEGIN_NAMESPACE_STD
54
55#if _LIBCPP_STD_VER >= 23
56
57// Helper for lightweight test checking that one did pass a layout policy as LayoutPolicy template argument
58namespace __mdspan_detail {
59template <class _Layout, class _Extents>
60concept __has_invalid_mapping = !requires { typename _Layout::template mapping<_Extents>; };
61} // namespace __mdspan_detail
62
63template <class _ElementType,
64          class _Extents,
65          class _LayoutPolicy   = layout_right,
66          class _AccessorPolicy = default_accessor<_ElementType> >
67class mdspan {
68private:
69  static_assert(__mdspan_detail::__is_extents_v<_Extents>,
70                "mdspan: Extents template parameter must be a specialization of extents.");
71  static_assert(!is_array_v<_ElementType>, "mdspan: ElementType template parameter may not be an array type");
72  static_assert(!is_abstract_v<_ElementType>, "mdspan: ElementType template parameter may not be an abstract class");
73  static_assert(is_same_v<_ElementType, typename _AccessorPolicy::element_type>,
74                "mdspan: ElementType template parameter must match AccessorPolicy::element_type");
75  static_assert(!__mdspan_detail::__has_invalid_mapping<_LayoutPolicy, _Extents>,
76                "mdspan: LayoutPolicy template parameter is invalid. A common mistake is to pass a layout mapping "
77                "instead of a layout policy");
78
79public:
80  using extents_type     = _Extents;
81  using layout_type      = _LayoutPolicy;
82  using accessor_type    = _AccessorPolicy;
83  using mapping_type     = typename layout_type::template mapping<extents_type>;
84  using element_type     = _ElementType;
85  using value_type       = remove_cv_t<element_type>;
86  using index_type       = typename extents_type::index_type;
87  using size_type        = typename extents_type::size_type;
88  using rank_type        = typename extents_type::rank_type;
89  using data_handle_type = typename accessor_type::data_handle_type;
90  using reference        = typename accessor_type::reference;
91
92  _LIBCPP_HIDE_FROM_ABI static constexpr rank_type rank() noexcept { return extents_type::rank(); }
93  _LIBCPP_HIDE_FROM_ABI static constexpr rank_type rank_dynamic() noexcept { return extents_type::rank_dynamic(); }
94  _LIBCPP_HIDE_FROM_ABI static constexpr size_t static_extent(rank_type __r) noexcept {
95    return extents_type::static_extent(__r);
96  }
97  _LIBCPP_HIDE_FROM_ABI constexpr index_type extent(rank_type __r) const noexcept {
98    return __map_.extents().extent(__r);
99  };
100
101public:
102  //--------------------------------------------------------------------------------
103  // [mdspan.mdspan.cons], mdspan constructors, assignment, and destructor
104
105  _LIBCPP_HIDE_FROM_ABI constexpr mdspan()
106    requires((extents_type::rank_dynamic() > 0) && is_default_constructible_v<data_handle_type> &&
107             is_default_constructible_v<mapping_type> && is_default_constructible_v<accessor_type>)
108  = default;
109  _LIBCPP_HIDE_FROM_ABI constexpr mdspan(const mdspan&) = default;
110  _LIBCPP_HIDE_FROM_ABI constexpr mdspan(mdspan&&)      = default;
111
112  template <class... _OtherIndexTypes>
113    requires((is_convertible_v<_OtherIndexTypes, index_type> && ...) &&
114             (is_nothrow_constructible_v<index_type, _OtherIndexTypes> && ...) &&
115             ((sizeof...(_OtherIndexTypes) == rank()) || (sizeof...(_OtherIndexTypes) == rank_dynamic())) &&
116             is_constructible_v<mapping_type, extents_type> && is_default_constructible_v<accessor_type>)
117  _LIBCPP_HIDE_FROM_ABI explicit constexpr mdspan(data_handle_type __p, _OtherIndexTypes... __exts)
118      : __ptr_(std::move(__p)), __map_(extents_type(static_cast<index_type>(std::move(__exts))...)), __acc_{} {}
119
120  template <class _OtherIndexType, size_t _Size>
121    requires(is_convertible_v<const _OtherIndexType&, index_type> &&
122             is_nothrow_constructible_v<index_type, const _OtherIndexType&> &&
123             ((_Size == rank()) || (_Size == rank_dynamic())) && is_constructible_v<mapping_type, extents_type> &&
124             is_default_constructible_v<accessor_type>)
125  explicit(_Size != rank_dynamic())
126      _LIBCPP_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const array<_OtherIndexType, _Size>& __exts)
127      : __ptr_(std::move(__p)), __map_(extents_type(__exts)), __acc_{} {}
128
129  template <class _OtherIndexType, size_t _Size>
130    requires(is_convertible_v<const _OtherIndexType&, index_type> &&
131             is_nothrow_constructible_v<index_type, const _OtherIndexType&> &&
132             ((_Size == rank()) || (_Size == rank_dynamic())) && is_constructible_v<mapping_type, extents_type> &&
133             is_default_constructible_v<accessor_type>)
134  explicit(_Size != rank_dynamic())
135      _LIBCPP_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, span<_OtherIndexType, _Size> __exts)
136      : __ptr_(std::move(__p)), __map_(extents_type(__exts)), __acc_{} {}
137
138  _LIBCPP_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const extents_type& __exts)
139    requires(is_default_constructible_v<accessor_type> && is_constructible_v<mapping_type, const extents_type&>)
140      : __ptr_(std::move(__p)), __map_(__exts), __acc_{} {}
141
142  _LIBCPP_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const mapping_type& __m)
143    requires(is_default_constructible_v<accessor_type>)
144      : __ptr_(std::move(__p)), __map_(__m), __acc_{} {}
145
146  _LIBCPP_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const mapping_type& __m, const accessor_type& __a)
147      : __ptr_(std::move(__p)), __map_(__m), __acc_(__a) {}
148
149  template <class _OtherElementType, class _OtherExtents, class _OtherLayoutPolicy, class _OtherAccessor>
150    requires(is_constructible_v<mapping_type, const typename _OtherLayoutPolicy::template mapping<_OtherExtents>&> &&
151             is_constructible_v<accessor_type, const _OtherAccessor&>)
152  explicit(!is_convertible_v<const typename _OtherLayoutPolicy::template mapping<_OtherExtents>&, mapping_type> ||
153           !is_convertible_v<const _OtherAccessor&, accessor_type>)
154      _LIBCPP_HIDE_FROM_ABI constexpr mdspan(
155          const mdspan<_OtherElementType, _OtherExtents, _OtherLayoutPolicy, _OtherAccessor>& __other)
156      : __ptr_(__other.__ptr_), __map_(__other.__map_), __acc_(__other.__acc_) {
157    static_assert(is_constructible_v<data_handle_type, const typename _OtherAccessor::data_handle_type&>,
158                  "mdspan: incompatible data_handle_type for mdspan construction");
159    static_assert(
160        is_constructible_v<extents_type, _OtherExtents>, "mdspan: incompatible extents for mdspan construction");
161
162    // The following precondition is part of the standard, but is unlikely to be triggered.
163    // The extents constructor checks this and the mapping must be storing the extents, since
164    // its extents() function returns a const reference to extents_type.
165    // The only way this can be triggered is if the mapping conversion constructor would for example
166    // always construct its extents() only from the dynamic extents, instead of from the other extents.
167    if constexpr (rank() > 0) {
168      for (size_t __r = 0; __r < rank(); __r++) {
169        // Not catching this could lead to out of bounds errors later
170        // e.g. mdspan<int, dextents<char,1>, non_checking_layout> m =
171        //        mdspan<int, dextents<unsigned, 1>, non_checking_layout>(ptr, 200); leads to an extent of -56 on m
172        _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
173            (static_extent(__r) == dynamic_extent) ||
174                (static_cast<index_type>(__other.extent(__r)) == static_cast<index_type>(static_extent(__r))),
175            "mdspan: conversion mismatch of source dynamic extents with static extents");
176      }
177    }
178  }
179
180  _LIBCPP_HIDE_FROM_ABI constexpr mdspan& operator=(const mdspan&) = default;
181  _LIBCPP_HIDE_FROM_ABI constexpr mdspan& operator=(mdspan&&)      = default;
182
183  //--------------------------------------------------------------------------------
184  // [mdspan.mdspan.members], members
185
186  template <class... _OtherIndexTypes>
187    requires((is_convertible_v<_OtherIndexTypes, index_type> && ...) &&
188             (is_nothrow_constructible_v<index_type, _OtherIndexTypes> && ...) &&
189             (sizeof...(_OtherIndexTypes) == rank()))
190  _LIBCPP_HIDE_FROM_ABI constexpr reference operator[](_OtherIndexTypes... __indices) const {
191    // Note the standard layouts would also check this, but user provided ones may not, so we
192    // check the precondition here
193    _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(__mdspan_detail::__is_multidimensional_index_in(extents(), __indices...),
194                                        "mdspan: operator[] out of bounds access");
195    return __acc_.access(__ptr_, __map_(static_cast<index_type>(std::move(__indices))...));
196  }
197
198  template <class _OtherIndexType>
199    requires(is_convertible_v<const _OtherIndexType&, index_type> &&
200             is_nothrow_constructible_v<index_type, const _OtherIndexType&>)
201  _LIBCPP_HIDE_FROM_ABI constexpr reference operator[](const array< _OtherIndexType, rank()>& __indices) const {
202    return __acc_.access(__ptr_, [&]<size_t... _Idxs>(index_sequence<_Idxs...>) {
203      return __map_(__indices[_Idxs]...);
204    }(make_index_sequence<rank()>()));
205  }
206
207  template <class _OtherIndexType>
208    requires(is_convertible_v<const _OtherIndexType&, index_type> &&
209             is_nothrow_constructible_v<index_type, const _OtherIndexType&>)
210  _LIBCPP_HIDE_FROM_ABI constexpr reference operator[](span<_OtherIndexType, rank()> __indices) const {
211    return __acc_.access(__ptr_, [&]<size_t... _Idxs>(index_sequence<_Idxs...>) {
212      return __map_(__indices[_Idxs]...);
213    }(make_index_sequence<rank()>()));
214  }
215
216  _LIBCPP_HIDE_FROM_ABI constexpr size_type size() const noexcept {
217    // Could leave this as only checked in debug mode: semantically size() is never
218    // guaranteed to be related to any accessible range
219    _LIBCPP_ASSERT_UNCATEGORIZED(
220        false == ([&]<size_t... _Idxs>(index_sequence<_Idxs...>) {
221          size_type __prod = 1;
222          return (__builtin_mul_overflow(__prod, extent(_Idxs), &__prod) || ... || false);
223        }(make_index_sequence<rank()>())),
224        "mdspan: size() is not representable as size_type");
225    return [&]<size_t... _Idxs>(index_sequence<_Idxs...>) {
226      return ((static_cast<size_type>(__map_.extents().extent(_Idxs))) * ... * size_type(1));
227    }(make_index_sequence<rank()>());
228  }
229
230  [[nodiscard]] _LIBCPP_HIDE_FROM_ABI constexpr bool empty() const noexcept {
231    return [&]<size_t... _Idxs>(index_sequence<_Idxs...>) {
232      return (rank() > 0) && ((__map_.extents().extent(_Idxs) == index_type(0)) || ... || false);
233    }(make_index_sequence<rank()>());
234  }
235
236  _LIBCPP_HIDE_FROM_ABI friend constexpr void swap(mdspan& __x, mdspan& __y) noexcept {
237    swap(__x.__ptr_, __y.__ptr_);
238    swap(__x.__map_, __y.__map_);
239    swap(__x.__acc_, __y.__acc_);
240  }
241
242  _LIBCPP_HIDE_FROM_ABI constexpr const extents_type& extents() const noexcept { return __map_.extents(); };
243  _LIBCPP_HIDE_FROM_ABI constexpr const data_handle_type& data_handle() const noexcept { return __ptr_; };
244  _LIBCPP_HIDE_FROM_ABI constexpr const mapping_type& mapping() const noexcept { return __map_; };
245  _LIBCPP_HIDE_FROM_ABI constexpr const accessor_type& accessor() const noexcept { return __acc_; };
246
247  // per LWG-4021 "mdspan::is_always_meow() should be noexcept"
248  _LIBCPP_HIDE_FROM_ABI static constexpr bool is_always_unique() noexcept { return mapping_type::is_always_unique(); };
249  _LIBCPP_HIDE_FROM_ABI static constexpr bool is_always_exhaustive() noexcept {
250    return mapping_type::is_always_exhaustive();
251  };
252  _LIBCPP_HIDE_FROM_ABI static constexpr bool is_always_strided() noexcept {
253    return mapping_type::is_always_strided();
254  };
255
256  _LIBCPP_HIDE_FROM_ABI constexpr bool is_unique() const { return __map_.is_unique(); };
257  _LIBCPP_HIDE_FROM_ABI constexpr bool is_exhaustive() const { return __map_.is_exhaustive(); };
258  _LIBCPP_HIDE_FROM_ABI constexpr bool is_strided() const { return __map_.is_strided(); };
259  _LIBCPP_HIDE_FROM_ABI constexpr index_type stride(rank_type __r) const { return __map_.stride(__r); };
260
261private:
262  _LIBCPP_NO_UNIQUE_ADDRESS data_handle_type __ptr_{};
263  _LIBCPP_NO_UNIQUE_ADDRESS mapping_type __map_{};
264  _LIBCPP_NO_UNIQUE_ADDRESS accessor_type __acc_{};
265
266  template <class, class, class, class>
267  friend class mdspan;
268};
269
270template <class _ElementType, class... _OtherIndexTypes>
271  requires((is_convertible_v<_OtherIndexTypes, size_t> && ...) && (sizeof...(_OtherIndexTypes) > 0))
272explicit mdspan(_ElementType*, _OtherIndexTypes...)
273    -> mdspan<_ElementType, dextents<size_t, sizeof...(_OtherIndexTypes)>>;
274
275template <class _Pointer>
276  requires(is_pointer_v<remove_reference_t<_Pointer>>)
277mdspan(_Pointer&&) -> mdspan<remove_pointer_t<remove_reference_t<_Pointer>>, extents<size_t>>;
278
279template <class _CArray>
280  requires(is_array_v<_CArray> && (rank_v<_CArray> == 1))
281mdspan(_CArray&) -> mdspan<remove_all_extents_t<_CArray>, extents<size_t, extent_v<_CArray, 0>>>;
282
283template <class _ElementType, class _OtherIndexType, size_t _Size>
284mdspan(_ElementType*, const array<_OtherIndexType, _Size>&) -> mdspan<_ElementType, dextents<size_t, _Size>>;
285
286template <class _ElementType, class _OtherIndexType, size_t _Size>
287mdspan(_ElementType*, span<_OtherIndexType, _Size>) -> mdspan<_ElementType, dextents<size_t, _Size>>;
288
289// This one is necessary because all the constructors take `data_handle_type`s, not
290// `_ElementType*`s, and `data_handle_type` is taken from `accessor_type::data_handle_type`, which
291// seems to throw off automatic deduction guides.
292template <class _ElementType, class _OtherIndexType, size_t... _ExtentsPack>
293mdspan(_ElementType*, const extents<_OtherIndexType, _ExtentsPack...>&)
294    -> mdspan<_ElementType, extents<_OtherIndexType, _ExtentsPack...>>;
295
296template <class _ElementType, class _MappingType>
297mdspan(_ElementType*, const _MappingType&)
298    -> mdspan<_ElementType, typename _MappingType::extents_type, typename _MappingType::layout_type>;
299
300template <class _MappingType, class _AccessorType>
301mdspan(const typename _AccessorType::data_handle_type, const _MappingType&, const _AccessorType&)
302    -> mdspan<typename _AccessorType::element_type,
303              typename _MappingType::extents_type,
304              typename _MappingType::layout_type,
305              _AccessorType>;
306
307#endif // _LIBCPP_STD_VER >= 23
308
309_LIBCPP_END_NAMESPACE_STD
310
311_LIBCPP_POP_MACROS
312
313#endif // _LIBCPP___MDSPAN_MDSPAN_H
314