use crate::{

common::{

atomic::{self, AtomicUsize},

lock::{

PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock,

PyRwLockReadGuard, PyRwLockWriteGuard,

},

str::wchar_t,

},

vm::{

builtins::{

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

PyIntRef, PyList, PyListRef, PyStr, PyStrRef, PyTupleRef, PyTypeRef,

},

class_or_notimplemented,

convert::{ToPyObject, ToPyResult, TryFromObject},

function::{ArgBytesLike, ArgIntoFloat, ArgIterable, OptionalArg, PyComparisonValue},

protocol::{

BufferDescriptor, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,

PyMappingMethods,

},

sequence::{OptionalRangeArgs, SequenceExt, SequenceMutExt},

sliceable::{

SaturatedSlice, SequenceIndex, SequenceIndexOp, SliceableSequenceMutOp,

SliceableSequenceOp,

},

types::{

AsBuffer, AsMapping, Comparable, Constructor, IterNext, IterNextIterable, Iterable,

PyComparisonOp,

},

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

},

};

use itertools::Itertools;

use num_traits::ToPrimitive;

use std::{cmp::Ordering, fmt, os::raw};

macro_rules! def_array_enum {

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

#[derive(Debug, Clone)]

pub enum ArrayContentType {

$($n(Vec<$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, h, H, i, I, l, L, q, Q, f or d)".into()

),

}

}

fn typecode(&self) -> char {

match self {

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

}

}

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

match self {

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

}

}

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

match c {

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

_ => None,

}

}

fn itemsize(&self) -> usize {

match self {

$(ArrayContentType::$n(_) => std::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 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 = std::mem::ManuallyDrop::new(b);

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

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

let capacity = b.capacity() / std::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

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

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

let slice = unsafe { std::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() * std::mem::size_of::<$t>();

unsafe { std::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() * std::mem::size_of::<$t>();

unsafe { std::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)

})*

}

}

fn setitem_by_index(

&mut self,

i: isize,

value: PyObjectRef,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

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

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

v.setitem_by_index(vm, i, value)

})*

}

}

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.del_item_by_index(vm, i)

})*

}

}

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

match self {

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

v.del_item_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"),

(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);

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_array_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_array_element!(

(i8, i8::try_from_object, i8::swap_bytes, PyInt::from),

(u8, u8::try_from_object, u8::swap_bytes, PyInt::from),

(i16, i16::try_from_object, i16::swap_bytes, PyInt::from),

(u16, u16::try_from_object, u16::swap_bytes, PyInt::from),

(i32, i32::try_from_object, i32::swap_bytes, PyInt::from),

(u32, u32::try_from_object, u32::swap_bytes, PyInt::from),

(i64, i64::try_from_object, i64::swap_bytes, PyInt::from),

(u64, u64::try_from_object, u64::swap_bytes, PyInt::from),

(

f32,

f32_try_into_from_object,

f32_swap_bytes,

pyfloat_from_f32

),

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

);

fn f32_swap_bytes(x: f32) -> f32 {

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

}

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 as f32)

}

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

ArgIntoFloat::try_from_object(vm, obj).map(Into::into)

}

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> {

PyStrRef::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".into()))

}

fn byteswap(self) -> Self {

Self(self.0.swap_bytes())

}

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

unreachable!()

}

}

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

if ch > 0x10ffff {

return Err(format!(

"character U+{:4x} is not in range [U+0000; U+10ffff]",

ch

));

};

char::from_u32(ch).ok_or_else(|| {

format!(

"'utf-8' codec can't encode character '\\u{:x}' \

ch

)

})

}

impl TryFrom<WideChar> for char {

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 u32)

}

}

