use crate::{

common::{

atomic::{self, AtomicUsize},

lock::{

PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,

PyRwLockReadGuard, PyRwLockWriteGuard,

},

str::wchar_t,

},

vm::{

AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,

atomic_func,

builtins::{

PositionIterInternal, PyByteArray, PyBytes, PyBytesRef, PyDictRef, PyFloat,

PyGenericAlias, PyInt, PyList, PyListRef, PyStr, PyStrRef, PyTupleRef, PyType,

PyTypeRef, PyUtf8StrRef, builtins_iter,

},

class_or_notimplemented,

convert::{ToPyObject, ToPyResult, TryFromBorrowedObject, TryFromObject},

function::{

ArgBytesLike, ArgIntoFloat, ArgIterable, KwArgs, OptionalArg, PyComparisonValue,

},

protocol::{

BufferDescriptor, BufferFlags, BufferMethods, BufferResizeGuard, PyBuffer,

PyIterReturn, PyMappingMethods, PySequenceMethods,

},

sequence::{OptionalRangeArgs, SequenceExt, SequenceMutExt},

sliceable::{

SaturatedSlice, SequenceIndex, SequenceIndexOp, SliceableSequenceMutOp,

SliceableSequenceOp,

},

stdlib::_warnings,

types::{

AsBuffer, AsMapping, AsSequence, Comparable, Constructor, IterNext, Iterable,

PyComparisonOp, Representable, SelfIter,

},

},

};

use alloc::fmt;

use core::cmp::Ordering;

use itertools::Itertools;

use num_traits::ToPrimitive;

use rustpython_common::wtf8::{CodePoint, Wtf8, Wtf8Buf};

use std::os::raw;

macro_rules! def_array_enum {

($(($n:ident, $t:ty, $c:literal, $scode:literal)),*$(,)?) => {

#[derive(Debug, Clone)]

pub enum ArrayContentType {

$($n(Vec<$t>),)*

}

/// One item, already converted to the array's element type.

enum ArrayItem {

$($n($t),)*

}

impl ArrayContentType {

fn from_char(c: char) -> Result<Self, String> {

match c {

$($c => Ok(ArrayContentType::$n(Vec::new())),)*

_ => Err(

"bad typecode (must be b, B, u, w, h, H, i, I, l, L, q, Q, f or d)".into()

),

}

}

const fn typecode(&self) -> char {

match self {

$(ArrayContentType::$n(_) => $c,)*

}

}

const fn typecode_str(&self) -> &'static str {

match self {

$(ArrayContentType::$n(_) => $scode,)*

}

}

const fn itemsize_of_typecode(c: char) -> Option<usize> {

match c {

$($c => Some(core::mem::size_of::<$t>()),)*

_ => None,

}

}

const fn itemsize(&self) -> usize {

match self {

$(ArrayContentType::$n(_) => core::mem::size_of::<$t>(),)*

}

}

fn addr(&self) -> usize {

match self {

$(ArrayContentType::$n(v) => v.as_ptr() as usize,)*

}

}

fn len(&self) -> usize {

match self {

$(ArrayContentType::$n(v) => v.len(),)*

}

}

fn reserve(&mut self, len: usize) {

match self {

$(ArrayContentType::$n(v) => v.reserve(len),)*

}

}

fn push(&mut self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => {

let val = <$t>::try_into_from_object(vm, obj)?;

v.push(val);

})*

}

Ok(())

}

fn pop(&mut self, i: isize, vm: &VirtualMachine) -> PyResult {

match self {

$(ArrayContentType::$n(v) => {

let i = v.wrap_index(i).ok_or_else(|| {

vm.new_index_error("pop index out of range".to_owned())

})?;

v.remove(i).to_pyresult(vm)

})*

}

}

fn insert(

&mut self,

i: isize,

obj: PyObjectRef,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => {

let val = <$t>::try_into_from_object(vm, obj)?;

v.insert(i.saturated_at(v.len()), val);

})*

}

