111986Sandreas.sandberg@arm.com/*
211986Sandreas.sandberg@arm.com    pybind11/numpy.h: Basic NumPy support, vectorize() wrapper
311986Sandreas.sandberg@arm.com
411986Sandreas.sandberg@arm.com    Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
511986Sandreas.sandberg@arm.com
611986Sandreas.sandberg@arm.com    All rights reserved. Use of this source code is governed by a
711986Sandreas.sandberg@arm.com    BSD-style license that can be found in the LICENSE file.
811986Sandreas.sandberg@arm.com*/
911986Sandreas.sandberg@arm.com
1011986Sandreas.sandberg@arm.com#pragma once
1111986Sandreas.sandberg@arm.com
1211986Sandreas.sandberg@arm.com#include "pybind11.h"
1311986Sandreas.sandberg@arm.com#include "complex.h"
1411986Sandreas.sandberg@arm.com#include <numeric>
1511986Sandreas.sandberg@arm.com#include <algorithm>
1611986Sandreas.sandberg@arm.com#include <array>
1714299Sbbruce@ucdavis.edu#include <cstdint>
1811986Sandreas.sandberg@arm.com#include <cstdlib>
1911986Sandreas.sandberg@arm.com#include <cstring>
2011986Sandreas.sandberg@arm.com#include <sstream>
2111986Sandreas.sandberg@arm.com#include <string>
2211986Sandreas.sandberg@arm.com#include <functional>
2311986Sandreas.sandberg@arm.com#include <utility>
2414299Sbbruce@ucdavis.edu#include <vector>
2511986Sandreas.sandberg@arm.com#include <typeindex>
2611986Sandreas.sandberg@arm.com
2711986Sandreas.sandberg@arm.com#if defined(_MSC_VER)
2811986Sandreas.sandberg@arm.com#  pragma warning(push)
2911986Sandreas.sandberg@arm.com#  pragma warning(disable: 4127) // warning C4127: Conditional expression is constant
3011986Sandreas.sandberg@arm.com#endif
3111986Sandreas.sandberg@arm.com
3211986Sandreas.sandberg@arm.com/* This will be true on all flat address space platforms and allows us to reduce the
3312391Sjason@lowepower.com   whole npy_intp / ssize_t / Py_intptr_t business down to just ssize_t for all size
3411986Sandreas.sandberg@arm.com   and dimension types (e.g. shape, strides, indexing), instead of inflicting this
3511986Sandreas.sandberg@arm.com   upon the library user. */
3612391Sjason@lowepower.comstatic_assert(sizeof(ssize_t) == sizeof(Py_intptr_t), "ssize_t != Py_intptr_t");
3711986Sandreas.sandberg@arm.com
3812391Sjason@lowepower.comNAMESPACE_BEGIN(PYBIND11_NAMESPACE)
3912037Sandreas.sandberg@arm.com
4012037Sandreas.sandberg@arm.comclass array; // Forward declaration
4112037Sandreas.sandberg@arm.com
4211986Sandreas.sandberg@arm.comNAMESPACE_BEGIN(detail)
4312037Sandreas.sandberg@arm.comtemplate <typename type, typename SFINAE = void> struct npy_format_descriptor;
4411986Sandreas.sandberg@arm.com
4511986Sandreas.sandberg@arm.comstruct PyArrayDescr_Proxy {
4611986Sandreas.sandberg@arm.com    PyObject_HEAD
4711986Sandreas.sandberg@arm.com    PyObject *typeobj;
4811986Sandreas.sandberg@arm.com    char kind;
4911986Sandreas.sandberg@arm.com    char type;
5011986Sandreas.sandberg@arm.com    char byteorder;
5111986Sandreas.sandberg@arm.com    char flags;
5211986Sandreas.sandberg@arm.com    int type_num;
5311986Sandreas.sandberg@arm.com    int elsize;
5411986Sandreas.sandberg@arm.com    int alignment;
5511986Sandreas.sandberg@arm.com    char *subarray;
5611986Sandreas.sandberg@arm.com    PyObject *fields;
5711986Sandreas.sandberg@arm.com    PyObject *names;
5811986Sandreas.sandberg@arm.com};
5911986Sandreas.sandberg@arm.com
6011986Sandreas.sandberg@arm.comstruct PyArray_Proxy {
6111986Sandreas.sandberg@arm.com    PyObject_HEAD
6211986Sandreas.sandberg@arm.com    char *data;
6311986Sandreas.sandberg@arm.com    int nd;
6411986Sandreas.sandberg@arm.com    ssize_t *dimensions;
6511986Sandreas.sandberg@arm.com    ssize_t *strides;
6611986Sandreas.sandberg@arm.com    PyObject *base;
6711986Sandreas.sandberg@arm.com    PyObject *descr;
6811986Sandreas.sandberg@arm.com    int flags;
6911986Sandreas.sandberg@arm.com};
7011986Sandreas.sandberg@arm.com
7111986Sandreas.sandberg@arm.comstruct PyVoidScalarObject_Proxy {
7211986Sandreas.sandberg@arm.com    PyObject_VAR_HEAD
7311986Sandreas.sandberg@arm.com    char *obval;
7411986Sandreas.sandberg@arm.com    PyArrayDescr_Proxy *descr;
7511986Sandreas.sandberg@arm.com    int flags;
7611986Sandreas.sandberg@arm.com    PyObject *base;
7711986Sandreas.sandberg@arm.com};
7811986Sandreas.sandberg@arm.com
7911986Sandreas.sandberg@arm.comstruct numpy_type_info {
8011986Sandreas.sandberg@arm.com    PyObject* dtype_ptr;
8111986Sandreas.sandberg@arm.com    std::string format_str;
8211986Sandreas.sandberg@arm.com};
8311986Sandreas.sandberg@arm.com
8411986Sandreas.sandberg@arm.comstruct numpy_internals {
8511986Sandreas.sandberg@arm.com    std::unordered_map<std::type_index, numpy_type_info> registered_dtypes;
8611986Sandreas.sandberg@arm.com
8711986Sandreas.sandberg@arm.com    numpy_type_info *get_type_info(const std::type_info& tinfo, bool throw_if_missing = true) {
8811986Sandreas.sandberg@arm.com        auto it = registered_dtypes.find(std::type_index(tinfo));
8911986Sandreas.sandberg@arm.com        if (it != registered_dtypes.end())
9011986Sandreas.sandberg@arm.com            return &(it->second);
9111986Sandreas.sandberg@arm.com        if (throw_if_missing)
9211986Sandreas.sandberg@arm.com            pybind11_fail(std::string("NumPy type info missing for ") + tinfo.name());
9311986Sandreas.sandberg@arm.com        return nullptr;
9411986Sandreas.sandberg@arm.com    }
9511986Sandreas.sandberg@arm.com
9611986Sandreas.sandberg@arm.com    template<typename T> numpy_type_info *get_type_info(bool throw_if_missing = true) {
9711986Sandreas.sandberg@arm.com        return get_type_info(typeid(typename std::remove_cv<T>::type), throw_if_missing);
9811986Sandreas.sandberg@arm.com    }
9911986Sandreas.sandberg@arm.com};
10011986Sandreas.sandberg@arm.com
10111986Sandreas.sandberg@arm.cominline PYBIND11_NOINLINE void load_numpy_internals(numpy_internals* &ptr) {
10211986Sandreas.sandberg@arm.com    ptr = &get_or_create_shared_data<numpy_internals>("_numpy_internals");
10311986Sandreas.sandberg@arm.com}
10411986Sandreas.sandberg@arm.com
10511986Sandreas.sandberg@arm.cominline numpy_internals& get_numpy_internals() {
10611986Sandreas.sandberg@arm.com    static numpy_internals* ptr = nullptr;
10711986Sandreas.sandberg@arm.com    if (!ptr)
10811986Sandreas.sandberg@arm.com        load_numpy_internals(ptr);
10911986Sandreas.sandberg@arm.com    return *ptr;
11011986Sandreas.sandberg@arm.com}
11111986Sandreas.sandberg@arm.com
11214299Sbbruce@ucdavis.edutemplate <typename T> struct same_size {
11314299Sbbruce@ucdavis.edu    template <typename U> using as = bool_constant<sizeof(T) == sizeof(U)>;
11414299Sbbruce@ucdavis.edu};
11514299Sbbruce@ucdavis.edu
11614299Sbbruce@ucdavis.edu// Lookup a type according to its size, and return a value corresponding to the NumPy typenum.
11714299Sbbruce@ucdavis.edutemplate <typename Concrete, typename... Check, typename... Int>
11814299Sbbruce@ucdavis.educonstexpr int platform_lookup(Int... codes) {
11914299Sbbruce@ucdavis.edu    using code_index = std::integral_constant<int, constexpr_first<same_size<Concrete>::template as, Check...>()>;
12014299Sbbruce@ucdavis.edu    static_assert(code_index::value != sizeof...(Check), "Unable to match type on this platform");
12114299Sbbruce@ucdavis.edu    return std::get<code_index::value>(std::make_tuple(codes...));
12214299Sbbruce@ucdavis.edu}
12314299Sbbruce@ucdavis.edu
12411986Sandreas.sandberg@arm.comstruct npy_api {
12511986Sandreas.sandberg@arm.com    enum constants {
12612037Sandreas.sandberg@arm.com        NPY_ARRAY_C_CONTIGUOUS_ = 0x0001,
12712037Sandreas.sandberg@arm.com        NPY_ARRAY_F_CONTIGUOUS_ = 0x0002,
12811986Sandreas.sandberg@arm.com        NPY_ARRAY_OWNDATA_ = 0x0004,
12911986Sandreas.sandberg@arm.com        NPY_ARRAY_FORCECAST_ = 0x0010,
13012037Sandreas.sandberg@arm.com        NPY_ARRAY_ENSUREARRAY_ = 0x0040,
13111986Sandreas.sandberg@arm.com        NPY_ARRAY_ALIGNED_ = 0x0100,
13211986Sandreas.sandberg@arm.com        NPY_ARRAY_WRITEABLE_ = 0x0400,
13311986Sandreas.sandberg@arm.com        NPY_BOOL_ = 0,
13411986Sandreas.sandberg@arm.com        NPY_BYTE_, NPY_UBYTE_,
13511986Sandreas.sandberg@arm.com        NPY_SHORT_, NPY_USHORT_,
13611986Sandreas.sandberg@arm.com        NPY_INT_, NPY_UINT_,
13711986Sandreas.sandberg@arm.com        NPY_LONG_, NPY_ULONG_,
13811986Sandreas.sandberg@arm.com        NPY_LONGLONG_, NPY_ULONGLONG_,
13911986Sandreas.sandberg@arm.com        NPY_FLOAT_, NPY_DOUBLE_, NPY_LONGDOUBLE_,
14011986Sandreas.sandberg@arm.com        NPY_CFLOAT_, NPY_CDOUBLE_, NPY_CLONGDOUBLE_,
14111986Sandreas.sandberg@arm.com        NPY_OBJECT_ = 17,
14214299Sbbruce@ucdavis.edu        NPY_STRING_, NPY_UNICODE_, NPY_VOID_,
14314299Sbbruce@ucdavis.edu        // Platform-dependent normalization
14414299Sbbruce@ucdavis.edu        NPY_INT8_ = NPY_BYTE_,
14514299Sbbruce@ucdavis.edu        NPY_UINT8_ = NPY_UBYTE_,
14614299Sbbruce@ucdavis.edu        NPY_INT16_ = NPY_SHORT_,
14714299Sbbruce@ucdavis.edu        NPY_UINT16_ = NPY_USHORT_,
14814299Sbbruce@ucdavis.edu        // `npy_common.h` defines the integer aliases. In order, it checks:
14914299Sbbruce@ucdavis.edu        // NPY_BITSOF_LONG, NPY_BITSOF_LONGLONG, NPY_BITSOF_INT, NPY_BITSOF_SHORT, NPY_BITSOF_CHAR
15014299Sbbruce@ucdavis.edu        // and assigns the alias to the first matching size, so we should check in this order.
15114299Sbbruce@ucdavis.edu        NPY_INT32_ = platform_lookup<std::int32_t, long, int, short>(
15214299Sbbruce@ucdavis.edu            NPY_LONG_, NPY_INT_, NPY_SHORT_),
15314299Sbbruce@ucdavis.edu        NPY_UINT32_ = platform_lookup<std::uint32_t, unsigned long, unsigned int, unsigned short>(
15414299Sbbruce@ucdavis.edu            NPY_ULONG_, NPY_UINT_, NPY_USHORT_),
15514299Sbbruce@ucdavis.edu        NPY_INT64_ = platform_lookup<std::int64_t, long, long long, int>(
15614299Sbbruce@ucdavis.edu            NPY_LONG_, NPY_LONGLONG_, NPY_INT_),
15714299Sbbruce@ucdavis.edu        NPY_UINT64_ = platform_lookup<std::uint64_t, unsigned long, unsigned long long, unsigned int>(
15814299Sbbruce@ucdavis.edu            NPY_ULONG_, NPY_ULONGLONG_, NPY_UINT_),
15911986Sandreas.sandberg@arm.com    };
16011986Sandreas.sandberg@arm.com
16112391Sjason@lowepower.com    typedef struct {
16212391Sjason@lowepower.com        Py_intptr_t *ptr;
16312391Sjason@lowepower.com        int len;
16412391Sjason@lowepower.com    } PyArray_Dims;
16512391Sjason@lowepower.com
16611986Sandreas.sandberg@arm.com    static npy_api& get() {
16711986Sandreas.sandberg@arm.com        static npy_api api = lookup();
16811986Sandreas.sandberg@arm.com        return api;
16911986Sandreas.sandberg@arm.com    }
17011986Sandreas.sandberg@arm.com
17111986Sandreas.sandberg@arm.com    bool PyArray_Check_(PyObject *obj) const {
17211986Sandreas.sandberg@arm.com        return (bool) PyObject_TypeCheck(obj, PyArray_Type_);
17311986Sandreas.sandberg@arm.com    }
17411986Sandreas.sandberg@arm.com    bool PyArrayDescr_Check_(PyObject *obj) const {
17511986Sandreas.sandberg@arm.com        return (bool) PyObject_TypeCheck(obj, PyArrayDescr_Type_);
17611986Sandreas.sandberg@arm.com    }
17711986Sandreas.sandberg@arm.com
17812391Sjason@lowepower.com    unsigned int (*PyArray_GetNDArrayCFeatureVersion_)();
17911986Sandreas.sandberg@arm.com    PyObject *(*PyArray_DescrFromType_)(int);
18011986Sandreas.sandberg@arm.com    PyObject *(*PyArray_NewFromDescr_)
18111986Sandreas.sandberg@arm.com        (PyTypeObject *, PyObject *, int, Py_intptr_t *,
18211986Sandreas.sandberg@arm.com         Py_intptr_t *, void *, int, PyObject *);
18311986Sandreas.sandberg@arm.com    PyObject *(*PyArray_DescrNewFromType_)(int);
18412391Sjason@lowepower.com    int (*PyArray_CopyInto_)(PyObject *, PyObject *);
18511986Sandreas.sandberg@arm.com    PyObject *(*PyArray_NewCopy_)(PyObject *, int);
18611986Sandreas.sandberg@arm.com    PyTypeObject *PyArray_Type_;
18711986Sandreas.sandberg@arm.com    PyTypeObject *PyVoidArrType_Type_;
18811986Sandreas.sandberg@arm.com    PyTypeObject *PyArrayDescr_Type_;
18911986Sandreas.sandberg@arm.com    PyObject *(*PyArray_DescrFromScalar_)(PyObject *);
19011986Sandreas.sandberg@arm.com    PyObject *(*PyArray_FromAny_) (PyObject *, PyObject *, int, int, int, PyObject *);
19111986Sandreas.sandberg@arm.com    int (*PyArray_DescrConverter_) (PyObject *, PyObject **);
19211986Sandreas.sandberg@arm.com    bool (*PyArray_EquivTypes_) (PyObject *, PyObject *);
19311986Sandreas.sandberg@arm.com    int (*PyArray_GetArrayParamsFromObject_)(PyObject *, PyObject *, char, PyObject **, int *,
19411986Sandreas.sandberg@arm.com                                             Py_ssize_t *, PyObject **, PyObject *);
19511986Sandreas.sandberg@arm.com    PyObject *(*PyArray_Squeeze_)(PyObject *);
19612037Sandreas.sandberg@arm.com    int (*PyArray_SetBaseObject_)(PyObject *, PyObject *);
19712391Sjason@lowepower.com    PyObject* (*PyArray_Resize_)(PyObject*, PyArray_Dims*, int, int);
19811986Sandreas.sandberg@arm.comprivate:
19911986Sandreas.sandberg@arm.com    enum functions {
20012391Sjason@lowepower.com        API_PyArray_GetNDArrayCFeatureVersion = 211,
20111986Sandreas.sandberg@arm.com        API_PyArray_Type = 2,
20211986Sandreas.sandberg@arm.com        API_PyArrayDescr_Type = 3,
20311986Sandreas.sandberg@arm.com        API_PyVoidArrType_Type = 39,
20411986Sandreas.sandberg@arm.com        API_PyArray_DescrFromType = 45,
20511986Sandreas.sandberg@arm.com        API_PyArray_DescrFromScalar = 57,
20611986Sandreas.sandberg@arm.com        API_PyArray_FromAny = 69,
20712391Sjason@lowepower.com        API_PyArray_Resize = 80,
20812391Sjason@lowepower.com        API_PyArray_CopyInto = 82,
20911986Sandreas.sandberg@arm.com        API_PyArray_NewCopy = 85,
21011986Sandreas.sandberg@arm.com        API_PyArray_NewFromDescr = 94,
21111986Sandreas.sandberg@arm.com        API_PyArray_DescrNewFromType = 9,
21211986Sandreas.sandberg@arm.com        API_PyArray_DescrConverter = 174,
21311986Sandreas.sandberg@arm.com        API_PyArray_EquivTypes = 182,
21411986Sandreas.sandberg@arm.com        API_PyArray_GetArrayParamsFromObject = 278,
21512037Sandreas.sandberg@arm.com        API_PyArray_Squeeze = 136,
21612037Sandreas.sandberg@arm.com        API_PyArray_SetBaseObject = 282
21711986Sandreas.sandberg@arm.com    };
21811986Sandreas.sandberg@arm.com
21911986Sandreas.sandberg@arm.com    static npy_api lookup() {
22011986Sandreas.sandberg@arm.com        module m = module::import("numpy.core.multiarray");
22111986Sandreas.sandberg@arm.com        auto c = m.attr("_ARRAY_API");
22211986Sandreas.sandberg@arm.com#if PY_MAJOR_VERSION >= 3
22311986Sandreas.sandberg@arm.com        void **api_ptr = (void **) PyCapsule_GetPointer(c.ptr(), NULL);
22411986Sandreas.sandberg@arm.com#else
22511986Sandreas.sandberg@arm.com        void **api_ptr = (void **) PyCObject_AsVoidPtr(c.ptr());
22611986Sandreas.sandberg@arm.com#endif
22711986Sandreas.sandberg@arm.com        npy_api api;
22811986Sandreas.sandberg@arm.com#define DECL_NPY_API(Func) api.Func##_ = (decltype(api.Func##_)) api_ptr[API_##Func];
22912391Sjason@lowepower.com        DECL_NPY_API(PyArray_GetNDArrayCFeatureVersion);
23012391Sjason@lowepower.com        if (api.PyArray_GetNDArrayCFeatureVersion_() < 0x7)
23112391Sjason@lowepower.com            pybind11_fail("pybind11 numpy support requires numpy >= 1.7.0");
23211986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_Type);
23311986Sandreas.sandberg@arm.com        DECL_NPY_API(PyVoidArrType_Type);
23411986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArrayDescr_Type);
23511986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_DescrFromType);
23611986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_DescrFromScalar);
23711986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_FromAny);
23812391Sjason@lowepower.com        DECL_NPY_API(PyArray_Resize);
23912391Sjason@lowepower.com        DECL_NPY_API(PyArray_CopyInto);
24011986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_NewCopy);
24111986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_NewFromDescr);
24211986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_DescrNewFromType);
24311986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_DescrConverter);
24411986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_EquivTypes);
24511986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_GetArrayParamsFromObject);
24611986Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_Squeeze);
24712037Sandreas.sandberg@arm.com        DECL_NPY_API(PyArray_SetBaseObject);
24811986Sandreas.sandberg@arm.com#undef DECL_NPY_API
24911986Sandreas.sandberg@arm.com        return api;
25011986Sandreas.sandberg@arm.com    }
25111986Sandreas.sandberg@arm.com};
25211986Sandreas.sandberg@arm.com
25311986Sandreas.sandberg@arm.cominline PyArray_Proxy* array_proxy(void* ptr) {
25411986Sandreas.sandberg@arm.com    return reinterpret_cast<PyArray_Proxy*>(ptr);
25511986Sandreas.sandberg@arm.com}
25611986Sandreas.sandberg@arm.com
25711986Sandreas.sandberg@arm.cominline const PyArray_Proxy* array_proxy(const void* ptr) {
25811986Sandreas.sandberg@arm.com    return reinterpret_cast<const PyArray_Proxy*>(ptr);
25911986Sandreas.sandberg@arm.com}
26011986Sandreas.sandberg@arm.com
26111986Sandreas.sandberg@arm.cominline PyArrayDescr_Proxy* array_descriptor_proxy(PyObject* ptr) {
26211986Sandreas.sandberg@arm.com   return reinterpret_cast<PyArrayDescr_Proxy*>(ptr);
26311986Sandreas.sandberg@arm.com}
26411986Sandreas.sandberg@arm.com
26511986Sandreas.sandberg@arm.cominline const PyArrayDescr_Proxy* array_descriptor_proxy(const PyObject* ptr) {
26611986Sandreas.sandberg@arm.com   return reinterpret_cast<const PyArrayDescr_Proxy*>(ptr);
26711986Sandreas.sandberg@arm.com}
26811986Sandreas.sandberg@arm.com
26911986Sandreas.sandberg@arm.cominline bool check_flags(const void* ptr, int flag) {
27011986Sandreas.sandberg@arm.com    return (flag == (array_proxy(ptr)->flags & flag));
27111986Sandreas.sandberg@arm.com}
27211986Sandreas.sandberg@arm.com
27312037Sandreas.sandberg@arm.comtemplate <typename T> struct is_std_array : std::false_type { };
27412037Sandreas.sandberg@arm.comtemplate <typename T, size_t N> struct is_std_array<std::array<T, N>> : std::true_type { };
27512037Sandreas.sandberg@arm.comtemplate <typename T> struct is_complex : std::false_type { };
27612037Sandreas.sandberg@arm.comtemplate <typename T> struct is_complex<std::complex<T>> : std::true_type { };
27712037Sandreas.sandberg@arm.com
27812391Sjason@lowepower.comtemplate <typename T> struct array_info_scalar {
27912391Sjason@lowepower.com    typedef T type;
28012391Sjason@lowepower.com    static constexpr bool is_array = false;
28112391Sjason@lowepower.com    static constexpr bool is_empty = false;
28214299Sbbruce@ucdavis.edu    static constexpr auto extents = _("");
28312391Sjason@lowepower.com    static void append_extents(list& /* shape */) { }
28412391Sjason@lowepower.com};
28512391Sjason@lowepower.com// Computes underlying type and a comma-separated list of extents for array
28612391Sjason@lowepower.com// types (any mix of std::array and built-in arrays). An array of char is
28712391Sjason@lowepower.com// treated as scalar because it gets special handling.
28812391Sjason@lowepower.comtemplate <typename T> struct array_info : array_info_scalar<T> { };
28912391Sjason@lowepower.comtemplate <typename T, size_t N> struct array_info<std::array<T, N>> {
29012391Sjason@lowepower.com    using type = typename array_info<T>::type;
29112391Sjason@lowepower.com    static constexpr bool is_array = true;
29212391Sjason@lowepower.com    static constexpr bool is_empty = (N == 0) || array_info<T>::is_empty;
29312391Sjason@lowepower.com    static constexpr size_t extent = N;
29412391Sjason@lowepower.com
29512391Sjason@lowepower.com    // appends the extents to shape
29612391Sjason@lowepower.com    static void append_extents(list& shape) {
29712391Sjason@lowepower.com        shape.append(N);
29812391Sjason@lowepower.com        array_info<T>::append_extents(shape);
29912391Sjason@lowepower.com    }
30012391Sjason@lowepower.com
30114299Sbbruce@ucdavis.edu    static constexpr auto extents = _<array_info<T>::is_array>(
30214299Sbbruce@ucdavis.edu        concat(_<N>(), array_info<T>::extents), _<N>()
30314299Sbbruce@ucdavis.edu    );
30412391Sjason@lowepower.com};
30512391Sjason@lowepower.com// For numpy we have special handling for arrays of characters, so we don't include
30612391Sjason@lowepower.com// the size in the array extents.
30712391Sjason@lowepower.comtemplate <size_t N> struct array_info<char[N]> : array_info_scalar<char[N]> { };
30812391Sjason@lowepower.comtemplate <size_t N> struct array_info<std::array<char, N>> : array_info_scalar<std::array<char, N>> { };
30912391Sjason@lowepower.comtemplate <typename T, size_t N> struct array_info<T[N]> : array_info<std::array<T, N>> { };
31012391Sjason@lowepower.comtemplate <typename T> using remove_all_extents_t = typename array_info<T>::type;
31112391Sjason@lowepower.com
31212037Sandreas.sandberg@arm.comtemplate <typename T> using is_pod_struct = all_of<
31312391Sjason@lowepower.com    std::is_standard_layout<T>,     // since we're accessing directly in memory we need a standard layout type
31412391Sjason@lowepower.com#if !defined(__GNUG__) || defined(_LIBCPP_VERSION) || defined(_GLIBCXX_USE_CXX11_ABI)
31512391Sjason@lowepower.com    // _GLIBCXX_USE_CXX11_ABI indicates that we're using libstdc++ from GCC 5 or newer, independent
31612391Sjason@lowepower.com    // of the actual compiler (Clang can also use libstdc++, but it always defines __GNUC__ == 4).
31712391Sjason@lowepower.com    std::is_trivially_copyable<T>,
31812391Sjason@lowepower.com#else
31912391Sjason@lowepower.com    // GCC 4 doesn't implement is_trivially_copyable, so approximate it
32012391Sjason@lowepower.com    std::is_trivially_destructible<T>,
32112391Sjason@lowepower.com    satisfies_any_of<T, std::has_trivial_copy_constructor, std::has_trivial_copy_assign>,
32212391Sjason@lowepower.com#endif
32312037Sandreas.sandberg@arm.com    satisfies_none_of<T, std::is_reference, std::is_array, is_std_array, std::is_arithmetic, is_complex, std::is_enum>
32412037Sandreas.sandberg@arm.com>;
32512037Sandreas.sandberg@arm.com
32612391Sjason@lowepower.comtemplate <ssize_t Dim = 0, typename Strides> ssize_t byte_offset_unsafe(const Strides &) { return 0; }
32712391Sjason@lowepower.comtemplate <ssize_t Dim = 0, typename Strides, typename... Ix>
32812391Sjason@lowepower.comssize_t byte_offset_unsafe(const Strides &strides, ssize_t i, Ix... index) {
32912037Sandreas.sandberg@arm.com    return i * strides[Dim] + byte_offset_unsafe<Dim + 1>(strides, index...);
33012037Sandreas.sandberg@arm.com}
33112037Sandreas.sandberg@arm.com
33212391Sjason@lowepower.com/**
33312391Sjason@lowepower.com * Proxy class providing unsafe, unchecked const access to array data.  This is constructed through
33412037Sandreas.sandberg@arm.com * the `unchecked<T, N>()` method of `array` or the `unchecked<N>()` method of `array_t<T>`.  `Dims`
33512037Sandreas.sandberg@arm.com * will be -1 for dimensions determined at runtime.
33612037Sandreas.sandberg@arm.com */
33712037Sandreas.sandberg@arm.comtemplate <typename T, ssize_t Dims>
33812037Sandreas.sandberg@arm.comclass unchecked_reference {
33912037Sandreas.sandberg@arm.comprotected:
34012037Sandreas.sandberg@arm.com    static constexpr bool Dynamic = Dims < 0;
34112037Sandreas.sandberg@arm.com    const unsigned char *data_;
34212037Sandreas.sandberg@arm.com    // Storing the shape & strides in local variables (i.e. these arrays) allows the compiler to
34312037Sandreas.sandberg@arm.com    // make large performance gains on big, nested loops, but requires compile-time dimensions
34412391Sjason@lowepower.com    conditional_t<Dynamic, const ssize_t *, std::array<ssize_t, (size_t) Dims>>
34512391Sjason@lowepower.com            shape_, strides_;
34612391Sjason@lowepower.com    const ssize_t dims_;
34712037Sandreas.sandberg@arm.com
34812037Sandreas.sandberg@arm.com    friend class pybind11::array;
34912037Sandreas.sandberg@arm.com    // Constructor for compile-time dimensions:
35012037Sandreas.sandberg@arm.com    template <bool Dyn = Dynamic>
35112391Sjason@lowepower.com    unchecked_reference(const void *data, const ssize_t *shape, const ssize_t *strides, enable_if_t<!Dyn, ssize_t>)
35212037Sandreas.sandberg@arm.com    : data_{reinterpret_cast<const unsigned char *>(data)}, dims_{Dims} {
35312391Sjason@lowepower.com        for (size_t i = 0; i < (size_t) dims_; i++) {
35412037Sandreas.sandberg@arm.com            shape_[i] = shape[i];
35512037Sandreas.sandberg@arm.com            strides_[i] = strides[i];
35612037Sandreas.sandberg@arm.com        }
35712037Sandreas.sandberg@arm.com    }
35812037Sandreas.sandberg@arm.com    // Constructor for runtime dimensions:
35912037Sandreas.sandberg@arm.com    template <bool Dyn = Dynamic>
36012391Sjason@lowepower.com    unchecked_reference(const void *data, const ssize_t *shape, const ssize_t *strides, enable_if_t<Dyn, ssize_t> dims)
36112037Sandreas.sandberg@arm.com    : data_{reinterpret_cast<const unsigned char *>(data)}, shape_{shape}, strides_{strides}, dims_{dims} {}
36212037Sandreas.sandberg@arm.com
36312037Sandreas.sandberg@arm.compublic:
36412391Sjason@lowepower.com    /**
36512391Sjason@lowepower.com     * Unchecked const reference access to data at the given indices.  For a compile-time known
36612037Sandreas.sandberg@arm.com     * number of dimensions, this requires the correct number of arguments; for run-time
36712037Sandreas.sandberg@arm.com     * dimensionality, this is not checked (and so is up to the caller to use safely).
36812037Sandreas.sandberg@arm.com     */
36912037Sandreas.sandberg@arm.com    template <typename... Ix> const T &operator()(Ix... index) const {
37012391Sjason@lowepower.com        static_assert(ssize_t{sizeof...(Ix)} == Dims || Dynamic,
37112037Sandreas.sandberg@arm.com                "Invalid number of indices for unchecked array reference");
37212391Sjason@lowepower.com        return *reinterpret_cast<const T *>(data_ + byte_offset_unsafe(strides_, ssize_t(index)...));
37312037Sandreas.sandberg@arm.com    }
37412391Sjason@lowepower.com    /**
37512391Sjason@lowepower.com     * Unchecked const reference access to data; this operator only participates if the reference
37612037Sandreas.sandberg@arm.com     * is to a 1-dimensional array.  When present, this is exactly equivalent to `obj(index)`.
37712037Sandreas.sandberg@arm.com     */
37812391Sjason@lowepower.com    template <ssize_t D = Dims, typename = enable_if_t<D == 1 || Dynamic>>
37912391Sjason@lowepower.com    const T &operator[](ssize_t index) const { return operator()(index); }
38012037Sandreas.sandberg@arm.com
38112037Sandreas.sandberg@arm.com    /// Pointer access to the data at the given indices.
38212391Sjason@lowepower.com    template <typename... Ix> const T *data(Ix... ix) const { return &operator()(ssize_t(ix)...); }
38312037Sandreas.sandberg@arm.com
38412037Sandreas.sandberg@arm.com    /// Returns the item size, i.e. sizeof(T)
38512391Sjason@lowepower.com    constexpr static ssize_t itemsize() { return sizeof(T); }
38612037Sandreas.sandberg@arm.com
38712037Sandreas.sandberg@arm.com    /// Returns the shape (i.e. size) of dimension `dim`
38812391Sjason@lowepower.com    ssize_t shape(ssize_t dim) const { return shape_[(size_t) dim]; }
38912037Sandreas.sandberg@arm.com
39012037Sandreas.sandberg@arm.com    /// Returns the number of dimensions of the array
39112391Sjason@lowepower.com    ssize_t ndim() const { return dims_; }
39212037Sandreas.sandberg@arm.com
39312037Sandreas.sandberg@arm.com    /// Returns the total number of elements in the referenced array, i.e. the product of the shapes
39412037Sandreas.sandberg@arm.com    template <bool Dyn = Dynamic>
39512391Sjason@lowepower.com    enable_if_t<!Dyn, ssize_t> size() const {
39612391Sjason@lowepower.com        return std::accumulate(shape_.begin(), shape_.end(), (ssize_t) 1, std::multiplies<ssize_t>());
39712037Sandreas.sandberg@arm.com    }
39812037Sandreas.sandberg@arm.com    template <bool Dyn = Dynamic>
39912391Sjason@lowepower.com    enable_if_t<Dyn, ssize_t> size() const {
40012391Sjason@lowepower.com        return std::accumulate(shape_, shape_ + ndim(), (ssize_t) 1, std::multiplies<ssize_t>());
40112037Sandreas.sandberg@arm.com    }
40212037Sandreas.sandberg@arm.com
40312037Sandreas.sandberg@arm.com    /// Returns the total number of bytes used by the referenced data.  Note that the actual span in
40412037Sandreas.sandberg@arm.com    /// memory may be larger if the referenced array has non-contiguous strides (e.g. for a slice).
40512391Sjason@lowepower.com    ssize_t nbytes() const {
40612037Sandreas.sandberg@arm.com        return size() * itemsize();
40712037Sandreas.sandberg@arm.com    }
40812037Sandreas.sandberg@arm.com};
40912037Sandreas.sandberg@arm.com
41012037Sandreas.sandberg@arm.comtemplate <typename T, ssize_t Dims>
41112037Sandreas.sandberg@arm.comclass unchecked_mutable_reference : public unchecked_reference<T, Dims> {
41212037Sandreas.sandberg@arm.com    friend class pybind11::array;
41312037Sandreas.sandberg@arm.com    using ConstBase = unchecked_reference<T, Dims>;
41412037Sandreas.sandberg@arm.com    using ConstBase::ConstBase;
41512037Sandreas.sandberg@arm.com    using ConstBase::Dynamic;
41612037Sandreas.sandberg@arm.compublic:
41712037Sandreas.sandberg@arm.com    /// Mutable, unchecked access to data at the given indices.
41812037Sandreas.sandberg@arm.com    template <typename... Ix> T& operator()(Ix... index) {
41912391Sjason@lowepower.com        static_assert(ssize_t{sizeof...(Ix)} == Dims || Dynamic,
42012037Sandreas.sandberg@arm.com                "Invalid number of indices for unchecked array reference");
42112037Sandreas.sandberg@arm.com        return const_cast<T &>(ConstBase::operator()(index...));
42212037Sandreas.sandberg@arm.com    }
42312391Sjason@lowepower.com    /**
42412391Sjason@lowepower.com     * Mutable, unchecked access data at the given index; this operator only participates if the
42512037Sandreas.sandberg@arm.com     * reference is to a 1-dimensional array (or has runtime dimensions).  When present, this is
42612037Sandreas.sandberg@arm.com     * exactly equivalent to `obj(index)`.
42712037Sandreas.sandberg@arm.com     */
42812391Sjason@lowepower.com    template <ssize_t D = Dims, typename = enable_if_t<D == 1 || Dynamic>>
42912391Sjason@lowepower.com    T &operator[](ssize_t index) { return operator()(index); }
43012037Sandreas.sandberg@arm.com
43112037Sandreas.sandberg@arm.com    /// Mutable pointer access to the data at the given indices.
43212391Sjason@lowepower.com    template <typename... Ix> T *mutable_data(Ix... ix) { return &operator()(ssize_t(ix)...); }
43312037Sandreas.sandberg@arm.com};
43412037Sandreas.sandberg@arm.com
43512391Sjason@lowepower.comtemplate <typename T, ssize_t Dim>
43612037Sandreas.sandberg@arm.comstruct type_caster<unchecked_reference<T, Dim>> {
43712037Sandreas.sandberg@arm.com    static_assert(Dim == 0 && Dim > 0 /* always fail */, "unchecked array proxy object is not castable");
43812037Sandreas.sandberg@arm.com};
43912391Sjason@lowepower.comtemplate <typename T, ssize_t Dim>
44012037Sandreas.sandberg@arm.comstruct type_caster<unchecked_mutable_reference<T, Dim>> : type_caster<unchecked_reference<T, Dim>> {};
44112037Sandreas.sandberg@arm.com
44211986Sandreas.sandberg@arm.comNAMESPACE_END(detail)
44311986Sandreas.sandberg@arm.com
44411986Sandreas.sandberg@arm.comclass dtype : public object {
44511986Sandreas.sandberg@arm.compublic:
44611986Sandreas.sandberg@arm.com    PYBIND11_OBJECT_DEFAULT(dtype, object, detail::npy_api::get().PyArrayDescr_Check_);
44711986Sandreas.sandberg@arm.com
44811986Sandreas.sandberg@arm.com    explicit dtype(const buffer_info &info) {
44911986Sandreas.sandberg@arm.com        dtype descr(_dtype_from_pep3118()(PYBIND11_STR_TYPE(info.format)));
45011986Sandreas.sandberg@arm.com        // If info.itemsize == 0, use the value calculated from the format string
45111986Sandreas.sandberg@arm.com        m_ptr = descr.strip_padding(info.itemsize ? info.itemsize : descr.itemsize()).release().ptr();
45211986Sandreas.sandberg@arm.com    }
45311986Sandreas.sandberg@arm.com
45411986Sandreas.sandberg@arm.com    explicit dtype(const std::string &format) {
45511986Sandreas.sandberg@arm.com        m_ptr = from_args(pybind11::str(format)).release().ptr();
45611986Sandreas.sandberg@arm.com    }
45711986Sandreas.sandberg@arm.com
45811986Sandreas.sandberg@arm.com    dtype(const char *format) : dtype(std::string(format)) { }
45911986Sandreas.sandberg@arm.com
46012391Sjason@lowepower.com    dtype(list names, list formats, list offsets, ssize_t itemsize) {
46111986Sandreas.sandberg@arm.com        dict args;
46211986Sandreas.sandberg@arm.com        args["names"] = names;
46311986Sandreas.sandberg@arm.com        args["formats"] = formats;
46411986Sandreas.sandberg@arm.com        args["offsets"] = offsets;
46511986Sandreas.sandberg@arm.com        args["itemsize"] = pybind11::int_(itemsize);
46611986Sandreas.sandberg@arm.com        m_ptr = from_args(args).release().ptr();
46711986Sandreas.sandberg@arm.com    }
46811986Sandreas.sandberg@arm.com
46911986Sandreas.sandberg@arm.com    /// This is essentially the same as calling numpy.dtype(args) in Python.
47011986Sandreas.sandberg@arm.com    static dtype from_args(object args) {
47111986Sandreas.sandberg@arm.com        PyObject *ptr = nullptr;
47214299Sbbruce@ucdavis.edu        if (!detail::npy_api::get().PyArray_DescrConverter_(args.ptr(), &ptr) || !ptr)
47311986Sandreas.sandberg@arm.com            throw error_already_set();
47411986Sandreas.sandberg@arm.com        return reinterpret_steal<dtype>(ptr);
47511986Sandreas.sandberg@arm.com    }
47611986Sandreas.sandberg@arm.com
47711986Sandreas.sandberg@arm.com    /// Return dtype associated with a C++ type.
47811986Sandreas.sandberg@arm.com    template <typename T> static dtype of() {
47911986Sandreas.sandberg@arm.com        return detail::npy_format_descriptor<typename std::remove_cv<T>::type>::dtype();
48011986Sandreas.sandberg@arm.com    }
48111986Sandreas.sandberg@arm.com
48211986Sandreas.sandberg@arm.com    /// Size of the data type in bytes.
48312391Sjason@lowepower.com    ssize_t itemsize() const {
48412391Sjason@lowepower.com        return detail::array_descriptor_proxy(m_ptr)->elsize;
48511986Sandreas.sandberg@arm.com    }
48611986Sandreas.sandberg@arm.com
48711986Sandreas.sandberg@arm.com    /// Returns true for structured data types.
48811986Sandreas.sandberg@arm.com    bool has_fields() const {
48911986Sandreas.sandberg@arm.com        return detail::array_descriptor_proxy(m_ptr)->names != nullptr;
49011986Sandreas.sandberg@arm.com    }
49111986Sandreas.sandberg@arm.com
49211986Sandreas.sandberg@arm.com    /// Single-character type code.
49311986Sandreas.sandberg@arm.com    char kind() const {
49411986Sandreas.sandberg@arm.com        return detail::array_descriptor_proxy(m_ptr)->kind;
49511986Sandreas.sandberg@arm.com    }
49611986Sandreas.sandberg@arm.com
49711986Sandreas.sandberg@arm.comprivate:
49811986Sandreas.sandberg@arm.com    static object _dtype_from_pep3118() {
49911986Sandreas.sandberg@arm.com        static PyObject *obj = module::import("numpy.core._internal")
50011986Sandreas.sandberg@arm.com            .attr("_dtype_from_pep3118").cast<object>().release().ptr();
50111986Sandreas.sandberg@arm.com        return reinterpret_borrow<object>(obj);
50211986Sandreas.sandberg@arm.com    }
50311986Sandreas.sandberg@arm.com
50412391Sjason@lowepower.com    dtype strip_padding(ssize_t itemsize) {
50511986Sandreas.sandberg@arm.com        // Recursively strip all void fields with empty names that are generated for
50611986Sandreas.sandberg@arm.com        // padding fields (as of NumPy v1.11).
50711986Sandreas.sandberg@arm.com        if (!has_fields())
50811986Sandreas.sandberg@arm.com            return *this;
50911986Sandreas.sandberg@arm.com
51011986Sandreas.sandberg@arm.com        struct field_descr { PYBIND11_STR_TYPE name; object format; pybind11::int_ offset; };
51111986Sandreas.sandberg@arm.com        std::vector<field_descr> field_descriptors;
51211986Sandreas.sandberg@arm.com
51311986Sandreas.sandberg@arm.com        for (auto field : attr("fields").attr("items")()) {
51411986Sandreas.sandberg@arm.com            auto spec = field.cast<tuple>();
51511986Sandreas.sandberg@arm.com            auto name = spec[0].cast<pybind11::str>();
51611986Sandreas.sandberg@arm.com            auto format = spec[1].cast<tuple>()[0].cast<dtype>();
51711986Sandreas.sandberg@arm.com            auto offset = spec[1].cast<tuple>()[1].cast<pybind11::int_>();
51811986Sandreas.sandberg@arm.com            if (!len(name) && format.kind() == 'V')
51911986Sandreas.sandberg@arm.com                continue;
52011986Sandreas.sandberg@arm.com            field_descriptors.push_back({(PYBIND11_STR_TYPE) name, format.strip_padding(format.itemsize()), offset});
52111986Sandreas.sandberg@arm.com        }
52211986Sandreas.sandberg@arm.com
52311986Sandreas.sandberg@arm.com        std::sort(field_descriptors.begin(), field_descriptors.end(),
52411986Sandreas.sandberg@arm.com                  [](const field_descr& a, const field_descr& b) {
52511986Sandreas.sandberg@arm.com                      return a.offset.cast<int>() < b.offset.cast<int>();
52611986Sandreas.sandberg@arm.com                  });
52711986Sandreas.sandberg@arm.com
52811986Sandreas.sandberg@arm.com        list names, formats, offsets;
52911986Sandreas.sandberg@arm.com        for (auto& descr : field_descriptors) {
53011986Sandreas.sandberg@arm.com            names.append(descr.name);
53111986Sandreas.sandberg@arm.com            formats.append(descr.format);
53211986Sandreas.sandberg@arm.com            offsets.append(descr.offset);
53311986Sandreas.sandberg@arm.com        }
53411986Sandreas.sandberg@arm.com        return dtype(names, formats, offsets, itemsize);
53511986Sandreas.sandberg@arm.com    }
53611986Sandreas.sandberg@arm.com};
53711986Sandreas.sandberg@arm.com
53811986Sandreas.sandberg@arm.comclass array : public buffer {
53911986Sandreas.sandberg@arm.compublic:
54011986Sandreas.sandberg@arm.com    PYBIND11_OBJECT_CVT(array, buffer, detail::npy_api::get().PyArray_Check_, raw_array)
54111986Sandreas.sandberg@arm.com
54211986Sandreas.sandberg@arm.com    enum {
54312037Sandreas.sandberg@arm.com        c_style = detail::npy_api::NPY_ARRAY_C_CONTIGUOUS_,
54412037Sandreas.sandberg@arm.com        f_style = detail::npy_api::NPY_ARRAY_F_CONTIGUOUS_,
54511986Sandreas.sandberg@arm.com        forcecast = detail::npy_api::NPY_ARRAY_FORCECAST_
54611986Sandreas.sandberg@arm.com    };
54711986Sandreas.sandberg@arm.com
54812391Sjason@lowepower.com    array() : array({{0}}, static_cast<const double *>(nullptr)) {}
54911986Sandreas.sandberg@arm.com
55012391Sjason@lowepower.com    using ShapeContainer = detail::any_container<ssize_t>;
55112391Sjason@lowepower.com    using StridesContainer = detail::any_container<ssize_t>;
55212391Sjason@lowepower.com
55312391Sjason@lowepower.com    // Constructs an array taking shape/strides from arbitrary container types
55412391Sjason@lowepower.com    array(const pybind11::dtype &dt, ShapeContainer shape, StridesContainer strides,
55512391Sjason@lowepower.com          const void *ptr = nullptr, handle base = handle()) {
55612391Sjason@lowepower.com
55712391Sjason@lowepower.com        if (strides->empty())
55812391Sjason@lowepower.com            *strides = c_strides(*shape, dt.itemsize());
55912391Sjason@lowepower.com
56012391Sjason@lowepower.com        auto ndim = shape->size();
56112391Sjason@lowepower.com        if (ndim != strides->size())
56211986Sandreas.sandberg@arm.com            pybind11_fail("NumPy: shape ndim doesn't match strides ndim");
56311986Sandreas.sandberg@arm.com        auto descr = dt;
56411986Sandreas.sandberg@arm.com
56511986Sandreas.sandberg@arm.com        int flags = 0;
56611986Sandreas.sandberg@arm.com        if (base && ptr) {
56711986Sandreas.sandberg@arm.com            if (isinstance<array>(base))
56812037Sandreas.sandberg@arm.com                /* Copy flags from base (except ownership bit) */
56911986Sandreas.sandberg@arm.com                flags = reinterpret_borrow<array>(base).flags() & ~detail::npy_api::NPY_ARRAY_OWNDATA_;
57011986Sandreas.sandberg@arm.com            else
57111986Sandreas.sandberg@arm.com                /* Writable by default, easy to downgrade later on if needed */
57211986Sandreas.sandberg@arm.com                flags = detail::npy_api::NPY_ARRAY_WRITEABLE_;
57311986Sandreas.sandberg@arm.com        }
57411986Sandreas.sandberg@arm.com
57512391Sjason@lowepower.com        auto &api = detail::npy_api::get();
57611986Sandreas.sandberg@arm.com        auto tmp = reinterpret_steal<object>(api.PyArray_NewFromDescr_(
57712391Sjason@lowepower.com            api.PyArray_Type_, descr.release().ptr(), (int) ndim, shape->data(), strides->data(),
57812037Sandreas.sandberg@arm.com            const_cast<void *>(ptr), flags, nullptr));
57911986Sandreas.sandberg@arm.com        if (!tmp)
58012391Sjason@lowepower.com            throw error_already_set();
58111986Sandreas.sandberg@arm.com        if (ptr) {
58211986Sandreas.sandberg@arm.com            if (base) {
58312037Sandreas.sandberg@arm.com                api.PyArray_SetBaseObject_(tmp.ptr(), base.inc_ref().ptr());
58411986Sandreas.sandberg@arm.com            } else {
58511986Sandreas.sandberg@arm.com                tmp = reinterpret_steal<object>(api.PyArray_NewCopy_(tmp.ptr(), -1 /* any order */));
58611986Sandreas.sandberg@arm.com            }
58711986Sandreas.sandberg@arm.com        }
58811986Sandreas.sandberg@arm.com        m_ptr = tmp.release().ptr();
58911986Sandreas.sandberg@arm.com    }
59011986Sandreas.sandberg@arm.com
59112391Sjason@lowepower.com    array(const pybind11::dtype &dt, ShapeContainer shape, const void *ptr = nullptr, handle base = handle())
59212391Sjason@lowepower.com        : array(dt, std::move(shape), {}, ptr, base) { }
59311986Sandreas.sandberg@arm.com
59412391Sjason@lowepower.com    template <typename T, typename = detail::enable_if_t<std::is_integral<T>::value && !std::is_same<bool, T>::value>>
59512391Sjason@lowepower.com    array(const pybind11::dtype &dt, T count, const void *ptr = nullptr, handle base = handle())
59612391Sjason@lowepower.com        : array(dt, {{count}}, ptr, base) { }
59711986Sandreas.sandberg@arm.com
59811986Sandreas.sandberg@arm.com    template <typename T>
59912391Sjason@lowepower.com    array(ShapeContainer shape, StridesContainer strides, const T *ptr, handle base = handle())
60012391Sjason@lowepower.com        : array(pybind11::dtype::of<T>(), std::move(shape), std::move(strides), ptr, base) { }
60111986Sandreas.sandberg@arm.com
60211986Sandreas.sandberg@arm.com    template <typename T>
60312391Sjason@lowepower.com    array(ShapeContainer shape, const T *ptr, handle base = handle())
60412391Sjason@lowepower.com        : array(std::move(shape), {}, ptr, base) { }
60512391Sjason@lowepower.com
60612391Sjason@lowepower.com    template <typename T>
60712391Sjason@lowepower.com    explicit array(ssize_t count, const T *ptr, handle base = handle()) : array({count}, {}, ptr, base) { }
60811986Sandreas.sandberg@arm.com
60911986Sandreas.sandberg@arm.com    explicit array(const buffer_info &info)
61011986Sandreas.sandberg@arm.com    : array(pybind11::dtype(info), info.shape, info.strides, info.ptr) { }
61111986Sandreas.sandberg@arm.com
61211986Sandreas.sandberg@arm.com    /// Array descriptor (dtype)
61311986Sandreas.sandberg@arm.com    pybind11::dtype dtype() const {
61411986Sandreas.sandberg@arm.com        return reinterpret_borrow<pybind11::dtype>(detail::array_proxy(m_ptr)->descr);
61511986Sandreas.sandberg@arm.com    }
61611986Sandreas.sandberg@arm.com
61711986Sandreas.sandberg@arm.com    /// Total number of elements
61812391Sjason@lowepower.com    ssize_t size() const {
61912391Sjason@lowepower.com        return std::accumulate(shape(), shape() + ndim(), (ssize_t) 1, std::multiplies<ssize_t>());
62011986Sandreas.sandberg@arm.com    }
62111986Sandreas.sandberg@arm.com
62211986Sandreas.sandberg@arm.com    /// Byte size of a single element
62312391Sjason@lowepower.com    ssize_t itemsize() const {
62412391Sjason@lowepower.com        return detail::array_descriptor_proxy(detail::array_proxy(m_ptr)->descr)->elsize;
62511986Sandreas.sandberg@arm.com    }
62611986Sandreas.sandberg@arm.com
62711986Sandreas.sandberg@arm.com    /// Total number of bytes
62812391Sjason@lowepower.com    ssize_t nbytes() const {
62911986Sandreas.sandberg@arm.com        return size() * itemsize();
63011986Sandreas.sandberg@arm.com    }
63111986Sandreas.sandberg@arm.com
63211986Sandreas.sandberg@arm.com    /// Number of dimensions
63312391Sjason@lowepower.com    ssize_t ndim() const {
63412391Sjason@lowepower.com        return detail::array_proxy(m_ptr)->nd;
63511986Sandreas.sandberg@arm.com    }
63611986Sandreas.sandberg@arm.com
63711986Sandreas.sandberg@arm.com    /// Base object
63811986Sandreas.sandberg@arm.com    object base() const {
63911986Sandreas.sandberg@arm.com        return reinterpret_borrow<object>(detail::array_proxy(m_ptr)->base);
64011986Sandreas.sandberg@arm.com    }
64111986Sandreas.sandberg@arm.com
64211986Sandreas.sandberg@arm.com    /// Dimensions of the array
64312391Sjason@lowepower.com    const ssize_t* shape() const {
64412391Sjason@lowepower.com        return detail::array_proxy(m_ptr)->dimensions;
64511986Sandreas.sandberg@arm.com    }
64611986Sandreas.sandberg@arm.com
64711986Sandreas.sandberg@arm.com    /// Dimension along a given axis
64812391Sjason@lowepower.com    ssize_t shape(ssize_t dim) const {
64911986Sandreas.sandberg@arm.com        if (dim >= ndim())
65011986Sandreas.sandberg@arm.com            fail_dim_check(dim, "invalid axis");
65111986Sandreas.sandberg@arm.com        return shape()[dim];
65211986Sandreas.sandberg@arm.com    }
65311986Sandreas.sandberg@arm.com
65411986Sandreas.sandberg@arm.com    /// Strides of the array
65512391Sjason@lowepower.com    const ssize_t* strides() const {
65612391Sjason@lowepower.com        return detail::array_proxy(m_ptr)->strides;
65711986Sandreas.sandberg@arm.com    }
65811986Sandreas.sandberg@arm.com
65911986Sandreas.sandberg@arm.com    /// Stride along a given axis
66012391Sjason@lowepower.com    ssize_t strides(ssize_t dim) const {
66111986Sandreas.sandberg@arm.com        if (dim >= ndim())
66211986Sandreas.sandberg@arm.com            fail_dim_check(dim, "invalid axis");
66311986Sandreas.sandberg@arm.com        return strides()[dim];
66411986Sandreas.sandberg@arm.com    }
66511986Sandreas.sandberg@arm.com
66611986Sandreas.sandberg@arm.com    /// Return the NumPy array flags
66711986Sandreas.sandberg@arm.com    int flags() const {
66811986Sandreas.sandberg@arm.com        return detail::array_proxy(m_ptr)->flags;
66911986Sandreas.sandberg@arm.com    }
67011986Sandreas.sandberg@arm.com
67111986Sandreas.sandberg@arm.com    /// If set, the array is writeable (otherwise the buffer is read-only)
67211986Sandreas.sandberg@arm.com    bool writeable() const {
67311986Sandreas.sandberg@arm.com        return detail::check_flags(m_ptr, detail::npy_api::NPY_ARRAY_WRITEABLE_);
67411986Sandreas.sandberg@arm.com    }
67511986Sandreas.sandberg@arm.com
67611986Sandreas.sandberg@arm.com    /// If set, the array owns the data (will be freed when the array is deleted)
67711986Sandreas.sandberg@arm.com    bool owndata() const {
67811986Sandreas.sandberg@arm.com        return detail::check_flags(m_ptr, detail::npy_api::NPY_ARRAY_OWNDATA_);
67911986Sandreas.sandberg@arm.com    }
68011986Sandreas.sandberg@arm.com
68111986Sandreas.sandberg@arm.com    /// Pointer to the contained data. If index is not provided, points to the
68211986Sandreas.sandberg@arm.com    /// beginning of the buffer. May throw if the index would lead to out of bounds access.
68311986Sandreas.sandberg@arm.com    template<typename... Ix> const void* data(Ix... index) const {
68411986Sandreas.sandberg@arm.com        return static_cast<const void *>(detail::array_proxy(m_ptr)->data + offset_at(index...));
68511986Sandreas.sandberg@arm.com    }
68611986Sandreas.sandberg@arm.com
68711986Sandreas.sandberg@arm.com    /// Mutable pointer to the contained data. If index is not provided, points to the
68811986Sandreas.sandberg@arm.com    /// beginning of the buffer. May throw if the index would lead to out of bounds access.
68911986Sandreas.sandberg@arm.com    /// May throw if the array is not writeable.
69011986Sandreas.sandberg@arm.com    template<typename... Ix> void* mutable_data(Ix... index) {
69111986Sandreas.sandberg@arm.com        check_writeable();
69211986Sandreas.sandberg@arm.com        return static_cast<void *>(detail::array_proxy(m_ptr)->data + offset_at(index...));
69311986Sandreas.sandberg@arm.com    }
69411986Sandreas.sandberg@arm.com
69511986Sandreas.sandberg@arm.com    /// Byte offset from beginning of the array to a given index (full or partial).
69611986Sandreas.sandberg@arm.com    /// May throw if the index would lead to out of bounds access.
69712391Sjason@lowepower.com    template<typename... Ix> ssize_t offset_at(Ix... index) const {
69812391Sjason@lowepower.com        if ((ssize_t) sizeof...(index) > ndim())
69911986Sandreas.sandberg@arm.com            fail_dim_check(sizeof...(index), "too many indices for an array");
70012391Sjason@lowepower.com        return byte_offset(ssize_t(index)...);
70111986Sandreas.sandberg@arm.com    }
70211986Sandreas.sandberg@arm.com
70312391Sjason@lowepower.com    ssize_t offset_at() const { return 0; }
70411986Sandreas.sandberg@arm.com
70511986Sandreas.sandberg@arm.com    /// Item count from beginning of the array to a given index (full or partial).
70611986Sandreas.sandberg@arm.com    /// May throw if the index would lead to out of bounds access.
70712391Sjason@lowepower.com    template<typename... Ix> ssize_t index_at(Ix... index) const {
70811986Sandreas.sandberg@arm.com        return offset_at(index...) / itemsize();
70911986Sandreas.sandberg@arm.com    }
71011986Sandreas.sandberg@arm.com
71112391Sjason@lowepower.com    /**
71212391Sjason@lowepower.com     * Returns a proxy object that provides access to the array's data without bounds or
71312037Sandreas.sandberg@arm.com     * dimensionality checking.  Will throw if the array is missing the `writeable` flag.  Use with
71412037Sandreas.sandberg@arm.com     * care: the array must not be destroyed or reshaped for the duration of the returned object,
71512037Sandreas.sandberg@arm.com     * and the caller must take care not to access invalid dimensions or dimension indices.
71612037Sandreas.sandberg@arm.com     */
71712391Sjason@lowepower.com    template <typename T, ssize_t Dims = -1> detail::unchecked_mutable_reference<T, Dims> mutable_unchecked() & {
71812391Sjason@lowepower.com        if (Dims >= 0 && ndim() != Dims)
71912037Sandreas.sandberg@arm.com            throw std::domain_error("array has incorrect number of dimensions: " + std::to_string(ndim()) +
72012037Sandreas.sandberg@arm.com                    "; expected " + std::to_string(Dims));
72112037Sandreas.sandberg@arm.com        return detail::unchecked_mutable_reference<T, Dims>(mutable_data(), shape(), strides(), ndim());
72212037Sandreas.sandberg@arm.com    }
72312037Sandreas.sandberg@arm.com
72412391Sjason@lowepower.com    /**
72512391Sjason@lowepower.com     * Returns a proxy object that provides const access to the array's data without bounds or
72612037Sandreas.sandberg@arm.com     * dimensionality checking.  Unlike `mutable_unchecked()`, this does not require that the
72712037Sandreas.sandberg@arm.com     * underlying array have the `writable` flag.  Use with care: the array must not be destroyed or
72812037Sandreas.sandberg@arm.com     * reshaped for the duration of the returned object, and the caller must take care not to access
72912037Sandreas.sandberg@arm.com     * invalid dimensions or dimension indices.
73012037Sandreas.sandberg@arm.com     */
73112391Sjason@lowepower.com    template <typename T, ssize_t Dims = -1> detail::unchecked_reference<T, Dims> unchecked() const & {
73212391Sjason@lowepower.com        if (Dims >= 0 && ndim() != Dims)
73312037Sandreas.sandberg@arm.com            throw std::domain_error("array has incorrect number of dimensions: " + std::to_string(ndim()) +
73412037Sandreas.sandberg@arm.com                    "; expected " + std::to_string(Dims));
73512037Sandreas.sandberg@arm.com        return detail::unchecked_reference<T, Dims>(data(), shape(), strides(), ndim());
73612037Sandreas.sandberg@arm.com    }
73712037Sandreas.sandberg@arm.com
73811986Sandreas.sandberg@arm.com    /// Return a new view with all of the dimensions of length 1 removed
73911986Sandreas.sandberg@arm.com    array squeeze() {
74011986Sandreas.sandberg@arm.com        auto& api = detail::npy_api::get();
74111986Sandreas.sandberg@arm.com        return reinterpret_steal<array>(api.PyArray_Squeeze_(m_ptr));
74211986Sandreas.sandberg@arm.com    }
74311986Sandreas.sandberg@arm.com
74412391Sjason@lowepower.com    /// Resize array to given shape
74512391Sjason@lowepower.com    /// If refcheck is true and more that one reference exist to this array
74612391Sjason@lowepower.com    /// then resize will succeed only if it makes a reshape, i.e. original size doesn't change
74712391Sjason@lowepower.com    void resize(ShapeContainer new_shape, bool refcheck = true) {
74812391Sjason@lowepower.com        detail::npy_api::PyArray_Dims d = {
74912391Sjason@lowepower.com            new_shape->data(), int(new_shape->size())
75012391Sjason@lowepower.com        };
75112391Sjason@lowepower.com        // try to resize, set ordering param to -1 cause it's not used anyway
75212391Sjason@lowepower.com        object new_array = reinterpret_steal<object>(
75312391Sjason@lowepower.com            detail::npy_api::get().PyArray_Resize_(m_ptr, &d, int(refcheck), -1)
75412391Sjason@lowepower.com        );
75512391Sjason@lowepower.com        if (!new_array) throw error_already_set();
75612391Sjason@lowepower.com        if (isinstance<array>(new_array)) { *this = std::move(new_array); }
75712391Sjason@lowepower.com    }
75812391Sjason@lowepower.com
75911986Sandreas.sandberg@arm.com    /// Ensure that the argument is a NumPy array
76011986Sandreas.sandberg@arm.com    /// In case of an error, nullptr is returned and the Python error is cleared.
76111986Sandreas.sandberg@arm.com    static array ensure(handle h, int ExtraFlags = 0) {
76211986Sandreas.sandberg@arm.com        auto result = reinterpret_steal<array>(raw_array(h.ptr(), ExtraFlags));
76311986Sandreas.sandberg@arm.com        if (!result)
76411986Sandreas.sandberg@arm.com            PyErr_Clear();
76511986Sandreas.sandberg@arm.com        return result;
76611986Sandreas.sandberg@arm.com    }
76711986Sandreas.sandberg@arm.com
76811986Sandreas.sandberg@arm.comprotected:
76911986Sandreas.sandberg@arm.com    template<typename, typename> friend struct detail::npy_format_descriptor;
77011986Sandreas.sandberg@arm.com
77112391Sjason@lowepower.com    void fail_dim_check(ssize_t dim, const std::string& msg) const {
77211986Sandreas.sandberg@arm.com        throw index_error(msg + ": " + std::to_string(dim) +
77311986Sandreas.sandberg@arm.com                          " (ndim = " + std::to_string(ndim()) + ")");
77411986Sandreas.sandberg@arm.com    }
77511986Sandreas.sandberg@arm.com
77612391Sjason@lowepower.com    template<typename... Ix> ssize_t byte_offset(Ix... index) const {
77711986Sandreas.sandberg@arm.com        check_dimensions(index...);
77812391Sjason@lowepower.com        return detail::byte_offset_unsafe(strides(), ssize_t(index)...);
77911986Sandreas.sandberg@arm.com    }
78011986Sandreas.sandberg@arm.com
78111986Sandreas.sandberg@arm.com    void check_writeable() const {
78211986Sandreas.sandberg@arm.com        if (!writeable())
78312037Sandreas.sandberg@arm.com            throw std::domain_error("array is not writeable");
78411986Sandreas.sandberg@arm.com    }
78511986Sandreas.sandberg@arm.com
78612391Sjason@lowepower.com    // Default, C-style strides
78712391Sjason@lowepower.com    static std::vector<ssize_t> c_strides(const std::vector<ssize_t> &shape, ssize_t itemsize) {
78811986Sandreas.sandberg@arm.com        auto ndim = shape.size();
78912391Sjason@lowepower.com        std::vector<ssize_t> strides(ndim, itemsize);
79014299Sbbruce@ucdavis.edu        if (ndim > 0)
79114299Sbbruce@ucdavis.edu            for (size_t i = ndim - 1; i > 0; --i)
79214299Sbbruce@ucdavis.edu                strides[i - 1] = strides[i] * shape[i];
79312391Sjason@lowepower.com        return strides;
79412391Sjason@lowepower.com    }
79512391Sjason@lowepower.com
79612391Sjason@lowepower.com    // F-style strides; default when constructing an array_t with `ExtraFlags & f_style`
79712391Sjason@lowepower.com    static std::vector<ssize_t> f_strides(const std::vector<ssize_t> &shape, ssize_t itemsize) {
79812391Sjason@lowepower.com        auto ndim = shape.size();
79912391Sjason@lowepower.com        std::vector<ssize_t> strides(ndim, itemsize);
80012391Sjason@lowepower.com        for (size_t i = 1; i < ndim; ++i)
80112391Sjason@lowepower.com            strides[i] = strides[i - 1] * shape[i - 1];
80211986Sandreas.sandberg@arm.com        return strides;
80311986Sandreas.sandberg@arm.com    }
80411986Sandreas.sandberg@arm.com
80511986Sandreas.sandberg@arm.com    template<typename... Ix> void check_dimensions(Ix... index) const {
80612391Sjason@lowepower.com        check_dimensions_impl(ssize_t(0), shape(), ssize_t(index)...);
80711986Sandreas.sandberg@arm.com    }
80811986Sandreas.sandberg@arm.com
80912391Sjason@lowepower.com    void check_dimensions_impl(ssize_t, const ssize_t*) const { }
81011986Sandreas.sandberg@arm.com
81112391Sjason@lowepower.com    template<typename... Ix> void check_dimensions_impl(ssize_t axis, const ssize_t* shape, ssize_t i, Ix... index) const {
81211986Sandreas.sandberg@arm.com        if (i >= *shape) {
81311986Sandreas.sandberg@arm.com            throw index_error(std::string("index ") + std::to_string(i) +
81411986Sandreas.sandberg@arm.com                              " is out of bounds for axis " + std::to_string(axis) +
81511986Sandreas.sandberg@arm.com                              " with size " + std::to_string(*shape));
81611986Sandreas.sandberg@arm.com        }
81711986Sandreas.sandberg@arm.com        check_dimensions_impl(axis + 1, shape + 1, index...);
81811986Sandreas.sandberg@arm.com    }
81911986Sandreas.sandberg@arm.com
82011986Sandreas.sandberg@arm.com    /// Create array from any object -- always returns a new reference
82111986Sandreas.sandberg@arm.com    static PyObject *raw_array(PyObject *ptr, int ExtraFlags = 0) {
82212391Sjason@lowepower.com        if (ptr == nullptr) {
82312391Sjason@lowepower.com            PyErr_SetString(PyExc_ValueError, "cannot create a pybind11::array from a nullptr");
82411986Sandreas.sandberg@arm.com            return nullptr;
82512391Sjason@lowepower.com        }
82611986Sandreas.sandberg@arm.com        return detail::npy_api::get().PyArray_FromAny_(
82712037Sandreas.sandberg@arm.com            ptr, nullptr, 0, 0, detail::npy_api::NPY_ARRAY_ENSUREARRAY_ | ExtraFlags, nullptr);
82811986Sandreas.sandberg@arm.com    }
82911986Sandreas.sandberg@arm.com};
83011986Sandreas.sandberg@arm.com
83111986Sandreas.sandberg@arm.comtemplate <typename T, int ExtraFlags = array::forcecast> class array_t : public array {
83212391Sjason@lowepower.comprivate:
83312391Sjason@lowepower.com    struct private_ctor {};
83412391Sjason@lowepower.com    // Delegating constructor needed when both moving and accessing in the same constructor
83512391Sjason@lowepower.com    array_t(private_ctor, ShapeContainer &&shape, StridesContainer &&strides, const T *ptr, handle base)
83612391Sjason@lowepower.com        : array(std::move(shape), std::move(strides), ptr, base) {}
83711986Sandreas.sandberg@arm.compublic:
83812391Sjason@lowepower.com    static_assert(!detail::array_info<T>::is_array, "Array types cannot be used with array_t");
83912391Sjason@lowepower.com
84012037Sandreas.sandberg@arm.com    using value_type = T;
84112037Sandreas.sandberg@arm.com
84211986Sandreas.sandberg@arm.com    array_t() : array(0, static_cast<const T *>(nullptr)) {}
84312391Sjason@lowepower.com    array_t(handle h, borrowed_t) : array(h, borrowed_t{}) { }
84412391Sjason@lowepower.com    array_t(handle h, stolen_t) : array(h, stolen_t{}) { }
84511986Sandreas.sandberg@arm.com
84611986Sandreas.sandberg@arm.com    PYBIND11_DEPRECATED("Use array_t<T>::ensure() instead")
84712391Sjason@lowepower.com    array_t(handle h, bool is_borrowed) : array(raw_array_t(h.ptr()), stolen_t{}) {
84811986Sandreas.sandberg@arm.com        if (!m_ptr) PyErr_Clear();
84911986Sandreas.sandberg@arm.com        if (!is_borrowed) Py_XDECREF(h.ptr());
85011986Sandreas.sandberg@arm.com    }
85111986Sandreas.sandberg@arm.com
85212391Sjason@lowepower.com    array_t(const object &o) : array(raw_array_t(o.ptr()), stolen_t{}) {
85311986Sandreas.sandberg@arm.com        if (!m_ptr) throw error_already_set();
85411986Sandreas.sandberg@arm.com    }
85511986Sandreas.sandberg@arm.com
85611986Sandreas.sandberg@arm.com    explicit array_t(const buffer_info& info) : array(info) { }
85711986Sandreas.sandberg@arm.com
85812391Sjason@lowepower.com    array_t(ShapeContainer shape, StridesContainer strides, const T *ptr = nullptr, handle base = handle())
85912391Sjason@lowepower.com        : array(std::move(shape), std::move(strides), ptr, base) { }
86011986Sandreas.sandberg@arm.com
86112391Sjason@lowepower.com    explicit array_t(ShapeContainer shape, const T *ptr = nullptr, handle base = handle())
86212391Sjason@lowepower.com        : array_t(private_ctor{}, std::move(shape),
86312391Sjason@lowepower.com                ExtraFlags & f_style ? f_strides(*shape, itemsize()) : c_strides(*shape, itemsize()),
86412391Sjason@lowepower.com                ptr, base) { }
86511986Sandreas.sandberg@arm.com
86611986Sandreas.sandberg@arm.com    explicit array_t(size_t count, const T *ptr = nullptr, handle base = handle())
86712391Sjason@lowepower.com        : array({count}, {}, ptr, base) { }
86811986Sandreas.sandberg@arm.com
86912391Sjason@lowepower.com    constexpr ssize_t itemsize() const {
87011986Sandreas.sandberg@arm.com        return sizeof(T);
87111986Sandreas.sandberg@arm.com    }
87211986Sandreas.sandberg@arm.com
87312391Sjason@lowepower.com    template<typename... Ix> ssize_t index_at(Ix... index) const {
87411986Sandreas.sandberg@arm.com        return offset_at(index...) / itemsize();
87511986Sandreas.sandberg@arm.com    }
87611986Sandreas.sandberg@arm.com
87711986Sandreas.sandberg@arm.com    template<typename... Ix> const T* data(Ix... index) const {
87811986Sandreas.sandberg@arm.com        return static_cast<const T*>(array::data(index...));
87911986Sandreas.sandberg@arm.com    }
88011986Sandreas.sandberg@arm.com
88111986Sandreas.sandberg@arm.com    template<typename... Ix> T* mutable_data(Ix... index) {
88211986Sandreas.sandberg@arm.com        return static_cast<T*>(array::mutable_data(index...));
88311986Sandreas.sandberg@arm.com    }
88411986Sandreas.sandberg@arm.com
88511986Sandreas.sandberg@arm.com    // Reference to element at a given index
88611986Sandreas.sandberg@arm.com    template<typename... Ix> const T& at(Ix... index) const {
88714299Sbbruce@ucdavis.edu        if ((ssize_t) sizeof...(index) != ndim())
88811986Sandreas.sandberg@arm.com            fail_dim_check(sizeof...(index), "index dimension mismatch");
88912391Sjason@lowepower.com        return *(static_cast<const T*>(array::data()) + byte_offset(ssize_t(index)...) / itemsize());
89011986Sandreas.sandberg@arm.com    }
89111986Sandreas.sandberg@arm.com
89211986Sandreas.sandberg@arm.com    // Mutable reference to element at a given index
89311986Sandreas.sandberg@arm.com    template<typename... Ix> T& mutable_at(Ix... index) {
89414299Sbbruce@ucdavis.edu        if ((ssize_t) sizeof...(index) != ndim())
89511986Sandreas.sandberg@arm.com            fail_dim_check(sizeof...(index), "index dimension mismatch");
89612391Sjason@lowepower.com        return *(static_cast<T*>(array::mutable_data()) + byte_offset(ssize_t(index)...) / itemsize());
89711986Sandreas.sandberg@arm.com    }
89811986Sandreas.sandberg@arm.com
89912391Sjason@lowepower.com    /**
90012391Sjason@lowepower.com     * Returns a proxy object that provides access to the array's data without bounds or
90112037Sandreas.sandberg@arm.com     * dimensionality checking.  Will throw if the array is missing the `writeable` flag.  Use with
90212037Sandreas.sandberg@arm.com     * care: the array must not be destroyed or reshaped for the duration of the returned object,
90312037Sandreas.sandberg@arm.com     * and the caller must take care not to access invalid dimensions or dimension indices.
90412037Sandreas.sandberg@arm.com     */
90512391Sjason@lowepower.com    template <ssize_t Dims = -1> detail::unchecked_mutable_reference<T, Dims> mutable_unchecked() & {
90612037Sandreas.sandberg@arm.com        return array::mutable_unchecked<T, Dims>();
90712037Sandreas.sandberg@arm.com    }
90812037Sandreas.sandberg@arm.com
90912391Sjason@lowepower.com    /**
91012391Sjason@lowepower.com     * Returns a proxy object that provides const access to the array's data without bounds or
91112037Sandreas.sandberg@arm.com     * dimensionality checking.  Unlike `unchecked()`, this does not require that the underlying
91212037Sandreas.sandberg@arm.com     * array have the `writable` flag.  Use with care: the array must not be destroyed or reshaped
91312037Sandreas.sandberg@arm.com     * for the duration of the returned object, and the caller must take care not to access invalid
91412037Sandreas.sandberg@arm.com     * dimensions or dimension indices.
91512037Sandreas.sandberg@arm.com     */
91612391Sjason@lowepower.com    template <ssize_t Dims = -1> detail::unchecked_reference<T, Dims> unchecked() const & {
91712037Sandreas.sandberg@arm.com        return array::unchecked<T, Dims>();
91812037Sandreas.sandberg@arm.com    }
91912037Sandreas.sandberg@arm.com
92012037Sandreas.sandberg@arm.com    /// Ensure that the argument is a NumPy array of the correct dtype (and if not, try to convert
92112037Sandreas.sandberg@arm.com    /// it).  In case of an error, nullptr is returned and the Python error is cleared.
92211986Sandreas.sandberg@arm.com    static array_t ensure(handle h) {
92311986Sandreas.sandberg@arm.com        auto result = reinterpret_steal<array_t>(raw_array_t(h.ptr()));
92411986Sandreas.sandberg@arm.com        if (!result)
92511986Sandreas.sandberg@arm.com            PyErr_Clear();
92611986Sandreas.sandberg@arm.com        return result;
92711986Sandreas.sandberg@arm.com    }
92811986Sandreas.sandberg@arm.com
92912037Sandreas.sandberg@arm.com    static bool check_(handle h) {
93011986Sandreas.sandberg@arm.com        const auto &api = detail::npy_api::get();
93111986Sandreas.sandberg@arm.com        return api.PyArray_Check_(h.ptr())
93211986Sandreas.sandberg@arm.com               && api.PyArray_EquivTypes_(detail::array_proxy(h.ptr())->descr, dtype::of<T>().ptr());
93311986Sandreas.sandberg@arm.com    }
93411986Sandreas.sandberg@arm.com
93511986Sandreas.sandberg@arm.comprotected:
93611986Sandreas.sandberg@arm.com    /// Create array from any object -- always returns a new reference
93711986Sandreas.sandberg@arm.com    static PyObject *raw_array_t(PyObject *ptr) {
93812391Sjason@lowepower.com        if (ptr == nullptr) {
93912391Sjason@lowepower.com            PyErr_SetString(PyExc_ValueError, "cannot create a pybind11::array_t from a nullptr");
94011986Sandreas.sandberg@arm.com            return nullptr;
94112391Sjason@lowepower.com        }
94211986Sandreas.sandberg@arm.com        return detail::npy_api::get().PyArray_FromAny_(
94311986Sandreas.sandberg@arm.com            ptr, dtype::of<T>().release().ptr(), 0, 0,
94412037Sandreas.sandberg@arm.com            detail::npy_api::NPY_ARRAY_ENSUREARRAY_ | ExtraFlags, nullptr);
94511986Sandreas.sandberg@arm.com    }
94611986Sandreas.sandberg@arm.com};
94711986Sandreas.sandberg@arm.com
94811986Sandreas.sandberg@arm.comtemplate <typename T>
94911986Sandreas.sandberg@arm.comstruct format_descriptor<T, detail::enable_if_t<detail::is_pod_struct<T>::value>> {
95011986Sandreas.sandberg@arm.com    static std::string format() {
95111986Sandreas.sandberg@arm.com        return detail::npy_format_descriptor<typename std::remove_cv<T>::type>::format();
95211986Sandreas.sandberg@arm.com    }
95311986Sandreas.sandberg@arm.com};
95411986Sandreas.sandberg@arm.com
95511986Sandreas.sandberg@arm.comtemplate <size_t N> struct format_descriptor<char[N]> {
95611986Sandreas.sandberg@arm.com    static std::string format() { return std::to_string(N) + "s"; }
95711986Sandreas.sandberg@arm.com};
95811986Sandreas.sandberg@arm.comtemplate <size_t N> struct format_descriptor<std::array<char, N>> {
95911986Sandreas.sandberg@arm.com    static std::string format() { return std::to_string(N) + "s"; }
96011986Sandreas.sandberg@arm.com};
96111986Sandreas.sandberg@arm.com
96211986Sandreas.sandberg@arm.comtemplate <typename T>
96311986Sandreas.sandberg@arm.comstruct format_descriptor<T, detail::enable_if_t<std::is_enum<T>::value>> {
96411986Sandreas.sandberg@arm.com    static std::string format() {
96511986Sandreas.sandberg@arm.com        return format_descriptor<
96611986Sandreas.sandberg@arm.com            typename std::remove_cv<typename std::underlying_type<T>::type>::type>::format();
96711986Sandreas.sandberg@arm.com    }
96811986Sandreas.sandberg@arm.com};
96911986Sandreas.sandberg@arm.com
97012391Sjason@lowepower.comtemplate <typename T>
97112391Sjason@lowepower.comstruct format_descriptor<T, detail::enable_if_t<detail::array_info<T>::is_array>> {
97212391Sjason@lowepower.com    static std::string format() {
97314299Sbbruce@ucdavis.edu        using namespace detail;
97414299Sbbruce@ucdavis.edu        static constexpr auto extents = _("(") + array_info<T>::extents + _(")");
97514299Sbbruce@ucdavis.edu        return extents.text + format_descriptor<remove_all_extents_t<T>>::format();
97612391Sjason@lowepower.com    }
97712391Sjason@lowepower.com};
97812391Sjason@lowepower.com
97911986Sandreas.sandberg@arm.comNAMESPACE_BEGIN(detail)
98011986Sandreas.sandberg@arm.comtemplate <typename T, int ExtraFlags>
98111986Sandreas.sandberg@arm.comstruct pyobject_caster<array_t<T, ExtraFlags>> {
98211986Sandreas.sandberg@arm.com    using type = array_t<T, ExtraFlags>;
98311986Sandreas.sandberg@arm.com
98412037Sandreas.sandberg@arm.com    bool load(handle src, bool convert) {
98512037Sandreas.sandberg@arm.com        if (!convert && !type::check_(src))
98612037Sandreas.sandberg@arm.com            return false;
98711986Sandreas.sandberg@arm.com        value = type::ensure(src);
98811986Sandreas.sandberg@arm.com        return static_cast<bool>(value);
98911986Sandreas.sandberg@arm.com    }
99011986Sandreas.sandberg@arm.com
99111986Sandreas.sandberg@arm.com    static handle cast(const handle &src, return_value_policy /* policy */, handle /* parent */) {
99211986Sandreas.sandberg@arm.com        return src.inc_ref();
99311986Sandreas.sandberg@arm.com    }
99414299Sbbruce@ucdavis.edu    PYBIND11_TYPE_CASTER(type, handle_type_name<type>::name);
99511986Sandreas.sandberg@arm.com};
99611986Sandreas.sandberg@arm.com
99711986Sandreas.sandberg@arm.comtemplate <typename T>
99812037Sandreas.sandberg@arm.comstruct compare_buffer_info<T, detail::enable_if_t<detail::is_pod_struct<T>::value>> {
99912037Sandreas.sandberg@arm.com    static bool compare(const buffer_info& b) {
100012037Sandreas.sandberg@arm.com        return npy_api::get().PyArray_EquivTypes_(dtype::of<T>().ptr(), dtype(b).ptr());
100112037Sandreas.sandberg@arm.com    }
100211986Sandreas.sandberg@arm.com};
100311986Sandreas.sandberg@arm.com
100414299Sbbruce@ucdavis.edutemplate <typename T, typename = void>
100514299Sbbruce@ucdavis.edustruct npy_format_descriptor_name;
100614299Sbbruce@ucdavis.edu
100714299Sbbruce@ucdavis.edutemplate <typename T>
100814299Sbbruce@ucdavis.edustruct npy_format_descriptor_name<T, enable_if_t<std::is_integral<T>::value>> {
100914299Sbbruce@ucdavis.edu    static constexpr auto name = _<std::is_same<T, bool>::value>(
101014299Sbbruce@ucdavis.edu        _("bool"), _<std::is_signed<T>::value>("int", "uint") + _<sizeof(T)*8>()
101114299Sbbruce@ucdavis.edu    );
101214299Sbbruce@ucdavis.edu};
101314299Sbbruce@ucdavis.edu
101414299Sbbruce@ucdavis.edutemplate <typename T>
101514299Sbbruce@ucdavis.edustruct npy_format_descriptor_name<T, enable_if_t<std::is_floating_point<T>::value>> {
101614299Sbbruce@ucdavis.edu    static constexpr auto name = _<std::is_same<T, float>::value || std::is_same<T, double>::value>(
101714299Sbbruce@ucdavis.edu        _("float") + _<sizeof(T)*8>(), _("longdouble")
101814299Sbbruce@ucdavis.edu    );
101914299Sbbruce@ucdavis.edu};
102014299Sbbruce@ucdavis.edu
102114299Sbbruce@ucdavis.edutemplate <typename T>
102214299Sbbruce@ucdavis.edustruct npy_format_descriptor_name<T, enable_if_t<is_complex<T>::value>> {
102314299Sbbruce@ucdavis.edu    static constexpr auto name = _<std::is_same<typename T::value_type, float>::value
102414299Sbbruce@ucdavis.edu                                   || std::is_same<typename T::value_type, double>::value>(
102514299Sbbruce@ucdavis.edu        _("complex") + _<sizeof(typename T::value_type)*16>(), _("longcomplex")
102614299Sbbruce@ucdavis.edu    );
102714299Sbbruce@ucdavis.edu};
102814299Sbbruce@ucdavis.edu
102914299Sbbruce@ucdavis.edutemplate <typename T>
103014299Sbbruce@ucdavis.edustruct npy_format_descriptor<T, enable_if_t<satisfies_any_of<T, std::is_arithmetic, is_complex>::value>>
103114299Sbbruce@ucdavis.edu    : npy_format_descriptor_name<T> {
103211986Sandreas.sandberg@arm.comprivate:
103312037Sandreas.sandberg@arm.com    // NB: the order here must match the one in common.h
103412037Sandreas.sandberg@arm.com    constexpr static const int values[15] = {
103512037Sandreas.sandberg@arm.com        npy_api::NPY_BOOL_,
103614299Sbbruce@ucdavis.edu        npy_api::NPY_BYTE_,   npy_api::NPY_UBYTE_,   npy_api::NPY_INT16_,    npy_api::NPY_UINT16_,
103714299Sbbruce@ucdavis.edu        npy_api::NPY_INT32_,  npy_api::NPY_UINT32_,  npy_api::NPY_INT64_,    npy_api::NPY_UINT64_,
103812037Sandreas.sandberg@arm.com        npy_api::NPY_FLOAT_,  npy_api::NPY_DOUBLE_,  npy_api::NPY_LONGDOUBLE_,
103912037Sandreas.sandberg@arm.com        npy_api::NPY_CFLOAT_, npy_api::NPY_CDOUBLE_, npy_api::NPY_CLONGDOUBLE_
104012037Sandreas.sandberg@arm.com    };
104112037Sandreas.sandberg@arm.com
104211986Sandreas.sandberg@arm.compublic:
104312037Sandreas.sandberg@arm.com    static constexpr int value = values[detail::is_fmt_numeric<T>::index];
104412037Sandreas.sandberg@arm.com
104511986Sandreas.sandberg@arm.com    static pybind11::dtype dtype() {
104611986Sandreas.sandberg@arm.com        if (auto ptr = npy_api::get().PyArray_DescrFromType_(value))
104714299Sbbruce@ucdavis.edu            return reinterpret_steal<pybind11::dtype>(ptr);
104811986Sandreas.sandberg@arm.com        pybind11_fail("Unsupported buffer format!");
104911986Sandreas.sandberg@arm.com    }
105011986Sandreas.sandberg@arm.com};
105111986Sandreas.sandberg@arm.com
105212037Sandreas.sandberg@arm.com#define PYBIND11_DECL_CHAR_FMT \
105314299Sbbruce@ucdavis.edu    static constexpr auto name = _("S") + _<N>(); \
105411986Sandreas.sandberg@arm.com    static pybind11::dtype dtype() { return pybind11::dtype(std::string("S") + std::to_string(N)); }
105512037Sandreas.sandberg@arm.comtemplate <size_t N> struct npy_format_descriptor<char[N]> { PYBIND11_DECL_CHAR_FMT };
105612037Sandreas.sandberg@arm.comtemplate <size_t N> struct npy_format_descriptor<std::array<char, N>> { PYBIND11_DECL_CHAR_FMT };
105712037Sandreas.sandberg@arm.com#undef PYBIND11_DECL_CHAR_FMT
105811986Sandreas.sandberg@arm.com
105912391Sjason@lowepower.comtemplate<typename T> struct npy_format_descriptor<T, enable_if_t<array_info<T>::is_array>> {
106012391Sjason@lowepower.comprivate:
106112391Sjason@lowepower.com    using base_descr = npy_format_descriptor<typename array_info<T>::type>;
106212391Sjason@lowepower.compublic:
106312391Sjason@lowepower.com    static_assert(!array_info<T>::is_empty, "Zero-sized arrays are not supported");
106412391Sjason@lowepower.com
106514299Sbbruce@ucdavis.edu    static constexpr auto name = _("(") + array_info<T>::extents + _(")") + base_descr::name;
106612391Sjason@lowepower.com    static pybind11::dtype dtype() {
106712391Sjason@lowepower.com        list shape;
106812391Sjason@lowepower.com        array_info<T>::append_extents(shape);
106912391Sjason@lowepower.com        return pybind11::dtype::from_args(pybind11::make_tuple(base_descr::dtype(), shape));
107012391Sjason@lowepower.com    }
107112391Sjason@lowepower.com};
107212391Sjason@lowepower.com
107311986Sandreas.sandberg@arm.comtemplate<typename T> struct npy_format_descriptor<T, enable_if_t<std::is_enum<T>::value>> {
107411986Sandreas.sandberg@arm.comprivate:
107511986Sandreas.sandberg@arm.com    using base_descr = npy_format_descriptor<typename std::underlying_type<T>::type>;
107611986Sandreas.sandberg@arm.compublic:
107714299Sbbruce@ucdavis.edu    static constexpr auto name = base_descr::name;
107811986Sandreas.sandberg@arm.com    static pybind11::dtype dtype() { return base_descr::dtype(); }
107911986Sandreas.sandberg@arm.com};
108011986Sandreas.sandberg@arm.com
108111986Sandreas.sandberg@arm.comstruct field_descriptor {
108211986Sandreas.sandberg@arm.com    const char *name;
108312391Sjason@lowepower.com    ssize_t offset;
108412391Sjason@lowepower.com    ssize_t size;
108511986Sandreas.sandberg@arm.com    std::string format;
108611986Sandreas.sandberg@arm.com    dtype descr;
108711986Sandreas.sandberg@arm.com};
108811986Sandreas.sandberg@arm.com
108911986Sandreas.sandberg@arm.cominline PYBIND11_NOINLINE void register_structured_dtype(
109014299Sbbruce@ucdavis.edu    any_container<field_descriptor> fields,
109112391Sjason@lowepower.com    const std::type_info& tinfo, ssize_t itemsize,
109211986Sandreas.sandberg@arm.com    bool (*direct_converter)(PyObject *, void *&)) {
109311986Sandreas.sandberg@arm.com
109411986Sandreas.sandberg@arm.com    auto& numpy_internals = get_numpy_internals();
109511986Sandreas.sandberg@arm.com    if (numpy_internals.get_type_info(tinfo, false))
109611986Sandreas.sandberg@arm.com        pybind11_fail("NumPy: dtype is already registered");
109711986Sandreas.sandberg@arm.com
109814299Sbbruce@ucdavis.edu    // Use ordered fields because order matters as of NumPy 1.14:
109914299Sbbruce@ucdavis.edu    // https://docs.scipy.org/doc/numpy/release.html#multiple-field-indexing-assignment-of-structured-arrays
110014299Sbbruce@ucdavis.edu    std::vector<field_descriptor> ordered_fields(std::move(fields));
110114299Sbbruce@ucdavis.edu    std::sort(ordered_fields.begin(), ordered_fields.end(),
110214299Sbbruce@ucdavis.edu        [](const field_descriptor &a, const field_descriptor &b) { return a.offset < b.offset; });
110314299Sbbruce@ucdavis.edu
110411986Sandreas.sandberg@arm.com    list names, formats, offsets;
110514299Sbbruce@ucdavis.edu    for (auto& field : ordered_fields) {
110611986Sandreas.sandberg@arm.com        if (!field.descr)
110711986Sandreas.sandberg@arm.com            pybind11_fail(std::string("NumPy: unsupported field dtype: `") +
110811986Sandreas.sandberg@arm.com                            field.name + "` @ " + tinfo.name());
110911986Sandreas.sandberg@arm.com        names.append(PYBIND11_STR_TYPE(field.name));
111011986Sandreas.sandberg@arm.com        formats.append(field.descr);
111111986Sandreas.sandberg@arm.com        offsets.append(pybind11::int_(field.offset));
111211986Sandreas.sandberg@arm.com    }
111311986Sandreas.sandberg@arm.com    auto dtype_ptr = pybind11::dtype(names, formats, offsets, itemsize).release().ptr();
111411986Sandreas.sandberg@arm.com
111511986Sandreas.sandberg@arm.com    // There is an existing bug in NumPy (as of v1.11): trailing bytes are
111611986Sandreas.sandberg@arm.com    // not encoded explicitly into the format string. This will supposedly
111711986Sandreas.sandberg@arm.com    // get fixed in v1.12; for further details, see these:
111811986Sandreas.sandberg@arm.com    // - https://github.com/numpy/numpy/issues/7797
111911986Sandreas.sandberg@arm.com    // - https://github.com/numpy/numpy/pull/7798
112011986Sandreas.sandberg@arm.com    // Because of this, we won't use numpy's logic to generate buffer format
112111986Sandreas.sandberg@arm.com    // strings and will just do it ourselves.
112212391Sjason@lowepower.com    ssize_t offset = 0;
112311986Sandreas.sandberg@arm.com    std::ostringstream oss;
112412391Sjason@lowepower.com    // mark the structure as unaligned with '^', because numpy and C++ don't
112512391Sjason@lowepower.com    // always agree about alignment (particularly for complex), and we're
112612391Sjason@lowepower.com    // explicitly listing all our padding. This depends on none of the fields
112712391Sjason@lowepower.com    // overriding the endianness. Putting the ^ in front of individual fields
112812391Sjason@lowepower.com    // isn't guaranteed to work due to https://github.com/numpy/numpy/issues/9049
112912391Sjason@lowepower.com    oss << "^T{";
113011986Sandreas.sandberg@arm.com    for (auto& field : ordered_fields) {
113111986Sandreas.sandberg@arm.com        if (field.offset > offset)
113211986Sandreas.sandberg@arm.com            oss << (field.offset - offset) << 'x';
113311986Sandreas.sandberg@arm.com        oss << field.format << ':' << field.name << ':';
113411986Sandreas.sandberg@arm.com        offset = field.offset + field.size;
113511986Sandreas.sandberg@arm.com    }
113611986Sandreas.sandberg@arm.com    if (itemsize > offset)
113711986Sandreas.sandberg@arm.com        oss << (itemsize - offset) << 'x';
113811986Sandreas.sandberg@arm.com    oss << '}';
113911986Sandreas.sandberg@arm.com    auto format_str = oss.str();
114011986Sandreas.sandberg@arm.com
114111986Sandreas.sandberg@arm.com    // Sanity check: verify that NumPy properly parses our buffer format string
114211986Sandreas.sandberg@arm.com    auto& api = npy_api::get();
114311986Sandreas.sandberg@arm.com    auto arr =  array(buffer_info(nullptr, itemsize, format_str, 1));
114411986Sandreas.sandberg@arm.com    if (!api.PyArray_EquivTypes_(dtype_ptr, arr.dtype().ptr()))
114511986Sandreas.sandberg@arm.com        pybind11_fail("NumPy: invalid buffer descriptor!");
114611986Sandreas.sandberg@arm.com
114711986Sandreas.sandberg@arm.com    auto tindex = std::type_index(tinfo);
114811986Sandreas.sandberg@arm.com    numpy_internals.registered_dtypes[tindex] = { dtype_ptr, format_str };
114911986Sandreas.sandberg@arm.com    get_internals().direct_conversions[tindex].push_back(direct_converter);
115011986Sandreas.sandberg@arm.com}
115111986Sandreas.sandberg@arm.com
115212037Sandreas.sandberg@arm.comtemplate <typename T, typename SFINAE> struct npy_format_descriptor {
115312037Sandreas.sandberg@arm.com    static_assert(is_pod_struct<T>::value, "Attempt to use a non-POD or unimplemented POD type as a numpy dtype");
115412037Sandreas.sandberg@arm.com
115514299Sbbruce@ucdavis.edu    static constexpr auto name = make_caster<T>::name;
115611986Sandreas.sandberg@arm.com
115711986Sandreas.sandberg@arm.com    static pybind11::dtype dtype() {
115811986Sandreas.sandberg@arm.com        return reinterpret_borrow<pybind11::dtype>(dtype_ptr());
115911986Sandreas.sandberg@arm.com    }
116011986Sandreas.sandberg@arm.com
116111986Sandreas.sandberg@arm.com    static std::string format() {
116211986Sandreas.sandberg@arm.com        static auto format_str = get_numpy_internals().get_type_info<T>(true)->format_str;
116311986Sandreas.sandberg@arm.com        return format_str;
116411986Sandreas.sandberg@arm.com    }
116511986Sandreas.sandberg@arm.com
116614299Sbbruce@ucdavis.edu    static void register_dtype(any_container<field_descriptor> fields) {
116714299Sbbruce@ucdavis.edu        register_structured_dtype(std::move(fields), typeid(typename std::remove_cv<T>::type),
116811986Sandreas.sandberg@arm.com                                  sizeof(T), &direct_converter);
116911986Sandreas.sandberg@arm.com    }
117011986Sandreas.sandberg@arm.com
117111986Sandreas.sandberg@arm.comprivate:
117211986Sandreas.sandberg@arm.com    static PyObject* dtype_ptr() {
117311986Sandreas.sandberg@arm.com        static PyObject* ptr = get_numpy_internals().get_type_info<T>(true)->dtype_ptr;
117411986Sandreas.sandberg@arm.com        return ptr;
117511986Sandreas.sandberg@arm.com    }
117611986Sandreas.sandberg@arm.com
117711986Sandreas.sandberg@arm.com    static bool direct_converter(PyObject *obj, void*& value) {
117811986Sandreas.sandberg@arm.com        auto& api = npy_api::get();
117911986Sandreas.sandberg@arm.com        if (!PyObject_TypeCheck(obj, api.PyVoidArrType_Type_))
118011986Sandreas.sandberg@arm.com            return false;
118111986Sandreas.sandberg@arm.com        if (auto descr = reinterpret_steal<object>(api.PyArray_DescrFromScalar_(obj))) {
118211986Sandreas.sandberg@arm.com            if (api.PyArray_EquivTypes_(dtype_ptr(), descr.ptr())) {
118311986Sandreas.sandberg@arm.com                value = ((PyVoidScalarObject_Proxy *) obj)->obval;
118411986Sandreas.sandberg@arm.com                return true;
118511986Sandreas.sandberg@arm.com            }
118611986Sandreas.sandberg@arm.com        }
118711986Sandreas.sandberg@arm.com        return false;
118811986Sandreas.sandberg@arm.com    }
118911986Sandreas.sandberg@arm.com};
119011986Sandreas.sandberg@arm.com
119112391Sjason@lowepower.com#ifdef __CLION_IDE__ // replace heavy macro with dummy code for the IDE (doesn't affect code)
119212391Sjason@lowepower.com# define PYBIND11_NUMPY_DTYPE(Type, ...) ((void)0)
119312391Sjason@lowepower.com# define PYBIND11_NUMPY_DTYPE_EX(Type, ...) ((void)0)
119412391Sjason@lowepower.com#else
119512391Sjason@lowepower.com
119611986Sandreas.sandberg@arm.com#define PYBIND11_FIELD_DESCRIPTOR_EX(T, Field, Name)                                          \
119711986Sandreas.sandberg@arm.com    ::pybind11::detail::field_descriptor {                                                    \
119811986Sandreas.sandberg@arm.com        Name, offsetof(T, Field), sizeof(decltype(std::declval<T>().Field)),                  \
119911986Sandreas.sandberg@arm.com        ::pybind11::format_descriptor<decltype(std::declval<T>().Field)>::format(),           \
120011986Sandreas.sandberg@arm.com        ::pybind11::detail::npy_format_descriptor<decltype(std::declval<T>().Field)>::dtype() \
120111986Sandreas.sandberg@arm.com    }
120211986Sandreas.sandberg@arm.com
120311986Sandreas.sandberg@arm.com// Extract name, offset and format descriptor for a struct field
120411986Sandreas.sandberg@arm.com#define PYBIND11_FIELD_DESCRIPTOR(T, Field) PYBIND11_FIELD_DESCRIPTOR_EX(T, Field, #Field)
120511986Sandreas.sandberg@arm.com
120611986Sandreas.sandberg@arm.com// The main idea of this macro is borrowed from https://github.com/swansontec/map-macro
120711986Sandreas.sandberg@arm.com// (C) William Swanson, Paul Fultz
120811986Sandreas.sandberg@arm.com#define PYBIND11_EVAL0(...) __VA_ARGS__
120911986Sandreas.sandberg@arm.com#define PYBIND11_EVAL1(...) PYBIND11_EVAL0 (PYBIND11_EVAL0 (PYBIND11_EVAL0 (__VA_ARGS__)))
121011986Sandreas.sandberg@arm.com#define PYBIND11_EVAL2(...) PYBIND11_EVAL1 (PYBIND11_EVAL1 (PYBIND11_EVAL1 (__VA_ARGS__)))
121111986Sandreas.sandberg@arm.com#define PYBIND11_EVAL3(...) PYBIND11_EVAL2 (PYBIND11_EVAL2 (PYBIND11_EVAL2 (__VA_ARGS__)))
121211986Sandreas.sandberg@arm.com#define PYBIND11_EVAL4(...) PYBIND11_EVAL3 (PYBIND11_EVAL3 (PYBIND11_EVAL3 (__VA_ARGS__)))
121311986Sandreas.sandberg@arm.com#define PYBIND11_EVAL(...)  PYBIND11_EVAL4 (PYBIND11_EVAL4 (PYBIND11_EVAL4 (__VA_ARGS__)))
121411986Sandreas.sandberg@arm.com#define PYBIND11_MAP_END(...)
121511986Sandreas.sandberg@arm.com#define PYBIND11_MAP_OUT
121611986Sandreas.sandberg@arm.com#define PYBIND11_MAP_COMMA ,
121711986Sandreas.sandberg@arm.com#define PYBIND11_MAP_GET_END() 0, PYBIND11_MAP_END
121811986Sandreas.sandberg@arm.com#define PYBIND11_MAP_NEXT0(test, next, ...) next PYBIND11_MAP_OUT
121911986Sandreas.sandberg@arm.com#define PYBIND11_MAP_NEXT1(test, next) PYBIND11_MAP_NEXT0 (test, next, 0)
122011986Sandreas.sandberg@arm.com#define PYBIND11_MAP_NEXT(test, next)  PYBIND11_MAP_NEXT1 (PYBIND11_MAP_GET_END test, next)
122111986Sandreas.sandberg@arm.com#ifdef _MSC_VER // MSVC is not as eager to expand macros, hence this workaround
122211986Sandreas.sandberg@arm.com#define PYBIND11_MAP_LIST_NEXT1(test, next) \
122311986Sandreas.sandberg@arm.com    PYBIND11_EVAL0 (PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0))
122411986Sandreas.sandberg@arm.com#else
122511986Sandreas.sandberg@arm.com#define PYBIND11_MAP_LIST_NEXT1(test, next) \
122611986Sandreas.sandberg@arm.com    PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0)
122711986Sandreas.sandberg@arm.com#endif
122811986Sandreas.sandberg@arm.com#define PYBIND11_MAP_LIST_NEXT(test, next) \
122911986Sandreas.sandberg@arm.com    PYBIND11_MAP_LIST_NEXT1 (PYBIND11_MAP_GET_END test, next)
123011986Sandreas.sandberg@arm.com#define PYBIND11_MAP_LIST0(f, t, x, peek, ...) \
123111986Sandreas.sandberg@arm.com    f(t, x) PYBIND11_MAP_LIST_NEXT (peek, PYBIND11_MAP_LIST1) (f, t, peek, __VA_ARGS__)
123211986Sandreas.sandberg@arm.com#define PYBIND11_MAP_LIST1(f, t, x, peek, ...) \
123311986Sandreas.sandberg@arm.com    f(t, x) PYBIND11_MAP_LIST_NEXT (peek, PYBIND11_MAP_LIST0) (f, t, peek, __VA_ARGS__)
123411986Sandreas.sandberg@arm.com// PYBIND11_MAP_LIST(f, t, a1, a2, ...) expands to f(t, a1), f(t, a2), ...
123511986Sandreas.sandberg@arm.com#define PYBIND11_MAP_LIST(f, t, ...) \
123611986Sandreas.sandberg@arm.com    PYBIND11_EVAL (PYBIND11_MAP_LIST1 (f, t, __VA_ARGS__, (), 0))
123711986Sandreas.sandberg@arm.com
123811986Sandreas.sandberg@arm.com#define PYBIND11_NUMPY_DTYPE(Type, ...) \
123911986Sandreas.sandberg@arm.com    ::pybind11::detail::npy_format_descriptor<Type>::register_dtype \
124014299Sbbruce@ucdavis.edu        (::std::vector<::pybind11::detail::field_descriptor> \
124114299Sbbruce@ucdavis.edu         {PYBIND11_MAP_LIST (PYBIND11_FIELD_DESCRIPTOR, Type, __VA_ARGS__)})
124211986Sandreas.sandberg@arm.com
124311986Sandreas.sandberg@arm.com#ifdef _MSC_VER
124411986Sandreas.sandberg@arm.com#define PYBIND11_MAP2_LIST_NEXT1(test, next) \
124511986Sandreas.sandberg@arm.com    PYBIND11_EVAL0 (PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0))
124611986Sandreas.sandberg@arm.com#else
124711986Sandreas.sandberg@arm.com#define PYBIND11_MAP2_LIST_NEXT1(test, next) \
124811986Sandreas.sandberg@arm.com    PYBIND11_MAP_NEXT0 (test, PYBIND11_MAP_COMMA next, 0)
124911986Sandreas.sandberg@arm.com#endif
125011986Sandreas.sandberg@arm.com#define PYBIND11_MAP2_LIST_NEXT(test, next) \
125111986Sandreas.sandberg@arm.com    PYBIND11_MAP2_LIST_NEXT1 (PYBIND11_MAP_GET_END test, next)
125211986Sandreas.sandberg@arm.com#define PYBIND11_MAP2_LIST0(f, t, x1, x2, peek, ...) \
125311986Sandreas.sandberg@arm.com    f(t, x1, x2) PYBIND11_MAP2_LIST_NEXT (peek, PYBIND11_MAP2_LIST1) (f, t, peek, __VA_ARGS__)
125411986Sandreas.sandberg@arm.com#define PYBIND11_MAP2_LIST1(f, t, x1, x2, peek, ...) \
125511986Sandreas.sandberg@arm.com    f(t, x1, x2) PYBIND11_MAP2_LIST_NEXT (peek, PYBIND11_MAP2_LIST0) (f, t, peek, __VA_ARGS__)
125611986Sandreas.sandberg@arm.com// PYBIND11_MAP2_LIST(f, t, a1, a2, ...) expands to f(t, a1, a2), f(t, a3, a4), ...
125711986Sandreas.sandberg@arm.com#define PYBIND11_MAP2_LIST(f, t, ...) \
125811986Sandreas.sandberg@arm.com    PYBIND11_EVAL (PYBIND11_MAP2_LIST1 (f, t, __VA_ARGS__, (), 0))
125911986Sandreas.sandberg@arm.com
126011986Sandreas.sandberg@arm.com#define PYBIND11_NUMPY_DTYPE_EX(Type, ...) \
126111986Sandreas.sandberg@arm.com    ::pybind11::detail::npy_format_descriptor<Type>::register_dtype \
126214299Sbbruce@ucdavis.edu        (::std::vector<::pybind11::detail::field_descriptor> \
126314299Sbbruce@ucdavis.edu         {PYBIND11_MAP2_LIST (PYBIND11_FIELD_DESCRIPTOR_EX, Type, __VA_ARGS__)})
126411986Sandreas.sandberg@arm.com
126512391Sjason@lowepower.com#endif // __CLION_IDE__
126612391Sjason@lowepower.com
126711986Sandreas.sandberg@arm.comtemplate  <class T>
126811986Sandreas.sandberg@arm.comusing array_iterator = typename std::add_pointer<T>::type;
126911986Sandreas.sandberg@arm.com
127011986Sandreas.sandberg@arm.comtemplate <class T>
127111986Sandreas.sandberg@arm.comarray_iterator<T> array_begin(const buffer_info& buffer) {
127211986Sandreas.sandberg@arm.com    return array_iterator<T>(reinterpret_cast<T*>(buffer.ptr));
127311986Sandreas.sandberg@arm.com}
127411986Sandreas.sandberg@arm.com
127511986Sandreas.sandberg@arm.comtemplate <class T>
127611986Sandreas.sandberg@arm.comarray_iterator<T> array_end(const buffer_info& buffer) {
127711986Sandreas.sandberg@arm.com    return array_iterator<T>(reinterpret_cast<T*>(buffer.ptr) + buffer.size);
127811986Sandreas.sandberg@arm.com}
127911986Sandreas.sandberg@arm.com
128011986Sandreas.sandberg@arm.comclass common_iterator {
128111986Sandreas.sandberg@arm.compublic:
128212391Sjason@lowepower.com    using container_type = std::vector<ssize_t>;
128311986Sandreas.sandberg@arm.com    using value_type = container_type::value_type;
128411986Sandreas.sandberg@arm.com    using size_type = container_type::size_type;
128511986Sandreas.sandberg@arm.com
128611986Sandreas.sandberg@arm.com    common_iterator() : p_ptr(0), m_strides() {}
128711986Sandreas.sandberg@arm.com
128812391Sjason@lowepower.com    common_iterator(void* ptr, const container_type& strides, const container_type& shape)
128911986Sandreas.sandberg@arm.com        : p_ptr(reinterpret_cast<char*>(ptr)), m_strides(strides.size()) {
129011986Sandreas.sandberg@arm.com        m_strides.back() = static_cast<value_type>(strides.back());
129111986Sandreas.sandberg@arm.com        for (size_type i = m_strides.size() - 1; i != 0; --i) {
129211986Sandreas.sandberg@arm.com            size_type j = i - 1;
129311986Sandreas.sandberg@arm.com            value_type s = static_cast<value_type>(shape[i]);
129411986Sandreas.sandberg@arm.com            m_strides[j] = strides[j] + m_strides[i] - strides[i] * s;
129511986Sandreas.sandberg@arm.com        }
129611986Sandreas.sandberg@arm.com    }
129711986Sandreas.sandberg@arm.com
129811986Sandreas.sandberg@arm.com    void increment(size_type dim) {
129911986Sandreas.sandberg@arm.com        p_ptr += m_strides[dim];
130011986Sandreas.sandberg@arm.com    }
130111986Sandreas.sandberg@arm.com
130211986Sandreas.sandberg@arm.com    void* data() const {
130311986Sandreas.sandberg@arm.com        return p_ptr;
130411986Sandreas.sandberg@arm.com    }
130511986Sandreas.sandberg@arm.com
130611986Sandreas.sandberg@arm.comprivate:
130711986Sandreas.sandberg@arm.com    char* p_ptr;
130811986Sandreas.sandberg@arm.com    container_type m_strides;
130911986Sandreas.sandberg@arm.com};
131011986Sandreas.sandberg@arm.com
131111986Sandreas.sandberg@arm.comtemplate <size_t N> class multi_array_iterator {
131211986Sandreas.sandberg@arm.compublic:
131312391Sjason@lowepower.com    using container_type = std::vector<ssize_t>;
131411986Sandreas.sandberg@arm.com
131511986Sandreas.sandberg@arm.com    multi_array_iterator(const std::array<buffer_info, N> &buffers,
131612391Sjason@lowepower.com                         const container_type &shape)
131711986Sandreas.sandberg@arm.com        : m_shape(shape.size()), m_index(shape.size(), 0),
131811986Sandreas.sandberg@arm.com          m_common_iterator() {
131911986Sandreas.sandberg@arm.com
132011986Sandreas.sandberg@arm.com        // Manual copy to avoid conversion warning if using std::copy
132111986Sandreas.sandberg@arm.com        for (size_t i = 0; i < shape.size(); ++i)
132212391Sjason@lowepower.com            m_shape[i] = shape[i];
132311986Sandreas.sandberg@arm.com
132411986Sandreas.sandberg@arm.com        container_type strides(shape.size());
132511986Sandreas.sandberg@arm.com        for (size_t i = 0; i < N; ++i)
132611986Sandreas.sandberg@arm.com            init_common_iterator(buffers[i], shape, m_common_iterator[i], strides);
132711986Sandreas.sandberg@arm.com    }
132811986Sandreas.sandberg@arm.com
132911986Sandreas.sandberg@arm.com    multi_array_iterator& operator++() {
133011986Sandreas.sandberg@arm.com        for (size_t j = m_index.size(); j != 0; --j) {
133111986Sandreas.sandberg@arm.com            size_t i = j - 1;
133211986Sandreas.sandberg@arm.com            if (++m_index[i] != m_shape[i]) {
133311986Sandreas.sandberg@arm.com                increment_common_iterator(i);
133411986Sandreas.sandberg@arm.com                break;
133511986Sandreas.sandberg@arm.com            } else {
133611986Sandreas.sandberg@arm.com                m_index[i] = 0;
133711986Sandreas.sandberg@arm.com            }
133811986Sandreas.sandberg@arm.com        }
133911986Sandreas.sandberg@arm.com        return *this;
134011986Sandreas.sandberg@arm.com    }
134111986Sandreas.sandberg@arm.com
134212391Sjason@lowepower.com    template <size_t K, class T = void> T* data() const {
134312391Sjason@lowepower.com        return reinterpret_cast<T*>(m_common_iterator[K].data());
134411986Sandreas.sandberg@arm.com    }
134511986Sandreas.sandberg@arm.com
134611986Sandreas.sandberg@arm.comprivate:
134711986Sandreas.sandberg@arm.com
134811986Sandreas.sandberg@arm.com    using common_iter = common_iterator;
134911986Sandreas.sandberg@arm.com
135011986Sandreas.sandberg@arm.com    void init_common_iterator(const buffer_info &buffer,
135112391Sjason@lowepower.com                              const container_type &shape,
135212391Sjason@lowepower.com                              common_iter &iterator,
135312391Sjason@lowepower.com                              container_type &strides) {
135411986Sandreas.sandberg@arm.com        auto buffer_shape_iter = buffer.shape.rbegin();
135511986Sandreas.sandberg@arm.com        auto buffer_strides_iter = buffer.strides.rbegin();
135611986Sandreas.sandberg@arm.com        auto shape_iter = shape.rbegin();
135711986Sandreas.sandberg@arm.com        auto strides_iter = strides.rbegin();
135811986Sandreas.sandberg@arm.com
135911986Sandreas.sandberg@arm.com        while (buffer_shape_iter != buffer.shape.rend()) {
136011986Sandreas.sandberg@arm.com            if (*shape_iter == *buffer_shape_iter)
136112391Sjason@lowepower.com                *strides_iter = *buffer_strides_iter;
136211986Sandreas.sandberg@arm.com            else
136311986Sandreas.sandberg@arm.com                *strides_iter = 0;
136411986Sandreas.sandberg@arm.com
136511986Sandreas.sandberg@arm.com            ++buffer_shape_iter;
136611986Sandreas.sandberg@arm.com            ++buffer_strides_iter;
136711986Sandreas.sandberg@arm.com            ++shape_iter;
136811986Sandreas.sandberg@arm.com            ++strides_iter;
136911986Sandreas.sandberg@arm.com        }
137011986Sandreas.sandberg@arm.com
137111986Sandreas.sandberg@arm.com        std::fill(strides_iter, strides.rend(), 0);
137211986Sandreas.sandberg@arm.com        iterator = common_iter(buffer.ptr, strides, shape);
137311986Sandreas.sandberg@arm.com    }
137411986Sandreas.sandberg@arm.com
137511986Sandreas.sandberg@arm.com    void increment_common_iterator(size_t dim) {
137611986Sandreas.sandberg@arm.com        for (auto &iter : m_common_iterator)
137711986Sandreas.sandberg@arm.com            iter.increment(dim);
137811986Sandreas.sandberg@arm.com    }
137911986Sandreas.sandberg@arm.com
138011986Sandreas.sandberg@arm.com    container_type m_shape;
138111986Sandreas.sandberg@arm.com    container_type m_index;
138211986Sandreas.sandberg@arm.com    std::array<common_iter, N> m_common_iterator;
138311986Sandreas.sandberg@arm.com};
138411986Sandreas.sandberg@arm.com
138512037Sandreas.sandberg@arm.comenum class broadcast_trivial { non_trivial, c_trivial, f_trivial };
138612037Sandreas.sandberg@arm.com
138712037Sandreas.sandberg@arm.com// Populates the shape and number of dimensions for the set of buffers.  Returns a broadcast_trivial
138812037Sandreas.sandberg@arm.com// enum value indicating whether the broadcast is "trivial"--that is, has each buffer being either a
138912037Sandreas.sandberg@arm.com// singleton or a full-size, C-contiguous (`c_trivial`) or Fortran-contiguous (`f_trivial`) storage
139012037Sandreas.sandberg@arm.com// buffer; returns `non_trivial` otherwise.
139111986Sandreas.sandberg@arm.comtemplate <size_t N>
139212391Sjason@lowepower.combroadcast_trivial broadcast(const std::array<buffer_info, N> &buffers, ssize_t &ndim, std::vector<ssize_t> &shape) {
139312391Sjason@lowepower.com    ndim = std::accumulate(buffers.begin(), buffers.end(), ssize_t(0), [](ssize_t res, const buffer_info &buf) {
139411986Sandreas.sandberg@arm.com        return std::max(res, buf.ndim);
139511986Sandreas.sandberg@arm.com    });
139611986Sandreas.sandberg@arm.com
139712037Sandreas.sandberg@arm.com    shape.clear();
139812391Sjason@lowepower.com    shape.resize((size_t) ndim, 1);
139912037Sandreas.sandberg@arm.com
140012037Sandreas.sandberg@arm.com    // Figure out the output size, and make sure all input arrays conform (i.e. are either size 1 or
140112037Sandreas.sandberg@arm.com    // the full size).
140211986Sandreas.sandberg@arm.com    for (size_t i = 0; i < N; ++i) {
140311986Sandreas.sandberg@arm.com        auto res_iter = shape.rbegin();
140412037Sandreas.sandberg@arm.com        auto end = buffers[i].shape.rend();
140512037Sandreas.sandberg@arm.com        for (auto shape_iter = buffers[i].shape.rbegin(); shape_iter != end; ++shape_iter, ++res_iter) {
140612037Sandreas.sandberg@arm.com            const auto &dim_size_in = *shape_iter;
140712037Sandreas.sandberg@arm.com            auto &dim_size_out = *res_iter;
140811986Sandreas.sandberg@arm.com
140912037Sandreas.sandberg@arm.com            // Each input dimension can either be 1 or `n`, but `n` values must match across buffers
141012037Sandreas.sandberg@arm.com            if (dim_size_out == 1)
141112037Sandreas.sandberg@arm.com                dim_size_out = dim_size_in;
141212037Sandreas.sandberg@arm.com            else if (dim_size_in != 1 && dim_size_in != dim_size_out)
141311986Sandreas.sandberg@arm.com                pybind11_fail("pybind11::vectorize: incompatible size/dimension of inputs!");
141412037Sandreas.sandberg@arm.com        }
141512037Sandreas.sandberg@arm.com    }
141611986Sandreas.sandberg@arm.com
141712037Sandreas.sandberg@arm.com    bool trivial_broadcast_c = true;
141812037Sandreas.sandberg@arm.com    bool trivial_broadcast_f = true;
141912037Sandreas.sandberg@arm.com    for (size_t i = 0; i < N && (trivial_broadcast_c || trivial_broadcast_f); ++i) {
142012037Sandreas.sandberg@arm.com        if (buffers[i].size == 1)
142112037Sandreas.sandberg@arm.com            continue;
142212037Sandreas.sandberg@arm.com
142312037Sandreas.sandberg@arm.com        // Require the same number of dimensions:
142412037Sandreas.sandberg@arm.com        if (buffers[i].ndim != ndim)
142512037Sandreas.sandberg@arm.com            return broadcast_trivial::non_trivial;
142612037Sandreas.sandberg@arm.com
142712037Sandreas.sandberg@arm.com        // Require all dimensions be full-size:
142812037Sandreas.sandberg@arm.com        if (!std::equal(buffers[i].shape.cbegin(), buffers[i].shape.cend(), shape.cbegin()))
142912037Sandreas.sandberg@arm.com            return broadcast_trivial::non_trivial;
143012037Sandreas.sandberg@arm.com
143112037Sandreas.sandberg@arm.com        // Check for C contiguity (but only if previous inputs were also C contiguous)
143212037Sandreas.sandberg@arm.com        if (trivial_broadcast_c) {
143312391Sjason@lowepower.com            ssize_t expect_stride = buffers[i].itemsize;
143412037Sandreas.sandberg@arm.com            auto end = buffers[i].shape.crend();
143512037Sandreas.sandberg@arm.com            for (auto shape_iter = buffers[i].shape.crbegin(), stride_iter = buffers[i].strides.crbegin();
143612037Sandreas.sandberg@arm.com                    trivial_broadcast_c && shape_iter != end; ++shape_iter, ++stride_iter) {
143712037Sandreas.sandberg@arm.com                if (expect_stride == *stride_iter)
143812037Sandreas.sandberg@arm.com                    expect_stride *= *shape_iter;
143912037Sandreas.sandberg@arm.com                else
144012037Sandreas.sandberg@arm.com                    trivial_broadcast_c = false;
144112037Sandreas.sandberg@arm.com            }
144211986Sandreas.sandberg@arm.com        }
144312037Sandreas.sandberg@arm.com
144412037Sandreas.sandberg@arm.com        // Check for Fortran contiguity (if previous inputs were also F contiguous)
144512037Sandreas.sandberg@arm.com        if (trivial_broadcast_f) {
144612391Sjason@lowepower.com            ssize_t expect_stride = buffers[i].itemsize;
144712037Sandreas.sandberg@arm.com            auto end = buffers[i].shape.cend();
144812037Sandreas.sandberg@arm.com            for (auto shape_iter = buffers[i].shape.cbegin(), stride_iter = buffers[i].strides.cbegin();
144912037Sandreas.sandberg@arm.com                    trivial_broadcast_f && shape_iter != end; ++shape_iter, ++stride_iter) {
145012037Sandreas.sandberg@arm.com                if (expect_stride == *stride_iter)
145112037Sandreas.sandberg@arm.com                    expect_stride *= *shape_iter;
145212037Sandreas.sandberg@arm.com                else
145312037Sandreas.sandberg@arm.com                    trivial_broadcast_f = false;
145412037Sandreas.sandberg@arm.com            }
145512037Sandreas.sandberg@arm.com        }
145611986Sandreas.sandberg@arm.com    }
145712037Sandreas.sandberg@arm.com
145812037Sandreas.sandberg@arm.com    return
145912037Sandreas.sandberg@arm.com        trivial_broadcast_c ? broadcast_trivial::c_trivial :
146012037Sandreas.sandberg@arm.com        trivial_broadcast_f ? broadcast_trivial::f_trivial :
146112037Sandreas.sandberg@arm.com        broadcast_trivial::non_trivial;
146211986Sandreas.sandberg@arm.com}
146311986Sandreas.sandberg@arm.com
146412391Sjason@lowepower.comtemplate <typename T>
146512391Sjason@lowepower.comstruct vectorize_arg {
146612391Sjason@lowepower.com    static_assert(!std::is_rvalue_reference<T>::value, "Functions with rvalue reference arguments cannot be vectorized");
146712391Sjason@lowepower.com    // The wrapped function gets called with this type:
146812391Sjason@lowepower.com    using call_type = remove_reference_t<T>;
146912391Sjason@lowepower.com    // Is this a vectorized argument?
147012391Sjason@lowepower.com    static constexpr bool vectorize =
147112391Sjason@lowepower.com        satisfies_any_of<call_type, std::is_arithmetic, is_complex, std::is_pod>::value &&
147212391Sjason@lowepower.com        satisfies_none_of<call_type, std::is_pointer, std::is_array, is_std_array, std::is_enum>::value &&
147312391Sjason@lowepower.com        (!std::is_reference<T>::value ||
147412391Sjason@lowepower.com         (std::is_lvalue_reference<T>::value && std::is_const<call_type>::value));
147512391Sjason@lowepower.com    // Accept this type: an array for vectorized types, otherwise the type as-is:
147612391Sjason@lowepower.com    using type = conditional_t<vectorize, array_t<remove_cv_t<call_type>, array::forcecast>, T>;
147712391Sjason@lowepower.com};
147812391Sjason@lowepower.com
147911986Sandreas.sandberg@arm.comtemplate <typename Func, typename Return, typename... Args>
148011986Sandreas.sandberg@arm.comstruct vectorize_helper {
148112391Sjason@lowepower.comprivate:
148212037Sandreas.sandberg@arm.com    static constexpr size_t N = sizeof...(Args);
148312391Sjason@lowepower.com    static constexpr size_t NVectorized = constexpr_sum(vectorize_arg<Args>::vectorize...);
148412391Sjason@lowepower.com    static_assert(NVectorized >= 1,
148512391Sjason@lowepower.com            "pybind11::vectorize(...) requires a function with at least one vectorizable argument");
148611986Sandreas.sandberg@arm.com
148712391Sjason@lowepower.compublic:
148811986Sandreas.sandberg@arm.com    template <typename T>
148912391Sjason@lowepower.com    explicit vectorize_helper(T &&f) : f(std::forward<T>(f)) { }
149011986Sandreas.sandberg@arm.com
149112391Sjason@lowepower.com    object operator()(typename vectorize_arg<Args>::type... args) {
149212391Sjason@lowepower.com        return run(args...,
149312391Sjason@lowepower.com                   make_index_sequence<N>(),
149412391Sjason@lowepower.com                   select_indices<vectorize_arg<Args>::vectorize...>(),
149512391Sjason@lowepower.com                   make_index_sequence<NVectorized>());
149611986Sandreas.sandberg@arm.com    }
149711986Sandreas.sandberg@arm.com
149812391Sjason@lowepower.comprivate:
149912391Sjason@lowepower.com    remove_reference_t<Func> f;
150012391Sjason@lowepower.com
150114299Sbbruce@ucdavis.edu    // Internal compiler error in MSVC 19.16.27025.1 (Visual Studio 2017 15.9.4), when compiling with "/permissive-" flag
150214299Sbbruce@ucdavis.edu    // when arg_call_types is manually inlined.
150314299Sbbruce@ucdavis.edu    using arg_call_types = std::tuple<typename vectorize_arg<Args>::call_type...>;
150414299Sbbruce@ucdavis.edu    template <size_t Index> using param_n_t = typename std::tuple_element<Index, arg_call_types>::type;
150512391Sjason@lowepower.com
150612391Sjason@lowepower.com    // Runs a vectorized function given arguments tuple and three index sequences:
150712391Sjason@lowepower.com    //     - Index is the full set of 0 ... (N-1) argument indices;
150812391Sjason@lowepower.com    //     - VIndex is the subset of argument indices with vectorized parameters, letting us access
150912391Sjason@lowepower.com    //       vectorized arguments (anything not in this sequence is passed through)
151012391Sjason@lowepower.com    //     - BIndex is a incremental sequence (beginning at 0) of the same size as VIndex, so that
151112391Sjason@lowepower.com    //       we can store vectorized buffer_infos in an array (argument VIndex has its buffer at
151212391Sjason@lowepower.com    //       index BIndex in the array).
151312391Sjason@lowepower.com    template <size_t... Index, size_t... VIndex, size_t... BIndex> object run(
151412391Sjason@lowepower.com            typename vectorize_arg<Args>::type &...args,
151512391Sjason@lowepower.com            index_sequence<Index...> i_seq, index_sequence<VIndex...> vi_seq, index_sequence<BIndex...> bi_seq) {
151612391Sjason@lowepower.com
151712391Sjason@lowepower.com        // Pointers to values the function was called with; the vectorized ones set here will start
151812391Sjason@lowepower.com        // out as array_t<T> pointers, but they will be changed them to T pointers before we make
151912391Sjason@lowepower.com        // call the wrapped function.  Non-vectorized pointers are left as-is.
152012391Sjason@lowepower.com        std::array<void *, N> params{{ &args... }};
152112391Sjason@lowepower.com
152212391Sjason@lowepower.com        // The array of `buffer_info`s of vectorized arguments:
152312391Sjason@lowepower.com        std::array<buffer_info, NVectorized> buffers{{ reinterpret_cast<array *>(params[VIndex])->request()... }};
152411986Sandreas.sandberg@arm.com
152511986Sandreas.sandberg@arm.com        /* Determine dimensions parameters of output array */
152612391Sjason@lowepower.com        ssize_t nd = 0;
152712391Sjason@lowepower.com        std::vector<ssize_t> shape(0);
152812391Sjason@lowepower.com        auto trivial = broadcast(buffers, nd, shape);
152912391Sjason@lowepower.com        size_t ndim = (size_t) nd;
153011986Sandreas.sandberg@arm.com
153112391Sjason@lowepower.com        size_t size = std::accumulate(shape.begin(), shape.end(), (size_t) 1, std::multiplies<size_t>());
153212391Sjason@lowepower.com
153312391Sjason@lowepower.com        // If all arguments are 0-dimension arrays (i.e. single values) return a plain value (i.e.
153412391Sjason@lowepower.com        // not wrapped in an array).
153512391Sjason@lowepower.com        if (size == 1 && ndim == 0) {
153612391Sjason@lowepower.com            PYBIND11_EXPAND_SIDE_EFFECTS(params[VIndex] = buffers[BIndex].ptr);
153712391Sjason@lowepower.com            return cast(f(*reinterpret_cast<param_n_t<Index> *>(params[Index])...));
153811986Sandreas.sandberg@arm.com        }
153911986Sandreas.sandberg@arm.com
154012391Sjason@lowepower.com        array_t<Return> result;
154112391Sjason@lowepower.com        if (trivial == broadcast_trivial::f_trivial) result = array_t<Return, array::f_style>(shape);
154212391Sjason@lowepower.com        else result = array_t<Return>(shape);
154311986Sandreas.sandberg@arm.com
154414299Sbbruce@ucdavis.edu        if (size == 0) return std::move(result);
154511986Sandreas.sandberg@arm.com
154612037Sandreas.sandberg@arm.com        /* Call the function */
154712391Sjason@lowepower.com        if (trivial == broadcast_trivial::non_trivial)
154812391Sjason@lowepower.com            apply_broadcast(buffers, params, result, i_seq, vi_seq, bi_seq);
154912391Sjason@lowepower.com        else
155012391Sjason@lowepower.com            apply_trivial(buffers, params, result.mutable_data(), size, i_seq, vi_seq, bi_seq);
155111986Sandreas.sandberg@arm.com
155214299Sbbruce@ucdavis.edu        return std::move(result);
155311986Sandreas.sandberg@arm.com    }
155411986Sandreas.sandberg@arm.com
155512391Sjason@lowepower.com    template <size_t... Index, size_t... VIndex, size_t... BIndex>
155612391Sjason@lowepower.com    void apply_trivial(std::array<buffer_info, NVectorized> &buffers,
155712391Sjason@lowepower.com                       std::array<void *, N> &params,
155812391Sjason@lowepower.com                       Return *out,
155912391Sjason@lowepower.com                       size_t size,
156012391Sjason@lowepower.com                       index_sequence<Index...>, index_sequence<VIndex...>, index_sequence<BIndex...>) {
156111986Sandreas.sandberg@arm.com
156212391Sjason@lowepower.com        // Initialize an array of mutable byte references and sizes with references set to the
156312391Sjason@lowepower.com        // appropriate pointer in `params`; as we iterate, we'll increment each pointer by its size
156412391Sjason@lowepower.com        // (except for singletons, which get an increment of 0).
156512391Sjason@lowepower.com        std::array<std::pair<unsigned char *&, const size_t>, NVectorized> vecparams{{
156612391Sjason@lowepower.com            std::pair<unsigned char *&, const size_t>(
156712391Sjason@lowepower.com                    reinterpret_cast<unsigned char *&>(params[VIndex] = buffers[BIndex].ptr),
156812391Sjason@lowepower.com                    buffers[BIndex].size == 1 ? 0 : sizeof(param_n_t<VIndex>)
156912391Sjason@lowepower.com            )...
157012391Sjason@lowepower.com        }};
157111986Sandreas.sandberg@arm.com
157212391Sjason@lowepower.com        for (size_t i = 0; i < size; ++i) {
157312391Sjason@lowepower.com            out[i] = f(*reinterpret_cast<param_n_t<Index> *>(params[Index])...);
157412391Sjason@lowepower.com            for (auto &x : vecparams) x.first += x.second;
157512391Sjason@lowepower.com        }
157612391Sjason@lowepower.com    }
157712391Sjason@lowepower.com
157812391Sjason@lowepower.com    template <size_t... Index, size_t... VIndex, size_t... BIndex>
157912391Sjason@lowepower.com    void apply_broadcast(std::array<buffer_info, NVectorized> &buffers,
158012391Sjason@lowepower.com                         std::array<void *, N> &params,
158112391Sjason@lowepower.com                         array_t<Return> &output_array,
158212391Sjason@lowepower.com                         index_sequence<Index...>, index_sequence<VIndex...>, index_sequence<BIndex...>) {
158312391Sjason@lowepower.com
158412391Sjason@lowepower.com        buffer_info output = output_array.request();
158512391Sjason@lowepower.com        multi_array_iterator<NVectorized> input_iter(buffers, output.shape);
158612391Sjason@lowepower.com
158712391Sjason@lowepower.com        for (array_iterator<Return> iter = array_begin<Return>(output), end = array_end<Return>(output);
158812391Sjason@lowepower.com             iter != end;
158912391Sjason@lowepower.com             ++iter, ++input_iter) {
159012391Sjason@lowepower.com            PYBIND11_EXPAND_SIDE_EFFECTS((
159112391Sjason@lowepower.com                params[VIndex] = input_iter.template data<BIndex>()
159212391Sjason@lowepower.com            ));
159312391Sjason@lowepower.com            *iter = f(*reinterpret_cast<param_n_t<Index> *>(std::get<Index>(params))...);
159411986Sandreas.sandberg@arm.com        }
159511986Sandreas.sandberg@arm.com    }
159611986Sandreas.sandberg@arm.com};
159711986Sandreas.sandberg@arm.com
159812391Sjason@lowepower.comtemplate <typename Func, typename Return, typename... Args>
159912391Sjason@lowepower.comvectorize_helper<Func, Return, Args...>
160012391Sjason@lowepower.comvectorize_extractor(const Func &f, Return (*) (Args ...)) {
160112391Sjason@lowepower.com    return detail::vectorize_helper<Func, Return, Args...>(f);
160212391Sjason@lowepower.com}
160312391Sjason@lowepower.com
160411986Sandreas.sandberg@arm.comtemplate <typename T, int Flags> struct handle_type_name<array_t<T, Flags>> {
160514299Sbbruce@ucdavis.edu    static constexpr auto name = _("numpy.ndarray[") + npy_format_descriptor<T>::name + _("]");
160611986Sandreas.sandberg@arm.com};
160711986Sandreas.sandberg@arm.com
160811986Sandreas.sandberg@arm.comNAMESPACE_END(detail)
160911986Sandreas.sandberg@arm.com
161012391Sjason@lowepower.com// Vanilla pointer vectorizer:
161112391Sjason@lowepower.comtemplate <typename Return, typename... Args>
161212391Sjason@lowepower.comdetail::vectorize_helper<Return (*)(Args...), Return, Args...>
161312391Sjason@lowepower.comvectorize(Return (*f) (Args ...)) {
161412391Sjason@lowepower.com    return detail::vectorize_helper<Return (*)(Args...), Return, Args...>(f);
161511986Sandreas.sandberg@arm.com}
161611986Sandreas.sandberg@arm.com
161712391Sjason@lowepower.com// lambda vectorizer:
161812391Sjason@lowepower.comtemplate <typename Func, detail::enable_if_t<detail::is_lambda<Func>::value, int> = 0>
161912391Sjason@lowepower.comauto vectorize(Func &&f) -> decltype(
162012391Sjason@lowepower.com        detail::vectorize_extractor(std::forward<Func>(f), (detail::function_signature_t<Func> *) nullptr)) {
162112391Sjason@lowepower.com    return detail::vectorize_extractor(std::forward<Func>(f), (detail::function_signature_t<Func> *) nullptr);
162211986Sandreas.sandberg@arm.com}
162311986Sandreas.sandberg@arm.com
162412391Sjason@lowepower.com// Vectorize a class method (non-const):
162512391Sjason@lowepower.comtemplate <typename Return, typename Class, typename... Args,
162612391Sjason@lowepower.com          typename Helper = detail::vectorize_helper<decltype(std::mem_fn(std::declval<Return (Class::*)(Args...)>())), Return, Class *, Args...>>
162712391Sjason@lowepower.comHelper vectorize(Return (Class::*f)(Args...)) {
162812391Sjason@lowepower.com    return Helper(std::mem_fn(f));
162911986Sandreas.sandberg@arm.com}
163011986Sandreas.sandberg@arm.com
163114299Sbbruce@ucdavis.edu// Vectorize a class method (const):
163212391Sjason@lowepower.comtemplate <typename Return, typename Class, typename... Args,
163312391Sjason@lowepower.com          typename Helper = detail::vectorize_helper<decltype(std::mem_fn(std::declval<Return (Class::*)(Args...) const>())), Return, const Class *, Args...>>
163412391Sjason@lowepower.comHelper vectorize(Return (Class::*f)(Args...) const) {
163512391Sjason@lowepower.com    return Helper(std::mem_fn(f));
163612391Sjason@lowepower.com}
163712391Sjason@lowepower.com
163812391Sjason@lowepower.comNAMESPACE_END(PYBIND11_NAMESPACE)
163911986Sandreas.sandberg@arm.com
164011986Sandreas.sandberg@arm.com#if defined(_MSC_VER)
164111986Sandreas.sandberg@arm.com#pragma warning(pop)
164211986Sandreas.sandberg@arm.com#endif
1643