diff --git a/Include/internal/pycore_pybuffer.h b/Include/internal/pycore_pybuffer.h index 9439d2bd770587..26f68946074b0c 100644 --- a/Include/internal/pycore_pybuffer.h +++ b/Include/internal/pycore_pybuffer.h @@ -15,6 +15,23 @@ PyAPI_FUNC(int) _PyBuffer_ReleaseInInterpreter( PyAPI_FUNC(int) _PyBuffer_ReleaseInInterpreterAndRawFree( PyInterpreterState *interp, Py_buffer *view); +/* Search the items of the one-dimensional buffer of obj for value, comparing + their bytes when that is equivalent to comparing them as objects. + + _PyBuffer_FindItem() returns the index of the first item in [start, stop) + equal to value, -1 if there is none, or -2 with an exception set on error. + Negative start and stop count from the end, and are clipped like slice + bounds. + + _PyBuffer_CountItem() returns the number of items equal to value as an + int, or NULL with an exception set on error. + + Implemented in memoryobject.c. */ +PyAPI_FUNC(Py_ssize_t) _PyBuffer_FindItem( + PyObject *obj, PyObject *value, + Py_ssize_t start, Py_ssize_t stop); +PyAPI_FUNC(PyObject *) _PyBuffer_CountItem(PyObject *obj, PyObject *value); + #ifdef __cplusplus } #endif diff --git a/Lib/test/seq_tests.py b/Lib/test/seq_tests.py index b7875fe8f2f75d..500e3247cfa18b 100644 --- a/Lib/test/seq_tests.py +++ b/Lib/test/seq_tests.py @@ -3,10 +3,12 @@ """ import unittest +import struct import sys import pickle from test import support from test.support import ALWAYS_EQ, NEVER_EQ +from test.support import warnings_helper # Various iterables # This is used for checking the constructor (here and in test_deque.py) @@ -442,3 +444,80 @@ def test_pickle(self): def test_free_after_iterating(self): support.check_free_after_iterating(self, iter, self.type2test) support.check_free_after_iterating(self, reversed, self.type2test) + + +class NativeItemsSearchTest: + """Tests of count(), index() and "in" for 1-d buffer protocol objects. + + Tests use "needle" for the value being looked for and "haystack" for + where to search. + """ + + # All the native single item formats of memoryview. + search_formats = ('b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q', + 'n', 'N', 'P', 'e', 'f', 'd', 'Zf', 'Zd', '?', 'c') + + def make_haystacks(self, fmt, raw): + """Make type-specific haystacks interpreting *raw* as *fmt*. + + Subclasses make one haystack per layout under test (for example a + contiguous, strided or unaligned memoryview). Yield (name, haystack) + pairs, where name identifies the haystack in subtests. + """ + raise NotImplementedError + + # Comparing bytes needles with int items emits BytesWarning with -b. + @warnings_helper.ignore_warnings(category=BytesWarning) + def run_search_cases(self, check): + """Call check(haystack, haystack_objects, needle) for every needle + of every haystack, with a subtest for each haystack. + + haystack_objects is the list of the items of haystack as Python + objects, to check results against. + """ + class AlwaysEqual(int): + def __eq__(self, other): + return True + __hash__ = int.__hash__ + + nan = float('nan') + # Every byte value, then the doubles 0.0, -0.0, NaN and 1.0, all + # twice so that values are found more than once. + raw = (bytes(range(256)) + struct.pack('4d', 0.0, -0.0, nan, 1)) * 2 + # Needles which may not be items of the format: of other types, out + # of range, or equal to items with different bytes. + other_needles = [0, 1, -1, 2**64, 0.0, -0.0, nan, True, False, 1.0, + b'a', AlwaysEqual(5)] + for fmt in self.search_formats: + for name, haystack in self.make_haystacks(fmt, raw): + with self.subTest(fmt=fmt, haystack=name): + haystack_objects = haystack.tolist() + for needle in haystack_objects + other_needles: + check(haystack, haystack_objects, needle) + + def test_search_formats(self): + self.run_search_cases(self.check_search) + + def check_search(self, haystack, haystack_objects, needle): + """Validate optimized search matches python object search. + + Compares search results between optimized format search and python + rich object comparison. + """ + msg = f'needle={needle!r}' + n = len(haystack_objects) + matches = [i for i, x in enumerate(haystack_objects) if x == needle] + self.assertEqual(haystack.count(needle), len(matches), msg) + self.assertEqual(needle in haystack, bool(matches), msg) + # index() bounds which are negative or past either end are clipped + # like slice bounds. + bounds = [(0, n), (5, -5), (-3, n), (2, n + 100), (-n - 5, n // 2)] + for start, stop in bounds: + r = range(n)[start:stop] + expected = [i for i in matches if i in r] + if expected: + self.assertEqual(haystack.index(needle, start, stop), + expected[0], msg) + else: + with self.assertRaises(ValueError, msg=msg): + haystack.index(needle, start, stop) diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py index 580966589fdd6b..600b10b6f7d68e 100755 --- a/Lib/test/test_array.py +++ b/Lib/test/test_array.py @@ -5,7 +5,8 @@ import collections.abc import unittest from test import support -from test.support import import_helper, script_helper +from test.seq_tests import NativeItemsSearchTest +from test.support import import_helper, script_helper, warnings_helper from test.support import os_helper from test.support import _2G from test.support import subTests @@ -130,6 +131,22 @@ def test_typecodes_documented(self): self.assertGreaterEqual(array.array(typecode).itemsize, minimum_size) + def test_search_mutated_during_compare(self): + # The buffer is held while searching, so the array cannot be resized + # by a comparison. + for search in (array.array.index, array.array.count, + array.array.remove, array.array.__contains__): + for mutate in (lambda a: a.append(9.0), lambda a: a.clear()): + with self.subTest(search=search.__name__): + a = array.array('d', [1.0, 2.0, 3.0]) + class Mutate: + def __eq__(self, other): + mutate(a) + return False + self.assertRaises(BufferError, search, a, Mutate()) + self.assertEqual(a.tolist(), [1.0, 2.0, 3.0]) + a.append(4.0) # The buffer was released. + # Machine format codes. # @@ -167,6 +184,47 @@ def test_typecodes_documented(self): MACHINE_FORMAT_CODE_MAX = 27 +class SearchTest(NativeItemsSearchTest, unittest.TestCase): + # 'w' is tested by test_search_unicode. The raw is not valid UCS4 so + # conversion to a list fails. + search_formats = tuple(fmt for fmt in array.typecodes if fmt != 'w') + + def make_haystacks(self, typecode, raw): + yield 'array', array.array(typecode, raw) + + def test_search_remove(self): + self.run_search_cases(self.check_remove) + + def check_remove(self, a, haystack_objects, needle): + """Check a.remove(needle) against removing the first item of + haystack_objects equal to needle.""" + matches = [i for i, x in enumerate(haystack_objects) if x == needle] + b = array.array(a.typecode, a) + if matches: + b.remove(needle) + after = array.array(a.typecode, a) + del after[matches[0]] + self.assertEqual(b.tobytes(), after.tobytes(), repr(needle)) + else: + with self.assertRaises(ValueError, msg=repr(needle)): + b.remove(needle) + + # Comparing bytes needles with str items emits BytesWarning with -b. + @warnings_helper.ignore_warnings(category=BytesWarning) + def test_search_unicode(self): + a = array.array('w', 'abc\U0010ffff\0ab' * 2) + haystack_objects = a.tolist() + for needle in ['a', 'b', '\U0010ffff', '\0', 'z', 'ab', '', b'a', 97]: + self.check_search(a, haystack_objects, needle) + self.check_remove(a, haystack_objects, needle) + # memoryview does not support 'w' items, but can search them. + m = memoryview(a) + self.assertEqual(m.count('a'), a.count('a')) + self.assertEqual(m.index('\0', 1), a.index('\0', 1)) + self.assertIn('\U0010ffff', m) + self.assertNotIn('z', m) + self.assertRaises(NotImplementedError, m.__getitem__, 0) + class ArrayReconstructorTest(unittest.TestCase): diff --git a/Lib/test/test_buffer.py b/Lib/test/test_buffer.py index f42a71cea319b6..713121eb67adfd 100644 --- a/Lib/test/test_buffer.py +++ b/Lib/test/test_buffer.py @@ -2963,6 +2963,16 @@ def assert_out_of_bounds_error(self, dim): self.assertEqual(str(cm.exception), "index out of bounds on dimension %d" % (dim,)) + def test_memoryview_search_unsupported_format(self): + nd = ndarray([1, 2], shape=[2], format=' None: diff --git a/Lib/test/test_memoryview.py b/Lib/test/test_memoryview.py index 40d4798efbf18f..9572c39807b68c 100644 --- a/Lib/test/test_memoryview.py +++ b/Lib/test/test_memoryview.py @@ -17,6 +17,7 @@ from itertools import product from test import support +from test.seq_tests import NativeItemsSearchTest from test.support import import_helper, threading_helper @@ -949,6 +950,26 @@ def test_overflows_in_floats(self): complex_view[0] = 1e300j self.assertEqual(complex_view[0], 123+321j) + def test_search_released_during_compare(self): + # The buffer is held while searching, so a comparison cannot release + # the memoryview. Leaving the with block fails if it is still held. + for name in 'index', 'count', '__contains__': + with memoryview(array.array('d', [1.0, 2.0, 3.0])) as m: + class Release: + def __eq__(self, other): + m.release() + return False + with self.assertRaises(BufferError, msg=name): + getattr(m, name)(Release()) + + +class SearchTest(NativeItemsSearchTest, unittest.TestCase): + def make_haystacks(self, fmt, raw): + mv = memoryview(raw).cast(fmt) + yield 'contiguous', mv + yield 'strided', mv[::3] + yield 'unaligned', memoryview(b'\0' + raw)[1:].cast(fmt) + @threading_helper.requires_working_threading() @support.requires_resource("cpu") diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2026-10-21-14-04-17.gh-issue-159716.mGLRXB.rst b/Misc/NEWS.d/next/Core_and_Builtins/2026-10-21-14-04-17.gh-issue-159716.mGLRXB.rst new file mode 100644 index 00000000000000..934e42156e7247 --- /dev/null +++ b/Misc/NEWS.d/next/Core_and_Builtins/2026-10-21-14-04-17.gh-issue-159716.mGLRXB.rst @@ -0,0 +1,6 @@ +Speed up ``in``, :meth:`memoryview.count` and :meth:`memoryview.index` for +:class:`memoryview` objects with integer formats. Attempts to resize during +these operations may now raise :exc:`BufferError`. Searching is 60-130x faster +for 100,000 items. :meth:`memoryview.index` now raises +:exc:`NotImplementedError` for an unsupported format even when there are no +items to search, like :meth:`~memoryview.count` and ``in``. diff --git a/Misc/NEWS.d/next/Library/2026-10-23-14-54-46.gh-issue-159716.ZKNJS2.rst b/Misc/NEWS.d/next/Library/2026-10-23-14-54-46.gh-issue-159716.ZKNJS2.rst new file mode 100644 index 00000000000000..ffddc3b79cf4c5 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-10-23-14-54-46.gh-issue-159716.ZKNJS2.rst @@ -0,0 +1,4 @@ +Speed up ``in``, :meth:`array.array.count`, :meth:`array.array.index` and +:meth:`array.array.remove` for :class:`array.array` objects with integer +typecodes. Attempts to resize during these operations may now raise +:exc:`BufferError`. Searching is 120-350x faster for 100,000 items. diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c index 64dacd1b8f1f61..c3d9166ec408ab 100644 --- a/Modules/arraymodule.c +++ b/Modules/arraymodule.c @@ -15,6 +15,7 @@ #include "pycore_modsupport.h" // _PyArg_NoKeywords() #include "pycore_moduleobject.h" // _PyModule_GetState() #include "pycore_object.h" // _PyObject_GC_TRACK() +#include "pycore_pybuffer.h" // _PyBuffer_FindItem() #include "pycore_tuple.h" // _PyTuple_FromPairSteal #include "pycore_weakref.h" // FT_CLEAR_WEAKREFS() @@ -1313,24 +1314,7 @@ static PyObject * array_array_count_impl(arrayobject *self, PyObject *v) /*[clinic end generated code: output=93ead26a2affb739 input=d9bce9d65e39d1f5]*/ { - Py_ssize_t count = 0; - Py_ssize_t i; - - for (i = 0; i < Py_SIZE(self); i++) { - PyObject *selfi; - int cmp; - - selfi = getarrayitem((PyObject *)self, i); - if (selfi == NULL) - return NULL; - cmp = PyObject_RichCompareBool(selfi, v, Py_EQ); - Py_DECREF(selfi); - if (cmp > 0) - count++; - else if (cmp < 0) - return NULL; - } - return PyLong_FromSsize_t(count); + return _PyBuffer_CountItem((PyObject *)self, v); } @@ -1352,50 +1336,25 @@ array_array_index_impl(arrayobject *self, PyObject *v, Py_ssize_t start, Py_ssize_t stop) /*[clinic end generated code: output=c45e777880c99f52 input=089dff7baa7e5a7e]*/ { - if (start < 0) { - start += Py_SIZE(self); - if (start < 0) { - start = 0; - } - } - if (stop < 0) { - stop += Py_SIZE(self); + Py_ssize_t i = _PyBuffer_FindItem((PyObject *)self, v, start, stop); + if (i == -2) { + return NULL; } - // Use Py_SIZE() for every iteration in case the array is mutated - // during PyObject_RichCompareBool() - for (Py_ssize_t i = start; i < stop && i < Py_SIZE(self); i++) { - PyObject *selfi; - int cmp; - - selfi = getarrayitem((PyObject *)self, i); - if (selfi == NULL) - return NULL; - cmp = PyObject_RichCompareBool(selfi, v, Py_EQ); - Py_DECREF(selfi); - if (cmp > 0) { - return PyLong_FromSsize_t(i); - } - else if (cmp < 0) - return NULL; + if (i == -1) { + PyErr_SetString(PyExc_ValueError, "array.index(x): x not in array"); + return NULL; } - PyErr_SetString(PyExc_ValueError, "array.index(x): x not in array"); - return NULL; + return PyLong_FromSsize_t(i); } static int array_contains(PyObject *self, PyObject *v) { - Py_ssize_t i; - int cmp; - - for (i = 0, cmp = 0 ; cmp == 0 && i < Py_SIZE(self); i++) { - PyObject *selfi = getarrayitem(self, i); - if (selfi == NULL) - return -1; - cmp = PyObject_RichCompareBool(selfi, v, Py_EQ); - Py_DECREF(selfi); + Py_ssize_t i = _PyBuffer_FindItem(self, v, 0, PY_SSIZE_T_MAX); + if (i == -2) { + return -1; } - return cmp; + return i >= 0; } /*[clinic input] @@ -1411,27 +1370,18 @@ static PyObject * array_array_remove_impl(arrayobject *self, PyObject *v) /*[clinic end generated code: output=f2a24e288ecb2a35 input=0b1e5aed25590027]*/ { - Py_ssize_t i; - - for (i = 0; i < Py_SIZE(self); i++) { - PyObject *selfi; - int cmp; - - selfi = getarrayitem((PyObject *)self,i); - if (selfi == NULL) - return NULL; - cmp = PyObject_RichCompareBool(selfi, v, Py_EQ); - Py_DECREF(selfi); - if (cmp > 0) { - if (array_del_slice(self, i, i+1) != 0) - return NULL; - Py_RETURN_NONE; - } - else if (cmp < 0) - return NULL; + Py_ssize_t i = _PyBuffer_FindItem((PyObject *)self, v, 0, PY_SSIZE_T_MAX); + if (i == -2) { + return NULL; } - PyErr_SetString(PyExc_ValueError, "array.remove(x): x not in array"); - return NULL; + if (i == -1) { + PyErr_SetString(PyExc_ValueError, "array.remove(x): x not in array"); + return NULL; + } + if (array_del_slice(self, i, i+1) != 0) { + return NULL; + } + Py_RETURN_NONE; } /*[clinic input] diff --git a/Objects/memoryobject.c b/Objects/memoryobject.c index e6d49b6c29a3e9..47929ee3b51fa0 100644 --- a/Objects/memoryobject.c +++ b/Objects/memoryobject.c @@ -12,8 +12,10 @@ #include "Python.h" #include "pycore_abstract.h" // _PyIndex_Check() +#include "pycore_long.h" // _PyLong_IsNegative() #include "pycore_memoryobject.h" // _PyManagedBuffer_Type #include "pycore_object.h" // _PyObject_GC_UNTRACK() +#include "pycore_pybuffer.h" // _PyBuffer_FindItem() #include "pycore_strhex.h" // _Py_strhex_with_sep() #include // offsetof() @@ -1086,6 +1088,294 @@ PyBuffer_ToContiguous(void *buf, const Py_buffer *src, Py_ssize_t len, char orde return ret; } +/****************************************************************************/ +/* Searching buffers */ +/****************************************************************************/ + +/* Defined with the memoryview pack and unpack functions below. */ +static inline Py_ssize_t get_native_fmtchar(const char **result, + const char *fmt); +static inline const char *adjust_fmt(const Py_buffer *view); +static int memory_getbuf_unrestricted(PyMemoryViewObject *self, + Py_buffer *view, int flags); +static inline PyObject *unpack_single(const char *ptr, const char *fmt); +static inline int is_int_fmtchar(char fmtchar); +static Py_ssize_t pack_int(char *buf, PyObject *v, char fmtchar, + Py_ssize_t size); + +/* Instantiate the stringlib search for items of 1, 2, 4 and 8 bytes. */ + +#define FASTSEARCH item1_fastsearch +#define STRINGLIB(F) item1_##F +#define STRINGLIB_CHAR uint8_t +#define STRINGLIB_SIZEOF_CHAR 1 +#define STRINGLIB_FAST_MEMCHR memchr +#include "stringlib/fastsearch.h" +#include "stringlib/undef.h" + +#define FASTSEARCH item2_fastsearch +#define STRINGLIB(F) item2_##F +#define STRINGLIB_CHAR uint16_t +#define STRINGLIB_SIZEOF_CHAR 2 +#include "stringlib/fastsearch.h" +#include "stringlib/undef.h" + +#define FASTSEARCH item4_fastsearch +#define STRINGLIB(F) item4_##F +#define STRINGLIB_CHAR uint32_t +#define STRINGLIB_SIZEOF_CHAR 4 +#include "stringlib/fastsearch.h" +#include "stringlib/undef.h" + +#define FASTSEARCH item8_fastsearch +#define STRINGLIB(F) item8_##F +#define STRINGLIB_CHAR uint64_t +#define STRINGLIB_SIZEOF_CHAR 8 +#include "stringlib/fastsearch.h" +#include "stringlib/undef.h" + +typedef union { + uint8_t u8; + uint16_t u16; + uint32_t u32; + uint64_t u64; +} item_needle; + +typedef enum { + NEEDLE_COMPARE_OBJECTS, /* the items have to be compared as objects */ + NEEDLE_NO_MATCH, /* no item can equal the value */ + NEEDLE_PACKED, /* the value was packed into the needle */ +} needle_result; + +/* Check whether the items of view can be value compared by their bytes. + + They can when the items are in a native integer format and value is an + exact int or a bool. This is not so for floats (0.0 == -0.0) or for int + subclasses, which can override __eq__(). The items must also be + C-contiguous and aligned, since the stringlib search reads them as uint8_t + to uint64_t. */ +static needle_result +pack_needle(const Py_buffer *view, PyObject *value, item_needle *needle) +{ + const char *fmt; + Py_ssize_t itemsize = view->itemsize; + if (!(PyLong_CheckExact(value) || PyBool_Check(value)) + || get_native_fmtchar(&fmt, view->format ? view->format : "B") != itemsize + || !is_int_fmtchar(fmt[0]) + || !PyBuffer_IsContiguous(view, 'C') + || !_Py_IS_ALIGNED(view->buf, itemsize)) + { + return NEEDLE_COMPARE_OBJECTS; + } + assert(itemsize <= (Py_ssize_t)sizeof(*needle)); + if (pack_int((char *)needle, value, fmt[0], itemsize) < 0) { + return NEEDLE_NO_MATCH; + } + return NEEDLE_PACKED; +} + +/* Return the index of the first of the n items at s equal to *needle, or -1 + if there is none. */ +static Py_ssize_t +find_needle(const char *s, Py_ssize_t n, Py_ssize_t itemsize, + const item_needle *needle) +{ + switch (itemsize) { + case 1: return item1_find_char((const uint8_t *)s, n, needle->u8); + case 2: return item2_find_char((const uint16_t *)s, n, needle->u16); + case 4: return item4_find_char((const uint32_t *)s, n, needle->u32); + case 8: return item8_find_char((const uint64_t *)s, n, needle->u64); + } + Py_UNREACHABLE(); +} + +/* Return the number of the n items at s equal to *needle. */ +static Py_ssize_t +count_needle(const char *s, Py_ssize_t n, Py_ssize_t itemsize, + const item_needle *needle) +{ + switch (itemsize) { + case 1: + return item1_count_char_no_maxcount((const uint8_t *)s, n, needle->u8); + case 2: + return item2_count_char_no_maxcount((const uint16_t *)s, n, needle->u16); + case 4: + return item4_count_char_no_maxcount((const uint32_t *)s, n, needle->u32); + case 8: + return item8_count_char_no_maxcount((const uint64_t *)s, n, needle->u64); + } + Py_UNREACHABLE(); +} + +/* Unpack the item at ptr like unpack_single(), but also support the 'w' + format, which array.array exports for Py_UCS4 items. unpack_single() does + not, as memoryview's equality uses struct, which has no 'w' format. */ +static PyObject * +unpack_item(const Py_buffer *view, const char *ptr, const char *fmt) +{ + if (fmt[0] == 'w' && fmt[1] == '\0' + && view->itemsize == (Py_ssize_t)sizeof(Py_UCS4)) + { + Py_UCS4 u; + memcpy(&u, ptr, sizeof(u)); + return PyUnicode_FromOrdinal(u); + } + return unpack_single(ptr, fmt); +} + +/* Compare item i of the one-dimensional view with value as objects. fmt is + the format of view as returned by adjust_fmt(). Return 1 if they are + equal, 0 if not, or -1 with an exception set on error. */ +static int +compare_item(const Py_buffer *view, const char *fmt, Py_ssize_t i, + PyObject *value) +{ + const char *ptr = (const char *)view->buf + i * view->strides[0]; + ptr = ADJUST_PTR(ptr, view->suboffsets, 0); + PyObject *item = unpack_item(view, ptr, fmt); + if (item == NULL) { + return -1; + } + int cmp = PyObject_RichCompareBool(item, value, Py_EQ); + Py_DECREF(item); + return cmp; +} + +/* Get the buffer of obj for searching. Set *fmt for compare_item(). + + The caller holds the buffer for the whole search, so a comparison which + tries to resize obj or release it raises BufferError instead of leaving + view->buf pointing at freed memory. + + A restricted memoryview, passed to __release_buffer__(), refuses to export + new views of its buffer, since they could outlive it. The buffer of the + search does not, so it is taken without that check. */ +static int +get_items_buffer(PyObject *obj, Py_buffer *view, const char **fmt) +{ + int res = PyMemoryView_Check(obj) + ? memory_getbuf_unrestricted((PyMemoryViewObject *)obj, view, + PyBUF_FULL_RO) + : PyObject_GetBuffer(obj, view, PyBUF_FULL_RO); + if (res < 0) { + return -1; + } + if (view->ndim != 1) { + if (view->ndim == 0) { + PyErr_SetString(PyExc_TypeError, + "invalid lookup on 0-dim memory"); + } + else { + PyErr_SetString(PyExc_NotImplementedError, + "multi-dimensional lookup is not implemented"); + } + PyBuffer_Release(view); + return -1; + } + *fmt = view->format ? adjust_fmt(view) : "B"; + if (*fmt == NULL) { + PyBuffer_Release(view); + return -1; + } + return 0; +} + +/* Return the index of the first item of view in [start, stop) equal to + value, -1 if there is none, or -2 with an exception set on error. fmt is + the format of view from get_items_buffer(). */ +static Py_ssize_t +find_item(const Py_buffer *view, const char *fmt, PyObject *value, + Py_ssize_t start, Py_ssize_t stop) +{ + if (start >= stop) { + return -1; + } + item_needle needle; + needle_result res = pack_needle(view, value, &needle); + if (res == NEEDLE_NO_MATCH) { + return -1; + } + // Searchable without making full objects. + if (res == NEEDLE_PACKED) { + Py_ssize_t i = find_needle((const char *)view->buf + + start * view->itemsize, + stop - start, view->itemsize, &needle); + return i < 0 ? -1 : start + i; + } + + for (Py_ssize_t i = start; i < stop; i++) { + int cmp = compare_item(view, fmt, i, value); + if (cmp != 0) { + return cmp > 0 ? i : -2; + } + } + return -1; +} + +/* Return the number of items of view equal to value, or -1 with an + exception set on error. fmt is the format of view from + get_items_buffer(). */ +static Py_ssize_t +count_item(const Py_buffer *view, const char *fmt, PyObject *value) +{ + Py_ssize_t n = view->shape[0]; + if (n == 0) { + return 0; + } + item_needle needle; + needle_result res = pack_needle(view, value, &needle); + if (res == NEEDLE_NO_MATCH) { + return 0; + } + if (res == NEEDLE_PACKED) { + return count_needle(view->buf, n, view->itemsize, &needle); + } + + Py_ssize_t count = 0; + for (Py_ssize_t i = 0; i < n; i++) { + int cmp = compare_item(view, fmt, i, value); + if (cmp < 0) { + return -1; + } + count += cmp; // no overflow since count <= n <= PY_SSIZE_T_MAX + } + return count; +} + +Py_ssize_t +_PyBuffer_FindItem(PyObject *obj, PyObject *value, + Py_ssize_t start, Py_ssize_t stop) +{ + Py_buffer view; + const char *fmt; + if (get_items_buffer(obj, &view, &fmt) < 0) { + return -2; + } + Py_ssize_t n = view.shape[0]; + if (start < 0) { + start = Py_MAX(start + n, 0); + } + if (stop < 0) { + stop = Py_MAX(stop + n, 0); + } + Py_ssize_t i = find_item(&view, fmt, value, start, Py_MIN(stop, n)); + PyBuffer_Release(&view); + return i; +} + +PyObject * +_PyBuffer_CountItem(PyObject *obj, PyObject *value) +{ + Py_buffer view; + const char *fmt; + if (get_items_buffer(obj, &view, &fmt) < 0) { + return NULL; + } + Py_ssize_t count = count_item(&view, fmt, value); + PyBuffer_Release(&view); + return count < 0 ? NULL : PyLong_FromSsize_t(count); +} + static inline Py_ssize_t get_exports(PyMemoryViewObject *buf) { @@ -1586,11 +1876,19 @@ static int memory_getbuf(PyObject *_self, Py_buffer *view, int flags) { PyMemoryViewObject *self = (PyMemoryViewObject *)_self; + CHECK_RESTRICTED_INT(self); + return memory_getbuf_unrestricted(self, view, flags); +} + +/* memory_getbuf() without the check for a restricted memoryview. */ +static int +memory_getbuf_unrestricted(PyMemoryViewObject *self, Py_buffer *view, + int flags) +{ Py_buffer *base = &self->view; int baseflags = self->flags; CHECK_RELEASED_INT(self); - CHECK_RESTRICTED_INT(self); /* start with complete information */ *view = *base; @@ -1725,97 +2023,6 @@ fix_error_int(const char *fmt) static const _Bool bool_false = 0; #define UNPACK_TO_BOOL(PTR) (memcmp((PTR), &bool_false, sizeof(_Bool)) != 0) -/* Accept integer objects or objects with an __index__() method. */ -static long -pylong_as_ld(PyObject *item) -{ - PyObject *tmp; - long ld; - - tmp = _PyNumber_Index(item); - if (tmp == NULL) - return -1; - - ld = PyLong_AsLong(tmp); - Py_DECREF(tmp); - return ld; -} - -static unsigned long -pylong_as_lu(PyObject *item) -{ - PyObject *tmp; - unsigned long lu; - - tmp = _PyNumber_Index(item); - if (tmp == NULL) - return (unsigned long)-1; - - lu = PyLong_AsUnsignedLong(tmp); - Py_DECREF(tmp); - return lu; -} - -static long long -pylong_as_lld(PyObject *item) -{ - PyObject *tmp; - long long lld; - - tmp = _PyNumber_Index(item); - if (tmp == NULL) - return -1; - - lld = PyLong_AsLongLong(tmp); - Py_DECREF(tmp); - return lld; -} - -static unsigned long long -pylong_as_llu(PyObject *item) -{ - PyObject *tmp; - unsigned long long llu; - - tmp = _PyNumber_Index(item); - if (tmp == NULL) - return (unsigned long long)-1; - - llu = PyLong_AsUnsignedLongLong(tmp); - Py_DECREF(tmp); - return llu; -} - -static Py_ssize_t -pylong_as_zd(PyObject *item) -{ - PyObject *tmp; - Py_ssize_t zd; - - tmp = _PyNumber_Index(item); - if (tmp == NULL) - return -1; - - zd = PyLong_AsSsize_t(tmp); - Py_DECREF(tmp); - return zd; -} - -static size_t -pylong_as_zu(PyObject *item) -{ - PyObject *tmp; - size_t zu; - - tmp = _PyNumber_Index(item); - if (tmp == NULL) - return (size_t)-1; - - zu = PyLong_AsSize_t(tmp); - Py_DECREF(tmp); - return zu; -} - /* Timings with the ndarray from _testbuffer.c indicate that using the struct module is around 15x slower than the two functions below. */ @@ -1830,7 +2037,7 @@ pylong_as_zu(PyObject *item) module syntax. This function is very sensitive to small changes. With this layout gcc automatically generates a fast jump table. */ static inline PyObject * -unpack_single(PyMemoryViewObject *self, const char *ptr, const char *fmt) +unpack_single(const char *ptr, const char *fmt) { unsigned long long llu; unsigned long lu; @@ -1842,8 +2049,6 @@ unpack_single(PyMemoryViewObject *self, const char *ptr, const char *fmt) unsigned char uc; void *p; - CHECK_RELEASED_AGAIN(self); - #if PY_LITTLE_ENDIAN int endian = 1; #else @@ -1946,17 +2151,69 @@ unpack_single(PyMemoryViewObject *self, const char *ptr, const char *fmt) memcpy(ptr, (char *)&x, sizeof x); \ } while (0) +/* Return 1 if fmtchar is a native integer format character, else 0. */ +static inline int +is_int_fmtchar(char fmtchar) +{ + switch (fmtchar) { + case 'b': case 'h': case 'i': case 'l': case 'q': case 'n': + case 'B': case 'H': case 'I': case 'L': case 'Q': case 'N': + return 1; + default: + return 0; + } +} + +/* Pack the int v into buf as a native integer of size bytes, signed if + fmtchar is lowercase. buf must have room for 8 bytes. Return size, or -1 + without setting an exception if v is out of range. */ +static Py_ssize_t +pack_int(char *buf, PyObject *v, char fmtchar, Py_ssize_t size) +{ + assert(PyLong_Check(v)); + assert(is_int_fmtchar(fmtchar)); + assert(0 < size && size <= 8); + int is_signed = Py_ISLOWER(fmtchar); + if (!is_signed && _PyLong_IsNegative((PyLongObject *)v)) { + return -1; + } + + if (_PyLong_IsCompact((PyLongObject *)v)) { + /* Fast path: a compact int has at most PyLong_SHIFT bits plus a sign, + and PyLong_SHIFT < 32, so it fits in any item of 4 or more bytes. + Negative values for unsigned formats were rejected above. */ + Py_BUILD_ASSERT(PyLong_SHIFT < 32); + Py_ssize_t ival = _PyLong_CompactValue((PyLongObject *)v); + if (size < 4) { + Py_ssize_t limit = (Py_ssize_t)1 << (size * 8 - is_signed); + if (ival >= limit || ival < (is_signed ? -limit : 0)) { + return -1; + } + } + switch (size) { + case 1: PACK_SINGLE(buf, ival, uint8_t); break; + case 2: PACK_SINGLE(buf, ival, uint16_t); break; + case 4: PACK_SINGLE(buf, ival, uint32_t); break; + default: PACK_SINGLE(buf, ival, uint64_t); break; + } + return size; + } + + int flags = Py_ASNATIVEBYTES_NATIVE_ENDIAN; + if (!is_signed) { + flags |= Py_ASNATIVEBYTES_UNSIGNED_BUFFER; + } + Py_ssize_t needed = PyLong_AsNativeBytes(v, buf, size, flags); + assert(needed > 0); + return needed <= size ? size : -1; +} + /* Pack a single item. 'fmt' can be any native format in struct module syntax. */ static int pack_single(PyMemoryViewObject *self, char *ptr, PyObject *item, const char *fmt) { - unsigned long long llu; - unsigned long lu; - size_t zu; - long long lld; long ld; - Py_ssize_t zd; double d; Py_complex c; void *p; @@ -1967,79 +2224,25 @@ pack_single(PyMemoryViewObject *self, char *ptr, PyObject *item, const char *fmt int endian = 0; #endif switch (fmt[0]) { - /* signed integers */ - case 'b': case 'h': case 'i': case 'l': - ld = pylong_as_ld(item); - if (ld == -1 && PyErr_Occurred()) + /* integers */ + case 'b': case 'h': case 'i': case 'l': case 'q': case 'n': + case 'B': case 'H': case 'I': case 'L': case 'Q': case 'N': { + char buf[8]; + const char *f; + Py_ssize_t size = get_native_fmtchar(&f, fmt); + PyObject *tmp = _PyNumber_Index(item); + if (tmp == NULL) { goto err_occurred; - CHECK_RELEASED_INT_AGAIN(self); - switch (fmt[0]) { - case 'b': - if (ld < SCHAR_MIN || ld > SCHAR_MAX) goto err_range; - *((signed char *)ptr) = (signed char)ld; break; - case 'h': - if (ld < SHRT_MIN || ld > SHRT_MAX) goto err_range; - PACK_SINGLE(ptr, ld, short); break; - case 'i': - if (ld < INT_MIN || ld > INT_MAX) goto err_range; - PACK_SINGLE(ptr, ld, int); break; - default: /* 'l' */ - PACK_SINGLE(ptr, ld, long); break; } - break; - - /* unsigned integers */ - case 'B': case 'H': case 'I': case 'L': - lu = pylong_as_lu(item); - if (lu == (unsigned long)-1 && PyErr_Occurred()) - goto err_occurred; + size = pack_int(buf, tmp, f[0], size); + Py_DECREF(tmp); CHECK_RELEASED_INT_AGAIN(self); - switch (fmt[0]) { - case 'B': - if (lu > UCHAR_MAX) goto err_range; - *((unsigned char *)ptr) = (unsigned char)lu; break; - case 'H': - if (lu > USHRT_MAX) goto err_range; - PACK_SINGLE(ptr, lu, unsigned short); break; - case 'I': - if (lu > UINT_MAX) goto err_range; - PACK_SINGLE(ptr, lu, unsigned int); break; - default: /* 'L' */ - PACK_SINGLE(ptr, lu, unsigned long); break; + if (size < 0) { + goto err_range; } + memcpy(ptr, buf, size); break; - - /* native 64-bit */ - case 'q': - lld = pylong_as_lld(item); - if (lld == -1 && PyErr_Occurred()) - goto err_occurred; - CHECK_RELEASED_INT_AGAIN(self); - PACK_SINGLE(ptr, lld, long long); - break; - case 'Q': - llu = pylong_as_llu(item); - if (llu == (unsigned long long)-1 && PyErr_Occurred()) - goto err_occurred; - CHECK_RELEASED_INT_AGAIN(self); - PACK_SINGLE(ptr, llu, unsigned long long); - break; - - /* ssize_t and size_t */ - case 'n': - zd = pylong_as_zd(item); - if (zd == -1 && PyErr_Occurred()) - goto err_occurred; - CHECK_RELEASED_INT_AGAIN(self); - PACK_SINGLE(ptr, zd, Py_ssize_t); - break; - case 'N': - zu = pylong_as_zu(item); - if (zu == (size_t)-1 && PyErr_Occurred()) - goto err_occurred; - CHECK_RELEASED_INT_AGAIN(self); - PACK_SINGLE(ptr, zu, size_t); - break; + } /* floats */ case 'f': case 'd': case 'e': @@ -2292,7 +2495,7 @@ tolist_base(PyMemoryViewObject *self, const char *ptr, const Py_ssize_t *shape, for (i = 0; i < shape[0]; ptr+=strides[0], i++) { const char *xptr = ADJUST_PTR(ptr, suboffsets, 0); - item = unpack_single(self, xptr, fmt); + item = unpack_single(xptr, fmt); if (item == NULL) { Py_DECREF(lst); return NULL; @@ -2360,7 +2563,7 @@ memoryview_tolist_impl(PyMemoryViewObject *self) if (fmt == NULL) return NULL; if (view->ndim == 0) { - return unpack_single(self, view->buf, fmt); + return unpack_single(view->buf, fmt); } else if (view->ndim == 1) { return tolist_base(self, view->buf, view->shape, @@ -2584,7 +2787,7 @@ memory_item(PyObject *_self, Py_ssize_t index) char *ptr = ptr_from_index(view, index); if (ptr == NULL) return NULL; - return unpack_single(self, ptr, fmt); + return unpack_single(ptr, fmt); } PyErr_SetString(PyExc_NotImplementedError, @@ -2615,7 +2818,9 @@ memory_item_multi(PyMemoryViewObject *self, PyObject *tup) ptr = ptr_from_tuple(view, tup); if (ptr == NULL) return NULL; - return unpack_single(self, ptr, fmt); + // ptr_from_tuple() can call __index__(), which can release the view. + CHECK_RELEASED_AGAIN(self); + return unpack_single(ptr, fmt); } static inline int @@ -2703,7 +2908,7 @@ memory_subscript(PyObject *_self, PyObject *key) const char *fmt = adjust_fmt(view); if (fmt == NULL) return NULL; - return unpack_single(self, view->buf, fmt); + return unpack_single(view->buf, fmt); } else if (key == Py_Ellipsis) { return Py_NewRef(self); @@ -2867,6 +3072,16 @@ memory_length(PyObject *_self) return self->view.shape[0]; } +static int +memory_contains(PyObject *_self, PyObject *value) +{ + Py_ssize_t i = _PyBuffer_FindItem(_self, value, 0, PY_SSIZE_T_MAX); + if (i == -2) { + return -1; + } + return i >= 0; +} + /* As mapping */ static PyMappingMethods memory_as_mapping = { memory_length, /* mp_length */ @@ -2880,6 +3095,10 @@ static PySequenceMethods memory_as_sequence = { 0, /* sq_concat */ 0, /* sq_repeat */ memory_item, /* sq_item */ + 0, /* was_sq_slice */ + 0, /* sq_ass_item */ + 0, /* was_sq_ass_slice */ + memory_contains, /* sq_contains */ }; @@ -2900,35 +3119,7 @@ static PyObject * memoryview_count_impl(PyMemoryViewObject *self, PyObject *value) /*[clinic end generated code: output=a15cb19311985063 input=e3036ce1ed7d1823]*/ { - PyObject *iter = PyObject_GetIter(_PyObject_CAST(self)); - if (iter == NULL) { - return NULL; - } - - Py_ssize_t count = 0; - PyObject *item = NULL; - while (PyIter_NextItem(iter, &item)) { - if (item == NULL) { - Py_DECREF(iter); - return NULL; - } - if (item == value) { - Py_DECREF(item); - count++; // no overflow since count <= len(mv) <= PY_SSIZE_T_MAX - continue; - } - int contained = PyObject_RichCompareBool(item, value, Py_EQ); - Py_DECREF(item); - if (contained > 0) { // more likely than 'contained < 0' - count++; // no overflow since count <= len(mv) <= PY_SSIZE_T_MAX - } - else if (contained < 0) { - Py_DECREF(iter); - return NULL; - } - } - Py_DECREF(iter); - return PyLong_FromSsize_t(count); + return _PyBuffer_CountItem(_PyObject_CAST(self), value); } @@ -2954,67 +3145,16 @@ memoryview_index_impl(PyMemoryViewObject *self, PyObject *value, Py_ssize_t start, Py_ssize_t stop) /*[clinic end generated code: output=e0185e3819e549df input=0697a0165bf90b5a]*/ { - const Py_buffer *view = &self->view; - CHECK_RELEASED(self); - - if (view->ndim == 0) { - PyErr_SetString(PyExc_TypeError, "invalid lookup on 0-dim memory"); - return NULL; + Py_ssize_t found = _PyBuffer_FindItem(_PyObject_CAST(self), value, + start, stop); + if (found >= 0) { + return PyLong_FromSsize_t(found); } - - if (view->ndim == 1) { - Py_ssize_t n = view->shape[0]; - - if (start < 0) { - start = Py_MAX(start + n, 0); - } - - if (stop < 0) { - stop = Py_MAX(stop + n, 0); - } - - stop = Py_MIN(stop, n); - assert(stop >= 0); - assert(stop <= n); - - start = Py_MIN(start, stop); - assert(0 <= start); - assert(start <= stop); - - PyObject *obj = _PyObject_CAST(self); - for (Py_ssize_t index = start; index < stop; index++) { - // Note: while memoryviews can be mutated during iterations - // when calling the == operator, their shape cannot. As such, - // it is safe to assume that the index remains valid for the - // entire loop. - assert(index < n); - - PyObject *item = memory_item(obj, index); - if (item == NULL) { - return NULL; - } - if (item == value) { - Py_DECREF(item); - return PyLong_FromSsize_t(index); - } - int contained = PyObject_RichCompareBool(item, value, Py_EQ); - Py_DECREF(item); - if (contained > 0) { // more likely than 'contained < 0' - return PyLong_FromSsize_t(index); - } - else if (contained < 0) { - return NULL; - } - } - - PyErr_SetString(PyExc_ValueError, "memoryview.index(x): x not found"); + if (found == -2) { return NULL; } - - PyErr_SetString(PyExc_NotImplementedError, - "multi-dimensional lookup is not implemented"); + PyErr_SetString(PyExc_ValueError, "memoryview.index(x): x not found"); return NULL; - } @@ -3663,7 +3803,7 @@ memoryiter_next(PyObject *self) if (ptr == NULL) { return NULL; } - return unpack_single(seq, ptr, it->it_fmt); + return unpack_single(ptr, it->it_fmt); } it->it_seq = NULL;