Ok(())

}

fn count(&self, obj: PyObjectRef, vm: &VirtualMachine) -> usize {

match self {

$(ArrayContentType::$n(v) => {

if let Ok(val) = <$t>::try_into_from_object(vm, obj) {

v.iter().filter(|&&a| a == val).count()

} else {

0

}

})*

}

}

fn clear(&mut self) -> PyResult<()>{

match self {

$(ArrayContentType::$n(v) => v.clear(),)*

};

Ok(())

}

fn remove(&mut self, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()>{

match self {

$(ArrayContentType::$n(v) => {

if let Ok(val) = <$t>::try_into_from_object(vm, obj) {

if let Some(pos) = v.iter().position(|&a| a == val) {

v.remove(pos);

return Ok(());

}

}

Err(vm.new_value_error("array.remove(x): x not in array".to_owned()))

})*

}

}

fn frombytes_move(&mut self, b: Vec<u8>) {

match self {

$(ArrayContentType::$n(v) => {

if v.is_empty() {

// safe because every configuration of bytes for the types we

// support are valid

let b = core::mem::ManuallyDrop::new(b);

let ptr = b.as_ptr() as *mut $t;

let len = b.len() / core::mem::size_of::<$t>();

let capacity = b.capacity() / core::mem::size_of::<$t>();

*v = unsafe { Vec::from_raw_parts(ptr, len, capacity) };

} else {

self.frombytes(&b);

}

})*

}

}

fn frombytes(&mut self, b: &[u8]) {

match self {

$(ArrayContentType::$n(v) => {

// safe because every configuration of bytes for the types we

// support are valid

if b.len() > 0 {

let ptr = b.as_ptr() as *const $t;

let ptr_len = b.len() / core::mem::size_of::<$t>();

let slice = unsafe { core::slice::from_raw_parts(ptr, ptr_len) };

v.extend_from_slice(slice);

}

})*

}

}

fn fromlist(&mut self, list: &PyList, vm: &VirtualMachine) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => {

// convert list before modify self

let mut list: Vec<$t> = list

.borrow_vec()

.iter()

.cloned()

.map(|value| <$t>::try_into_from_object(vm, value))

.try_collect()?;

v.append(&mut list);

Ok(())

})*

}

}

fn get_bytes(&self) -> &[u8] {

match self {

$(ArrayContentType::$n(v) => {

// safe because we're just reading memory as bytes

let ptr = v.as_ptr() as *const u8;

let ptr_len = v.len() * core::mem::size_of::<$t>();

unsafe { core::slice::from_raw_parts(ptr, ptr_len) }

})*

}

}

fn get_bytes_mut(&mut self) -> &mut [u8] {

match self {

$(ArrayContentType::$n(v) => {

// safe because we're just reading memory as bytes

let ptr = v.as_ptr() as *mut u8;

let ptr_len = v.len() * core::mem::size_of::<$t>();

unsafe { core::slice::from_raw_parts_mut(ptr, ptr_len) }

})*

}

}

fn index(

&self,

obj: PyObjectRef,

start: usize,

stop: usize,

vm: &VirtualMachine

) -> PyResult<usize> {

match self {

$(ArrayContentType::$n(v) => {

if let Ok(val) = <$t>::try_into_from_object(vm, obj) {

if let Some(pos) = v.iter().take(stop as _).skip(start as _).position(|&elem| elem == val) {

return Ok(pos + start);

}

}

Err(vm.new_value_error("array.index(x): x not in array".to_owned()))

})*

}

}

fn reverse(&mut self) {

match self {

$(ArrayContentType::$n(v) => v.reverse(),)*

}

}

fn get(

&self,

i: usize,

vm: &VirtualMachine

) -> Option<PyResult> {

match self {

$(ArrayContentType::$n(v) => {

v.get(i).map(|x| x.to_pyresult(vm))

})*

}

}

