use crate::common::{

borrow::{BorrowedValue, BorrowedValueMut},

lock::{

PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLock, PyRwLockReadGuard,

PyRwLockWriteGuard,

},

str::wchar_t,

};

use crate::{

builtins::{

IntoPyFloat, PyByteArray, PyBytes, PyBytesRef, PyIntRef, PyList, PyListRef, PySliceRef,

PyStr, PyStrRef, PyTypeRef,

},

byteslike::ArgBytesLike,

function::{ArgIterable, OptionalArg},

protocol::{BufferInternal, BufferOptions, PyBuffer, ResizeGuard},

sliceable::{saturate_index, PySliceableSequence, PySliceableSequenceMut, SequenceIndex},

slots::{

AsBuffer, Comparable, Iterable, IteratorIterable, PyComparisonOp, PyIter,

SlotConstructor,

},

IdProtocol, IntoPyObject, IntoPyResult, PyComparisonValue, PyObjectRef, PyRef, PyResult,

PyValue, TryFromObject, TypeProtocol, VirtualMachine,

};

use crossbeam_utils::atomic::AtomicCell;

use itertools::Itertools;

use lexical_core::Integer;

use num_traits::ToPrimitive;

use std::cmp::Ordering;

use std::convert::{TryFrom, TryInto};

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

macro_rules! def_array_enum {

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

#[derive(Debug, Clone)]

pub(crate) enum ArrayContentType {

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

}

#[allow(clippy::naive_bytecount, clippy::float_cmp)]

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).into_pyresult(vm)

})*

}

}

fn insert(

&mut self,

i: usize,

obj: PyObjectRef,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

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

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

v.insert(i, 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, 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().position(|&a| a == val) {

return Ok(pos);

}

}

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 getitem_by_idx(

&self,

i: usize,

vm: &VirtualMachine

) -> PyResult<Option<PyObjectRef>> {

match self {

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

v.get(i).map(|x| x.into_pyresult(vm)).transpose()

})*

}

}

fn getitem(

&self,

needle: PyObjectRef,

vm: &VirtualMachine

) -> PyResult {

match self {

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

match SequenceIndex::try_from_object_for(vm, needle, "array")? {

SequenceIndex::Int(i) => {

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

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

})?;

v.get(pos_index).unwrap().into_pyresult(vm)

}

SequenceIndex::Slice(slice) => {

let elements = v.get_slice_items(vm, &slice)?;

let array: PyArray = ArrayContentType::$n(elements).into();

Ok(array.into_object(vm))

}

}

})*

}

}

fn setitem_by_slice(

&mut self,

slice: PySliceRef,

items: &ArrayContentType,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

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

elements.set_slice_items(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: PySliceRef,

items: &ArrayContentType,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

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

elements.set_slice_items_no_resize(vm, &slice, items)

} else {

Err(vm.new_type_error(

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

))

},)*

}

}

fn setitem_by_idx(

&mut self,

i: isize,

value: PyObjectRef,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

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

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

vm.new_index_error("array assignment index out of range".to_owned())

})?;

v[i] = <$t>::try_into_from_object(vm, value)?

})*

}

Ok(())

}

fn delitem_by_idx(&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("array assignment index out of range".to_owned())

})?;

v.remove(i);

})*

}

Ok(())

}

fn delitem_by_slice(

&mut self,

slice: PySliceRef,

vm: &VirtualMachine

) -> PyResult<()> {

match self {

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

elements.delete_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, counter: usize) -> Self {

match self {

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

let elements = v.repeat(counter);

ArrayContentType::$n(elements)

})*

}

}

fn clear(&mut self) {

match self {

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

}

}

fn imul(&mut self, counter: usize) {

if counter == 0 {

self.clear();

} else if counter != 1 {

match self {

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

let old = v.clone();

v.reserve((counter - 1) * old.len());

for _ in 1..counter {

v.extend(&old);

}

})*

}

}

}

fn byteswap(&mut self) {

match self {

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

for element in v.iter_mut() {

let x = element.byteswap();

*element = x;

}

})*

}

}

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

// we don't need ReprGuard here

let s = match self {

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

if v.is_empty() {

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

} else {

format!("array('{}', [{}])", $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.getitem_by_idx(i, vm).map(Option::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(())

}

})*

}

}

}

};

}

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;

}

macro_rules! impl_array_element {

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

impl ArrayElement for $t {

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

$f_into(vm, obj)

}

fn byteswap(self) -> Self {

$f_swap(self)

}

}

)*};

}

impl_array_element!(

(i8, i8::try_from_object, i8::swap_bytes),

(u8, u8::try_from_object, u8::swap_bytes),

(i16, i16::try_from_object, i16::swap_bytes),

(u16, u16::try_from_object, u16::swap_bytes),

(i32, i32::try_from_object, i32::swap_bytes),

(u32, u32::try_from_object, u32::swap_bytes),

(i64, i64::try_from_object, i64::swap_bytes),

(u64, u64::try_from_object, u64::swap_bytes),

(f32, f32_try_into_from_object, f32_swap_bytes),

(f64, f64_try_into_from_object, f64_swap_bytes),

);

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

IntoPyFloat::try_from_object(vm, obj).map(|x| x.to_f64() as f32)

}

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

IntoPyFloat::try_from_object(vm, obj).map(|x| x.to_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 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 IntoPyResult for WideChar {

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

Ok(

String::from(char::try_from(self).map_err(|e| vm.new_unicode_encode_error(e))?)

.into_pyobject(vm),

)

}

}

impl fmt::Display for WideChar {

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

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

}

}

#[pyattr]

#[pyclass(name = "array")]

#[derive(Debug, PyValue)]

pub struct PyArray {

array: PyRwLock<ArrayContentType>,

exports: AtomicCell<usize>,

}

pub type PyArrayRef = PyRef<PyArray>;

impl From<ArrayContentType> for PyArray {

fn from(array: ArrayContentType) -> Self {

PyArray {

array: PyRwLock::new(array),

exports: AtomicCell::new(0),

}

}

}

#[derive(FromArgs)]

pub struct ArrayNewArgs {

#[pyarg(positional)]

spec: PyStrRef,

#[pyarg(positional, optional)]

init: OptionalArg<PyObjectRef>,

}

impl SlotConstructor 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_object())

}

}

#[pyimpl(flags(BASETYPE), with(Comparable, AsBuffer, Iterable, SlotConstructor))]

impl PyArray {

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

self.array.read()

}

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

self.array.write()

}

#[pyproperty]

fn typecode(&self) -> String {

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

}

#[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);

Ok(())

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

if self.itemsize() == 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 _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)

}

#[pymethod]

fn fromfile(&self, f: PyObjectRef, n: isize, vm: &VirtualMachine) -> PyResult<()> {

let itemsize = self.itemsize();

if n < 0 {

return Err(vm.new_value_error("negative count".to_owned()));

}

let n = vm.check_repeat_or_memory_error(itemsize, n)?;

let nbytes = n * itemsize;

let b = vm.call_method(&f, "read", (nbytes,))?;

let b = b

.downcast::<PyBytes>()

.map_err(|_| vm.new_type_error("read() didn't return bytes".to_owned()))?;

let not_enough_bytes = b.len() != nbytes;

self._from_bytes(b.as_bytes(), itemsize, vm)?;

if not_enough_bytes {

Err(vm.new_exception_msg(

vm.ctx.exceptions.eof_error.clone(),

"read() didn't return enough bytes".to_owned(),