impl ToPyResult for WideChar {

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

Ok(

String::from(char::try_from(self).map_err(|e| vm.new_unicode_encode_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 {

PyArray {

array: PyRwLock::new(array),

exports: AtomicUsize::new(0),

}

}

}

#[derive(FromArgs)]

pub struct ArrayNewArgs {

#[pyarg(positional)]

spec: PyStrRef,

#[pyarg(positional, optional)]

init: OptionalArg<PyObjectRef>,

}

impl Constructor for PyArray {

type Args = ArrayNewArgs;

fn py_new(

cls: PyTypeRef,

Self::Args { spec, init }: Self::Args,

vm: &VirtualMachine,

) -> PyResult {

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

vm.new_type_error(

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

)

})?;

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.payload::<PyArray>() {

match (spec, init.read().typecode()) {

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

(spec, 'u') => {

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(utf8) = init.payload::<PyStr>() {

if spec == 'u' {

let bytes = Self::_unicode_to_wchar_bytes(utf8.as_str(), 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.payload_is::<PyBytes>() || init.payload_is::<PyByteArray>() {

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

} else if let Ok(iter) = ArgIterable::try_from_object(vm, init.clone()) {

for obj in iter.iter(vm)? {

array.push(obj?, vm)?;

}

} else {

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

}

}

let zelf = Self::from(array).into_ref_with_type(vm, cls)?;

Ok(zelf.into())

}

}

#[pyimpl(

impl PyArray {

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

self.array.read()

}

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

self.array.write()

}

#[pyproperty]

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

vm.ctx

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

.to_owned()

}

#[pyproperty]

fn itemsize(&self) -> usize {

self.read().itemsize()

}

#[pymethod]

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

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

}

#[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: PyRef<Self>, x: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {

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

}

#[pymethod]

fn extend(zelf: PyRef<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.payload::<PyArray>() {

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,

) -> PyResult<String> {

if item_size == 2 {

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

let utf16 = unsafe {

std::slice::from_raw_parts(

bytes.as_ptr() as *const u16,

bytes.len() / std::mem::size_of::<u16>(),

)

};

Ok(String::from_utf16_lossy(utf16))

} else {

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

let chars = unsafe {

std::slice::from_raw_parts(

bytes.as_ptr() as *const u32,

bytes.len() / std::mem::size_of::<u32>(),

)

};

chars

.iter()

.map(|&ch| {

// cpython issue 17223

u32_to_char(ch).map_err(|msg| vm.new_value_error(msg))

})

.try_collect()

}

}

fn _unicode_to_wchar_bytes(utf8: &str, item_size: usize) -> Vec<u8> {

if item_size == 2 {

utf8.encode_utf16()

.flat_map(|ch| ch.to_ne_bytes())

.collect()

} else {

utf8.chars()

.flat_map(|ch| (ch as u32).to_ne_bytes())

.collect()

}

}

#[pymethod]

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

let utf8 = PyStrRef::try_from_object(vm, obj.clone()).map_err(|_| {

vm.new_type_error(format!(

"fromunicode() argument must be str, not {}",

obj.class().name()

))

})?;

if zelf.read().typecode() != 'u' {

return Err(vm.new_value_error(

"fromunicode() may only be called on unicode type arrays".into(),

));

}

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

let bytes = Self::_unicode_to_wchar_bytes(utf8.as_str(), w.itemsize());

w.frombytes_move(bytes);

Ok(())

}

#[pymethod]

fn tounicode(&self, vm: &VirtualMachine) -> PyResult<String> {

let array = self.array.read();

if array.typecode() != 'u' {

return Err(vm.new_value_error(

"tounicode() may only be called on unicode type arrays".into(),

));

}

let bytes = array.get_bytes();

Self::_wchar_bytes_to_string(bytes, self.itemsize(), vm)

}

fn _from_bytes(&self, b: &[u8], itemsize: usize, vm: &VirtualMachine) -> PyResult<()> {

if b.len() % itemsize != 0 {

return Err(

vm.new_value_error("bytes length not a multiple of item size".to_owned())

);

}

if b.len() / itemsize > 0 {

self.try_resizable(vm)?.frombytes(b);

}

Ok(())

}

#[pymethod]

fn frombytes(&self, b: ArgBytesLike, vm: &VirtualMachine) -> PyResult<()> {

let b = b.borrow_buf();

let itemsize = self.read().itemsize();

self._from_bytes(&b, itemsize, vm)

}