fn getitem_by_index(&self, i: isize, vm: &VirtualMachine) -> PyResult {

match self {

$(ArrayContentType::$n(v) => {

v.getitem_by_index(vm, i).map(|x| x.to_pyresult(vm))?

})*

}

}

fn getitem_by_slice(&self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult {

match self {

$(ArrayContentType::$n(v) => {

let r = v.getitem_by_slice(vm, slice)?;

let array = PyArray::from(ArrayContentType::$n(r));

array.to_pyresult(vm)

})*

}

}

/// Convert an object to the element type of the array with

/// this typecode. This runs the object's conversion methods,

/// which can reach the array, so it takes the typecode by

/// value and holds no lock on it.

fn item_from_object(

typecode: char,

value: PyObjectRef,

vm: &VirtualMachine

) -> PyResult<ArrayItem> {

match typecode {

$($c => Ok(ArrayItem::$n(<$t>::try_into_from_object(vm, value)?)),)*

_ => unreachable!("array has a typecode"),

}

}

fn setitem_by_item(

&mut self,

i: isize,

item: ArrayItem,

vm: &VirtualMachine

) -> PyResult<()> {

match (self, item) {

$((ArrayContentType::$n(v), ArrayItem::$n(value)) =>

v.setitem_by_index(vm, i, value),)*

_ => unreachable!("item was converted for this array"),

}

}

fn setitem_by_slice(

&mut self,

slice: SaturatedSlice,

items: &ArrayContentType,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

$(Self::$n(elements) => if let ArrayContentType::$n(items) = items {

elements.setitem_by_slice(vm, slice, items)

} else {

Err(vm.new_type_error(

"bad argument type for built-in operation".to_owned()

))

},)*

}

}

fn setitem_by_slice_no_resize(

&mut self,

slice: SaturatedSlice,

items: &ArrayContentType,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

$(Self::$n(elements) => if let ArrayContentType::$n(items) = items {

elements.setitem_by_slice_no_resize(vm, slice, items)

} else {

Err(vm.new_type_error(

"bad argument type for built-in operation".to_owned()

))

},)*

}

}

fn delitem_by_index(&mut self, i: isize, vm: &VirtualMachine) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => {

v.delitem_by_index(vm, i)

})*

}

}

fn delitem_by_slice(&mut self, slice: SaturatedSlice, vm: &VirtualMachine) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => {

v.delitem_by_slice(vm, slice)

})*

}

}

fn add(&self, other: &ArrayContentType, vm: &VirtualMachine) -> PyResult<Self> {

match self {

$(ArrayContentType::$n(v) => if let ArrayContentType::$n(other) = other {

let elements = v.iter().chain(other.iter()).cloned().collect();

Ok(ArrayContentType::$n(elements))

} else {

Err(vm.new_type_error(

"bad argument type for built-in operation".to_owned()

))

},)*

}

}

fn iadd(&mut self, other: &ArrayContentType, vm: &VirtualMachine) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => if let ArrayContentType::$n(other) = other {

v.extend(other);

Ok(())

} else {

Err(vm.new_type_error(

"can only extend with array of same kind".to_owned()

))

},)*

}

}

fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {

match self {

$(ArrayContentType::$n(v) => {

// MemoryError instead Overflow Error, hard to says it is right

// but it is how cpython doing right now

let elements = v.mul(vm, value).map_err(|_| vm.new_memory_error("".to_owned()))?;

Ok(ArrayContentType::$n(elements))

})*

}

}

fn imul(&mut self, value: isize, vm: &VirtualMachine) -> PyResult<()> {

match self {

$(ArrayContentType::$n(v) => {

// MemoryError instead Overflow Error, hard to says it is right

// but it is how cpython doing right now

v.imul(vm, value).map_err(|_| vm.new_memory_error("".to_owned()))

})*

}

}

