xref: /aosp_15_r20/external/pytorch/c10/util/ssize.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <c10/util/Exception.h>
4 #include <c10/util/TypeSafeSignMath.h>
5 
6 #include <cstddef>
7 #include <type_traits>
8 
9 namespace c10 {
10 
11 // Implementations of std::ssize() from C++ 20.
12 //
13 // This is useful in particular for avoiding -Werror=sign-compare
14 // issues.
15 //
16 // Use this with argument-dependent lookup, e.g.:
17 // use c10::ssize;
18 // auto size = ssize(container);
19 //
20 // As with the standard library version, containers are permitted to
21 // specialize this with a free function defined in the same namespace.
22 //
23 // See https://en.cppreference.com/w/cpp/iterator/size for more
24 // information as well as the source of our implementations.
25 //
26 // We augment the implementation by adding an assert() if an overflow
27 // would occur.
28 
29 template <typename C>
30 constexpr auto ssize(const C& c) -> std::
31     common_type_t<std::ptrdiff_t, std::make_signed_t<decltype(c.size())>> {
32   using R = std::
33       common_type_t<std::ptrdiff_t, std::make_signed_t<decltype(c.size())>>;
34   // We expect this to be exceedingly rare to fire and don't wish to
35   // pay a performance hit in release mode.
36   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!greater_than_max<R>(c.size()));
37   return static_cast<R>(c.size());
38 }
39 
40 template <typename T, std::ptrdiff_t N>
41 // NOLINTNEXTLINE(*-c-arrays)
42 constexpr auto ssize(const T (&array)[N]) noexcept -> std::ptrdiff_t {
43   return N;
44 }
45 
46 } // namespace c10
47