fn byteswap(&mut self) {

match self {

$(ArrayContentType::$n(v) => {

for element in v.iter_mut() {

let x = element.byteswap();

*element = x;

}

})*

}

}

fn repr(&self, class_name: &str, _vm: &VirtualMachine) -> PyResult<String> {

// we don't need ReprGuard here

let s = match self {

$(ArrayContentType::$n(v) => {

if v.is_empty() {

format!("{}('{}')", class_name, $c)

} else {

format!("{}('{}', [{}])", class_name, $c, v.iter().format(", "))

}

})*

};

Ok(s)

}

fn iter<'a, 'vm: 'a>(

&'a self,

vm: &'vm VirtualMachine

) -> impl Iterator<Item = PyResult> + 'a {

(0..self.len()).map(move |i| self.get(i, vm).unwrap())

}

fn cmp(&self, other: &ArrayContentType) -> Result<Option<Ordering>, ()> {

match self {

$(ArrayContentType::$n(v) => {

if let ArrayContentType::$n(other) = other {

Ok(PartialOrd::partial_cmp(v, other))

} else {

Err(())

}

})*

}

}

fn get_objects(&self, vm: &VirtualMachine) -> Vec<PyObjectRef> {

match self {

$(ArrayContentType::$n(v) => {

v.iter().map(|&x| x.to_object(vm)).collect()

})*

}

}

}

};

}

def_array_enum!(

(SignedByte, i8, 'b', "b"),

(UnsignedByte, u8, 'B', "B"),

(PyUnicode, WideChar, 'u', "u"),

(PyUcs4, Ucs4Char, 'w', "w"),

(SignedShort, raw::c_short, 'h', "h"),

(UnsignedShort, raw::c_ushort, 'H', "H"),

(SignedInt, raw::c_int, 'i', "i"),

(UnsignedInt, raw::c_uint, 'I', "I"),

(SignedLong, raw::c_long, 'l', "l"),

(UnsignedLong, raw::c_ulong, 'L', "L"),

(SignedLongLong, raw::c_longlong, 'q', "q"),

(UnsignedLongLong, raw::c_ulonglong, 'Q', "Q"),

(Float, f32, 'f', "f"),

(Double, f64, 'd', "d"),

);

#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]

pub struct WideChar(wchar_t);

/// Element type for the 'w' typecode: always a 4-byte unicode code point

/// (Py_UCS4), unlike 'u' which is platform-dependent `wchar_t`.

#[derive(Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Debug)]

pub struct Ucs4Char(u32);

trait ArrayElement: Sized {

fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self>;

fn byteswap(self) -> Self;

fn to_object(self, vm: &VirtualMachine) -> PyObjectRef;

}

macro_rules! impl_int_element {

($($t:ty,)*) => {$(

impl ArrayElement for $t {

fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

obj.try_index(vm)?.try_to_primitive_raw(vm)

}

fn byteswap(self) -> Self {

<$t>::swap_bytes(self)

}

fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {

self.to_pyobject(vm)

}

}

)*};

}

macro_rules! impl_float_element {

($(($t:ty, $f_from:path, $f_swap:path, $f_to:path),)*) => {$(

impl ArrayElement for $t {

fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

$f_from(vm, obj)

}

fn byteswap(self) -> Self {

$f_swap(self)

}

fn to_object(self, vm: &VirtualMachine) -> PyObjectRef {

$f_to(self).into_pyobject(vm)

}

}

)*};

}

impl_int_element!(i8, u8, i16, u16, i32, u32, i64, u64,);

impl_float_element!(

(

f32,

f32_try_into_from_object,

f32_swap_bytes,

pyfloat_from_f32

),

(f64, f64_try_into_from_object, f64_swap_bytes, PyFloat::from),

);

const fn f32_swap_bytes(x: f32) -> f32 {

f32::from_bits(x.to_bits().swap_bytes())

}

const fn f64_swap_bytes(x: f64) -> f64 {

f64::from_bits(x.to_bits().swap_bytes())

}

fn f32_try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<f32> {

ArgIntoFloat::try_from_object(vm, obj).map(|x| x.into_float() as f32)

}

fn f64_try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<f64> {

ArgIntoFloat::try_from_object(vm, obj).map(|x| x.into_float())

}

fn pyfloat_from_f32(value: f32) -> PyFloat {

PyFloat::from(value as f64)

}

impl ArrayElement for WideChar {

fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

PyUtf8StrRef::try_from_object(vm, obj)?

.as_str()

.chars()

.exactly_one()

.map(|ch| Self(ch as _))

.map_err(|_| vm.new_type_error("array item must be unicode character"))

}

fn byteswap(self) -> Self {

Self(self.0.swap_bytes())

}

fn to_object(self, _vm: &VirtualMachine) -> PyObjectRef {

unreachable!()

}

}

impl ArrayElement for Ucs4Char {

fn try_into_from_object(vm: &VirtualMachine, obj: PyObjectRef) -> PyResult<Self> {

let s = obj.downcast::<PyStr>().map_err(|obj| {

vm.new_type_error(format!(

"array item must be a unicode character, not {}",

obj.class().name()

))

})?;

s.as_wtf8()

.code_points()

.exactly_one()

.map(|ch| Self(ch.to_u32()))

.map_err(|e| {

vm.new_type_error(format!(

"array item must be a unicode character, not a string of length {}",

e.count()

))

})

}

fn byteswap(self) -> Self {

Self(self.0.swap_bytes())

}

fn to_object(self, _vm: &VirtualMachine) -> PyObjectRef {

unreachable!()

}

}

impl ToPyResult for Ucs4Char {

fn to_pyresult(self, vm: &VirtualMachine) -> PyResult {

Ok(u32_to_char(self.0)

.map_err(|msg| vm.new_value_error(msg))?

.to_pyobject(vm))

}

}

impl fmt::Display for Ucs4Char {

fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {

unreachable!("`repr(array('w'))` calls `PyStr::repr`")

}

}

fn u32_to_char(ch: u32) -> Result<CodePoint, String> {

CodePoint::from_u32(ch)

.ok_or_else(|| format!("character U+{ch:4x} is not in range [U+0000; U+10ffff]"))

}

impl TryFrom<WideChar> for CodePoint {

type Error = String;

fn try_from(ch: WideChar) -> Result<Self, Self::Error> {

// safe because every configuration of bytes for the types we support are valid

u32_to_char(ch.0 as _)

}

}

impl ToPyResult for WideChar {

fn to_pyresult(self, vm: &VirtualMachine) -> PyResult {

Ok(CodePoint::try_from(self)

.map_err(|e| vm.new_value_error(e))?

.to_pyobject(vm))

}

}

impl fmt::Display for WideChar {

fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {

unreachable!("`repr(array('u'))` calls `PyStr::repr`")

}

}

#[pyattr]

#[pyattr(name = "ArrayType")]

#[pyclass(name = "array")]

#[derive(Debug, PyPayload)]

pub struct PyArray {

array: PyRwLock<ArrayContentType>,

exports: AtomicUsize,

}

pub type PyArrayRef = PyRef<PyArray>;

impl From<ArrayContentType> for PyArray {

fn from(array: ArrayContentType) -> Self {

Self {

array: PyRwLock::new(array),

exports: AtomicUsize::new(0),

}

}

}

#[derive(FromArgs)]

pub struct ArrayNewArgs {

#[pyarg(positional)]

spec: PyUtf8StrRef,

#[pyarg(positional, optional)]

init: OptionalArg<PyObjectRef>,

}

impl Constructor for PyArray {

type Args = (ArrayNewArgs, KwArgs);

fn py_new(

cls: &Py<PyType>,

(ArrayNewArgs { spec, init }, kwargs): Self::Args,

vm: &VirtualMachine,

) -> PyResult<Self> {

let spec = spec.as_str().chars().exactly_one().map_err(|_| {

vm.new_type_error("array() argument 1 must be a unicode character, not str")

})?;

if cls.is(Self::class(&vm.ctx)) && !kwargs.is_empty() {

return Err(vm.new_type_error("array.array() takes no keyword arguments"));

}

if spec == 'u' {

_warnings::warn(

vm.ctx.exceptions.deprecation_warning,

"The 'u' type code is deprecated and will be removed in Python 3.16".to_owned(),

1,

vm,

)?;

}

let mut array =

ArrayContentType::from_char(spec).map_err(|err| vm.new_value_error(err))?;

if let OptionalArg::Present(init) = init {

if let Some(init) = init.downcast_ref::<Self>() {

let value = init.read().typecode();

match (spec, value) {

(spec, ch) if spec == ch => array.frombytes(&init.get_bytes()),

(spec, 'u' | 'w') if !matches!(spec, 'u' | 'w') => {

return Err(vm.new_type_error(format!(

"cannot use a unicode array to initialize an array with typecode '{spec}'"

)))

}

_ => {

for obj in init.read().iter(vm) {

array.push(obj?, vm)?;

}

}

}

} else if let Some(wtf8) = init.downcast_ref::<PyStr>() {

if matches!(spec, 'u' | 'w') {

let bytes = Self::_unicode_to_wchar_bytes(wtf8.as_wtf8(), array.itemsize());

array.frombytes_move(bytes);

} else {

return Err(vm.new_type_error(format!(

"cannot use a str to initialize an array with typecode '{spec}'"

)));

}

} else if init.downcastable::<PyBytes>() || init.downcastable::<PyByteArray>() {

init.try_bytes_like(vm, |x| array.frombytes(x))?;

} else {

// Everything else is taken item by item, buffer or not.

let iter = ArgIterable::try_from_object(vm, init)?;

for obj in iter.iter(vm)? {

array.push(obj?, vm)?;

}

}

}

Ok(Self::from(array))

}

}

#[pyclass(

impl PyArray {

fn read(&self) -> PyRwLockReadGuard<'_, ArrayContentType> {

self.array.read()

}

fn write(&self) -> PyRwLockWriteGuard<'_, ArrayContentType> {

self.array.write()

}

#[pygetset]

fn typecode(&self, vm: &VirtualMachine) -> PyStrRef {

vm.ctx

.intern_str(self.read().typecode().to_string())

.to_owned()

}

#[pygetset]

fn itemsize(&self) -> usize {

self.read().itemsize()

}

#[pymethod]

fn append(zelf: &Py<Self>, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

zelf.try_resizable(vm)?.push(x, vm)

}

#[pymethod]

fn clear(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {

zelf.try_resizable(vm)?.clear()

}

#[pymethod]

fn buffer_info(&self) -> (usize, usize) {

let array = self.read();

(array.addr(), array.len())

}

#[pymethod]

fn count(&self, x: PyObjectRef, vm: &VirtualMachine) -> usize {

self.read().count(x, vm)

}

#[pymethod]

fn remove(zelf: &Py<Self>, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

zelf.try_resizable(vm)?.remove(x, vm)

}

#[pymethod]

fn extend(zelf: &Py<Self>, obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

let mut w = zelf.try_resizable(vm)?;

if zelf.is(&obj) {

w.imul(2, vm)

} else if let Some(array) = obj.downcast_ref::<Self>() {

w.iadd(&array.read(), vm)

} else {

let iter = ArgIterable::try_from_object(vm, obj)?;

// zelf.extend_from_iterable(iter, vm)

for obj in iter.iter(vm)? {

w.push(obj?, vm)?;

}

Ok(())

}

}

fn _wchar_bytes_to_string(

bytes: &[u8],

item_size: usize,

vm: &VirtualMachine,