use crate::{
AsObject, Py, PyObject, PyObjectRef, PyResult, TryFromObject, VirtualMachine,
builtins::{PyBaseExceptionRef, PyBytesRef, PyComplex, PyTuple, PyTupleRef, PyType, PyTypeRef},
common::{lock::PyRwLock, rc::PyRc, static_cell, str::wchar_t},
convert::ToPyObject,
exceptions,
function::{ArgBytesLike, ArgIntoBool, ArgIntoComplex, ArgIntoFloat},
};
use rustpython_common::wtf8::Wtf8Buf;
use core::{fmt, iter::Peekable, mem};
use half::f16;
use itertools::Itertools;
use malachite_bigint::BigInt;
use num_complex::Complex64;
use num_traits::{PrimInt, ToPrimitive};
use std::{collections::HashMap, os::raw};
type PackFunc = fn(&VirtualMachine, FormatType, PyObjectRef, &mut [u8]) -> Result<(), PackError>;
type UnpackFunc = fn(&VirtualMachine, &[u8]) -> PyObjectRef;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PackErrorKind {
Type,
Value,
Raised,
}
pub struct PackError {
pub kind: PackErrorKind,
pub exception: PyBaseExceptionRef,
}
impl PackError {
fn new<T: Into<Wtf8Buf>>(kind: PackErrorKind, vm: &VirtualMachine, msg: T) -> Self {
Self {
kind,
exception: new_struct_error(vm, msg),
}
}
fn from_exception(exception: PyBaseExceptionRef, vm: &VirtualMachine) -> Self {
let kind = if exception.fast_isinstance(vm.ctx.exceptions.type_error) {
PackErrorKind::Type
} else if exception.fast_isinstance(vm.ctx.exceptions.overflow_error)
|| exception.fast_isinstance(vm.ctx.exceptions.value_error)
{
PackErrorKind::Value
} else {
PackErrorKind::Raised
};
Self { kind, exception }
}
fn raised(exception: PyBaseExceptionRef) -> Self {
Self {
kind: PackErrorKind::Raised,
exception,
}
}
}
static OVERFLOW_MSG: &str = "total struct size too long";
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum Endianness {
Native,
Little,
Big,
Host,
}
impl Endianness {
fn parse<I>(chars: &mut Peekable<I>) -> Self
where
I: Sized + Iterator<Item = u8>,
{
let e = match chars.peek() {
Some(b'@') => Self::Native,
Some(b'=') => Self::Host,
Some(b'<') => Self::Little,
Some(b'>' | b'!') => Self::Big,
_ => return Self::Native,
};
unsafe {
let _ = chars.next().unwrap_unchecked();
}
e
}
}
trait ByteOrder {
fn convert<I: PrimInt>(i: I) -> I;
}
enum BigEndian {}
impl ByteOrder for BigEndian {
fn convert<I: PrimInt>(i: I) -> I {
i.to_be()
}
}
enum LittleEndian {}
impl ByteOrder for LittleEndian {
fn convert<I: PrimInt>(i: I) -> I {
i.to_le()
}
}
type NativeEndian = cfg_select! {
target_endian = "big" => BigEndian,
target_endian = "little" => LittleEndian,
};
#[derive(Copy, Clone, num_enum::TryFromPrimitive, Eq, PartialEq)]
#[repr(u8)]
pub(crate) enum FormatType {
Pad = b'x',
SByte = b'b',
UByte = b'B',
Char = b'c',
WideChar = b'u',
Ucs4Char = b'w',
Str = b's',
Pascal = b'p',
Short = b'h',
UShort = b'H',
Int = b'i',
UInt = b'I',
Long = b'l',
ULong = b'L',
SSizeT = b'n',
SizeT = b'N',
LongLong = b'q',
ULongLong = b'Q',
Bool = b'?',
Half = b'e',
Float = b'f',
Double = b'd',
LongDouble = b'g',
FloatComplex = b'F',
DoubleComplex = b'D',
VoidP = b'P',
PyObject = b'O',
}
impl fmt::Debug for FormatType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(&(*self as u8 as char), f)
}
}
impl FormatType {
fn info(self, e: Endianness) -> &'static FormatInfo {
use mem::{align_of, size_of};
macro_rules! native_info {
($t:ty) => {{
&FormatInfo {
size: size_of::<$t>(),
align: align_of::<$t>(),
pack: Some(<$t as Packable>::pack::<NativeEndian>),
unpack: Some(<$t as Packable>::unpack::<NativeEndian>),
}
}};
}
macro_rules! nonnative_info {
($t:ty, $end:ty) => {{
&FormatInfo {
size: size_of::<$t>(),
align: 0,
pack: Some(<$t as Packable>::pack::<$end>),
unpack: Some(<$t as Packable>::unpack::<$end>),
}
}};
}
macro_rules! match_nonnative {
($zelf:expr, $end:ty) => {{
match $zelf {
Self::Pad | Self::Str | Self::Pascal => &FormatInfo {
size: size_of::<u8>(),
align: 0,
pack: None,
unpack: None,
},
Self::SByte => nonnative_info!(i8, $end),
Self::UByte => nonnative_info!(u8, $end),
Self::Char => &FormatInfo {
size: size_of::<u8>(),
align: 0,
pack: Some(pack_char),
unpack: Some(unpack_char),
},
Self::Short => nonnative_info!(i16, $end),
Self::UShort => nonnative_info!(u16, $end),
Self::Int | Self::Long => nonnative_info!(i32, $end),
Self::UInt | Self::ULong => nonnative_info!(u32, $end),
Self::LongLong => nonnative_info!(i64, $end),
Self::ULongLong => nonnative_info!(u64, $end),
Self::Bool => nonnative_info!(bool, $end),
Self::Half => nonnative_info!(f16, $end),
Self::Float => nonnative_info!(f32, $end),
Self::Double => nonnative_info!(f64, $end),
Self::LongDouble => nonnative_info!(f64, $end), Self::FloatComplex => nonnative_info!(PackFloatComplex, $end),
Self::DoubleComplex => nonnative_info!(PackDoubleComplex, $end),
Self::PyObject => nonnative_info!(usize, $end), _ => unreachable!(), }
}};
}
match e {
Endianness::Native => match self {
Self::Pad | Self::Str | Self::Pascal => &FormatInfo {
size: size_of::<raw::c_char>(),
align: 0,
pack: None,
unpack: None,
},
Self::SByte => native_info!(raw::c_schar),
Self::UByte => native_info!(raw::c_uchar),
Self::Char => &FormatInfo {
size: size_of::<raw::c_char>(),
align: 0,
pack: Some(pack_char),
unpack: Some(unpack_char),
},
Self::WideChar => native_info!(wchar_t),
Self::Ucs4Char => native_info!(u32),
Self::Short => native_info!(raw::c_short),
Self::UShort => native_info!(raw::c_ushort),
Self::Int => native_info!(raw::c_int),
Self::UInt => native_info!(raw::c_uint),
Self::Long => native_info!(raw::c_long),
Self::ULong => native_info!(raw::c_ulong),
Self::SSizeT => native_info!(isize), Self::SizeT => native_info!(usize), Self::LongLong => native_info!(raw::c_longlong),
Self::ULongLong => native_info!(raw::c_ulonglong),
Self::Bool => native_info!(bool),
Self::Half => native_info!(f16),
Self::Float => native_info!(raw::c_float),
Self::Double => native_info!(raw::c_double),
Self::LongDouble => native_info!(raw::c_double), Self::FloatComplex => native_info!(PackFloatComplex),
Self::DoubleComplex => native_info!(PackDoubleComplex),
Self::VoidP => native_info!(*mut raw::c_void),
Self::PyObject => native_info!(*mut raw::c_void), },
Endianness::Big => match_nonnative!(self, BigEndian),
Endianness::Little => match_nonnative!(self, LittleEndian),
Endianness::Host => match_nonnative!(self, NativeEndian),
}
}
}
#[derive(Debug, Clone)]
pub(crate) struct FormatCode {
pub repeat: usize,
pub code: FormatType,
pub info: &'static FormatInfo,
pub pre_padding: usize,
}
impl FormatCode {
pub(crate) const fn arg_count(&self) -> usize {
match self.code {
FormatType::Pad => 0,
FormatType::Str | FormatType::Pascal => 1,
_ => self.repeat,
}
}
pub(crate) fn parse<I>(
chars: &mut Peekable<I>,
endianness: Endianness,
) -> Result<(Vec<Self>, usize, usize), String>
where
I: Sized + Iterator<Item = u8>,
{
let mut offset = 0isize;
let mut arg_count = 0usize;
let mut codes = vec![];
while chars.peek().is_some() {
while let Some(b' ' | b'\t' | b'\n' | b'\r') = chars.peek() {
chars.next();
}
let repeat = match chars.peek() {
Some(b'0'..=b'9') => {
let mut repeat = 0isize;
while let Some(b'0'..=b'9') = chars.peek() {
if let Some(c) = chars.next() {
let current_digit = c - b'0';
repeat = repeat
.checked_mul(10)
.and_then(|r| r.checked_add(current_digit as _))
.ok_or_else(|| OVERFLOW_MSG.to_owned())?;
}
}
repeat
}
_ => 1,
};
let c = match chars.next() {
Some(c) => c,
None => {
if repeat != 1 {
return Err("repeat count given without format specifier".to_owned());
}
break;
}
};
if c == 0 {
return Err(exceptions::NulError.to_string());
}
if c == b'T' || c == b'X' {
if chars.peek() == Some(&b'{') {
chars.next(); let mut depth = 1;
while depth > 0 {
match chars.next() {
Some(b'{') => depth += 1,
Some(b'}') => depth -= 1,
None => return Err("unmatched '{' in format".to_owned()),
_ => {}
}
}
continue;
}
}
if c == b'(' {
let mut depth = 1;
while depth > 0 {
match chars.next() {
Some(b'(') => depth += 1,
Some(b')') => depth -= 1,
None => return Err("unmatched '(' in format".to_owned()),
_ => {}
}
}
continue;
}
if c == b':' {
loop {
match chars.next() {
Some(b':') => break,
None => return Err("unmatched ':' in format".to_owned()),
_ => {}
}
}
continue;
}
if c == b'{'
|| c == b'}'
|| c == b'&'
|| c == b'<'
|| c == b'>'
|| c == b'@'
|| c == b'='
|| c == b'!'
{
continue;
}
let code = FormatType::try_from(c)
.ok()
.filter(|c| match c {
FormatType::SSizeT
| FormatType::SizeT
| FormatType::VoidP
| FormatType::Ucs4Char => endianness == Endianness::Native,
_ => true,
})
.ok_or_else(|| "bad char in struct format".to_owned())?;
let info = code.info(endianness);
let padding = compensate_alignment(offset as usize, info.align)
.ok_or_else(|| OVERFLOW_MSG.to_owned())?;
offset = padding
.to_isize()
.and_then(|extra| offset.checked_add(extra))
.ok_or_else(|| OVERFLOW_MSG.to_owned())?;
let code = Self {
repeat: repeat as usize,
code,
info,
pre_padding: padding,
};
arg_count += code.arg_count();
codes.push(code);
offset = (info.size as isize)
.checked_mul(repeat)
.and_then(|item_size| offset.checked_add(item_size))
.ok_or_else(|| OVERFLOW_MSG.to_owned())?;
}
Ok((codes, offset as usize, arg_count))
}
}
const fn compensate_alignment(offset: usize, align: usize) -> Option<usize> {
if align != 0 && offset != 0 {
(align - 1).checked_sub((offset - 1) & (align - 1))
} else {
Some(0)
}
}
pub(crate) struct FormatInfo {
pub size: usize,
pub align: usize,
pub pack: Option<PackFunc>,
pub unpack: Option<UnpackFunc>,
}
impl fmt::Debug for FormatInfo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("FormatInfo")
.field("size", &self.size)
.field("align", &self.align)
.finish()
}
}
#[derive(Debug, Clone)]
pub struct FormatSpec {
#[allow(dead_code)]
pub(crate) endianness: Endianness,
pub(crate) codes: Vec<FormatCode>,
pub size: usize,
pub arg_count: usize,
}
#[derive(Default)]
pub struct FormatSpecCache {
entries: PyRwLock<HashMap<Box<[u8]>, PyRc<FormatSpec>>>,
}
impl FormatSpecCache {
pub fn get_or_parse(&self, format: &[u8], vm: &VirtualMachine) -> PyResult<PyRc<FormatSpec>> {
if let Some(spec) = self.entries.read().get(format) {
return Ok(spec.clone());
}
let spec = PyRc::new(FormatSpec::parse(format, vm)?);
let mut entries = self.entries.write();
if let Some(spec) = entries.get(format) {
return Ok(spec.clone());
}
if entries.len() >= 100 {
entries.clear();
}
entries.insert(format.into(), spec.clone());
Ok(spec)
}
pub fn clear(&self) {
*self.entries.write() = HashMap::new();
}
}
impl FormatSpec {
pub fn parse(fmt: &[u8], vm: &VirtualMachine) -> PyResult<Self> {
let mut chars = fmt.iter().copied().peekable();
let endianness = Endianness::parse(&mut chars);
let (codes, size, arg_count) =
FormatCode::parse(&mut chars, endianness).map_err(|err| new_struct_error(vm, err))?;
Ok(Self {
endianness,
codes,
size,
arg_count,
})
}
pub fn pack(&self, args: Vec<PyObjectRef>, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
self.try_pack(args, vm).map_err(|e| e.exception)
}
pub fn try_pack(
&self,
args: Vec<PyObjectRef>,
vm: &VirtualMachine,
) -> Result<Vec<u8>, PackError> {
let mut data = vm
.new_zeroed_bytes(self.size)
.map_err(|e| PackError::from_exception(e, vm))?;
self.try_pack_into(&mut data, args, vm)?;
Ok(data)
}
pub fn pack_into(
&self,
buffer: &mut [u8],
args: Vec<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<()> {
self.try_pack_into(buffer, args, vm)
.map_err(|e| e.exception)
}
pub fn try_pack_into(
&self,
mut buffer: &mut [u8],
args: Vec<PyObjectRef>,
vm: &VirtualMachine,
) -> Result<(), PackError> {
if self.arg_count != args.len() {
return Err(PackError::new(
PackErrorKind::Type,
vm,
format!(
"pack expected {} items for packing (got {})",
self.codes.len(),
args.len()
),
));
}
let mut args = args.into_iter();
for code in &self.codes {
buffer = &mut buffer[code.pre_padding..];
debug!("code: {code:?}");
match code.code {
FormatType::Str => {
let (buf, rest) = buffer.split_at_mut(code.repeat);
pack_string(vm, args.next().unwrap(), buf)
.map_err(|e| PackError::from_exception(e, vm))?;
buffer = rest;
}
FormatType::Pascal => {
let (buf, rest) = buffer.split_at_mut(code.repeat);
pack_pascal(vm, args.next().unwrap(), buf)
.map_err(|e| PackError::from_exception(e, vm))?;
buffer = rest;
}
FormatType::Pad => {
let (pad_buf, rest) = buffer.split_at_mut(code.repeat);
for el in pad_buf {
*el = 0
}
buffer = rest;
}
_ => {
let pack = code.info.pack.unwrap();
for arg in args.by_ref().take(code.repeat) {
let (item_buf, rest) = buffer.split_at_mut(code.info.size);
pack(vm, code.code, arg, item_buf)?;
buffer = rest;
}
}
}
}
Ok(())
}
pub fn unpack(&self, mut data: &[u8], vm: &VirtualMachine) -> PyResult<PyTupleRef> {
if self.size != data.len() {
return Err(new_struct_error(
vm,
format!("unpack requires a buffer of {} bytes", self.size),
));
}
let mut items = Vec::with_capacity(self.arg_count);
for code in &self.codes {
data = &data[code.pre_padding..];
debug!("unpack code: {code:?}");
match code.code {
FormatType::Pad => {
data = &data[code.repeat..];
}
FormatType::Str => {
let (str_data, rest) = data.split_at(code.repeat);
items.push(vm.ctx.new_bytes(str_data.to_vec()).into());
data = rest;
}
FormatType::Pascal => {
let (str_data, rest) = data.split_at(code.repeat);
items.push(unpack_pascal(vm, str_data));
data = rest;
}
_ => {
let unpack = code.info.unpack.unwrap();
for _ in 0..code.repeat {
let (item_data, rest) = data.split_at(code.info.size);
items.push(unpack(vm, item_data));
data = rest;
}
}
};
}
Ok(PyTuple::new_ref(items, &vm.ctx))
}
#[inline]
#[must_use]
pub const fn size(&self) -> usize {
self.size
}
#[must_use]
pub fn codes_sizeof(&self) -> usize {
core::mem::size_of::<FormatCode>() * (self.codes.len() + 1)
}
}
trait Packable {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError>;
fn unpack<E: ByteOrder>(vm: &VirtualMachine, data: &[u8]) -> PyObjectRef;
}
trait PackInt: PrimInt {
fn pack_int<E: ByteOrder>(self, data: &mut [u8]);
fn unpack_int<E: ByteOrder>(data: &[u8]) -> Self;
}
macro_rules! make_pack_prim_int {
($T:ty) => {
impl PackInt for $T {
fn pack_int<E: ByteOrder>(self, data: &mut [u8]) {
let i = E::convert(self);
data.copy_from_slice(&i.to_ne_bytes());
}
#[inline]
fn unpack_int<E: ByteOrder>(data: &[u8]) -> Self {
let mut x = [0; core::mem::size_of::<$T>()];
x.copy_from_slice(data);
E::convert(<$T>::from_ne_bytes(x))
}
}
impl Packable for $T {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
let i: $T = get_int_or_index(vm, code, &arg)?;
i.pack_int::<E>(data);
Ok(())
}
fn unpack<E: ByteOrder>(vm: &VirtualMachine, rdr: &[u8]) -> PyObjectRef {
let i = <$T>::unpack_int::<E>(rdr);
vm.ctx.new_int(i).into()
}
}
};
}
fn get_int_or_index<T>(
vm: &VirtualMachine,
code: FormatType,
arg: &PyObject,
) -> Result<T, PackError>
where
T: PrimInt + fmt::Display + for<'a> TryFrom<&'a BigInt>,
{
let index = match arg.try_index_opt(vm) {
None => {
return Err(PackError::new(
PackErrorKind::Type,
vm,
"required argument is not an integer",
));
}
Some(Err(e)) => return Err(PackError::from_exception(e, vm)),
Some(Ok(index)) => index,
};
index.try_to_primitive(vm).map_err(|_| {
let msg = if code == FormatType::VoidP {
"int too large to convert".to_owned()
} else {
format!(
"'{}' format requires {} <= number <= {}",
code as u8 as char,
T::min_value(),
T::max_value()
)
};
PackError::new(PackErrorKind::Value, vm, msg)
})
}
make_pack_prim_int!(i8);
make_pack_prim_int!(u8);
make_pack_prim_int!(i16);
make_pack_prim_int!(u16);
make_pack_prim_int!(i32);
make_pack_prim_int!(u32);
make_pack_prim_int!(i64);
make_pack_prim_int!(u64);
make_pack_prim_int!(usize);
make_pack_prim_int!(isize);
macro_rules! make_pack_float {
($T:ty, $fmt:literal) => {
impl Packable for $T {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
_code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
let f_64 = ArgIntoFloat::try_from_object(vm, arg)
.map_err(|e| PackError::from_exception(e, vm))?
.into_float();
let f = f_64 as $T;
if f.is_infinite() != f_64.is_infinite() {
return Err(PackError {
kind: PackErrorKind::Value,
exception: vm.new_overflow_error(concat!(
"float too large to pack with ",
$fmt,
" format"
)),
});
}
f.to_bits().pack_int::<E>(data);
Ok(())
}
fn unpack<E: ByteOrder>(vm: &VirtualMachine, rdr: &[u8]) -> PyObjectRef {
let i = PackInt::unpack_int::<E>(rdr);
<$T>::from_bits(i).to_pyobject(vm)
}
}
};
}
make_pack_float!(f32, "f");
make_pack_float!(f64, "d");
#[repr(C)]
struct PackFloatComplex(f32, f32);
#[repr(C)]
struct PackDoubleComplex(f64, f64);
macro_rules! make_pack_complex {
($T:ty, $Elem:ty, $Bits:ty, $fmt:literal) => {
impl Packable for $T {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
_code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
let c = if let Some(value) = arg.downcast_ref::<PyComplex>() {
value.as_complex()
} else {
ArgIntoComplex::try_from_object(vm, arg)
.map_err(|_| {
PackError::new(
PackErrorKind::Type,
vm,
"required argument is not a complex",
)
})?
.into_complex()
};
for (component, bytes) in [c.re, c.im]
.into_iter()
.zip(data.chunks_exact_mut(size_of::<$Elem>()))
{
let narrowed = component as $Elem;
if E::convert(1u16) != 1 && narrowed.is_infinite() != component.is_infinite() {
return Err(PackError {
kind: PackErrorKind::Value,
exception: vm.new_overflow_error(concat!(
"float too large to pack with ",
$fmt,
" format"
)),
});
}
narrowed.to_bits().pack_int::<E>(bytes);
}
Ok(())
}
fn unpack<E: ByteOrder>(vm: &VirtualMachine, rdr: &[u8]) -> PyObjectRef {
let half = size_of::<$Elem>();
let re = <$Elem>::from_bits(<$Bits>::unpack_int::<E>(&rdr[..half])) as f64;
let im = <$Elem>::from_bits(<$Bits>::unpack_int::<E>(&rdr[half..half * 2])) as f64;
vm.ctx.new_complex(Complex64::new(re, im)).into()
}
}
};
}
make_pack_complex!(PackFloatComplex, f32, u32, "f");
make_pack_complex!(PackDoubleComplex, f64, u64, "d");
impl Packable for f16 {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
_code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
let f_64 = ArgIntoFloat::try_from_object(vm, arg)
.map_err(|e| PackError::from_exception(e, vm))?
.into_float();
let f_16 = Self::from_f64_const(f_64);
if f_16.is_infinite() != f_64.is_infinite() {
return Err(PackError {
kind: PackErrorKind::Value,
exception: vm.new_overflow_error("float too large to pack with e format"),
});
}
f_16.to_bits().pack_int::<E>(data);
Ok(())
}
fn unpack<E: ByteOrder>(vm: &VirtualMachine, rdr: &[u8]) -> PyObjectRef {
let i = PackInt::unpack_int::<E>(rdr);
Self::from_bits(i).to_f64().to_pyobject(vm)
}
}
impl Packable for *mut raw::c_void {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
usize::pack::<E>(vm, code, arg, data)
}
fn unpack<E: ByteOrder>(vm: &VirtualMachine, rdr: &[u8]) -> PyObjectRef {
usize::unpack::<E>(vm, rdr)
}
}
impl Packable for bool {
fn pack<E: ByteOrder>(
vm: &VirtualMachine,
_code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
let v = ArgIntoBool::try_from_object(vm, arg)
.map_err(PackError::raised)?
.into_bool() as u8;
v.pack_int::<E>(data);
Ok(())
}
fn unpack<E: ByteOrder>(vm: &VirtualMachine, rdr: &[u8]) -> PyObjectRef {
let i = u8::unpack_int::<E>(rdr);
vm.ctx.new_bool(i != 0).into()
}
}
fn pack_char(
vm: &VirtualMachine,
_code: FormatType,
arg: PyObjectRef,
data: &mut [u8],
) -> Result<(), PackError> {
let v = PyBytesRef::try_from_object(vm, arg).map_err(|e| PackError::from_exception(e, vm))?;
let ch = *v.as_bytes().iter().exactly_one().map_err(|_| {
PackError::new(
PackErrorKind::Value,
vm,
"char format requires a bytes object of length 1",
)
})?;
data[0] = ch;
Ok(())
}
fn pack_string(vm: &VirtualMachine, arg: PyObjectRef, buf: &mut [u8]) -> PyResult<()> {
let b = ArgBytesLike::try_from_object(vm, arg)?;
b.with_ref(|data| write_string(buf, data));
Ok(())
}
fn pack_pascal(vm: &VirtualMachine, arg: PyObjectRef, buf: &mut [u8]) -> PyResult<()> {
if buf.is_empty() {
return Ok(());
}
let b = ArgBytesLike::try_from_object(vm, arg)?;
b.with_ref(|data| {
let string_length = core::cmp::min(core::cmp::min(data.len(), 255), buf.len() - 1);
buf[0] = string_length as u8;
write_string(&mut buf[1..], data);
});
Ok(())
}
fn write_string(buf: &mut [u8], data: &[u8]) {
let len_from_data = core::cmp::min(data.len(), buf.len());
buf[..len_from_data].copy_from_slice(&data[..len_from_data]);
for byte in &mut buf[len_from_data..] {
*byte = 0
}
}
fn unpack_char(vm: &VirtualMachine, data: &[u8]) -> PyObjectRef {
vm.ctx.new_bytes(vec![data[0]]).into()
}
fn unpack_pascal(vm: &VirtualMachine, data: &[u8]) -> PyObjectRef {
let (&len, data) = match data.split_first() {
Some(x) => x,
None => {
return vm.ctx.new_bytes(vec![]).into();
}
};
let len = core::cmp::min(len as usize, data.len());
vm.ctx.new_bytes(data[..len].to_vec()).into()
}
pub fn struct_error_type(vm: &VirtualMachine) -> &'static Py<PyType> {
static_cell! {
static INSTANCE: PyTypeRef;
}
INSTANCE.get_or_init(|| vm.ctx.new_exception_type("struct", "error", None))
}
pub fn new_struct_error<T: Into<Wtf8Buf>>(vm: &VirtualMachine, msg: T) -> PyBaseExceptionRef {
vm.new_exception_msg(struct_error_type(vm).to_owned(), msg.into())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Interpreter;
#[test]
fn format_cache_reuses_specs_and_releases_evicted_entries() {
Interpreter::without_stdlib(Default::default()).enter(|vm| {
let cache = &vm.state.struct_format_cache;
let spec = cache.get_or_parse(b"<IH", vm).unwrap();
assert!(PyRc::ptr_eq(
&spec,
&cache.get_or_parse(b"<IH", vm).unwrap()
));
let weak = PyRc::downgrade(&spec);
drop(spec);
for padding in 0..100 {
let format = format!("{padding}x");
cache.get_or_parse(format.as_bytes(), vm).unwrap();
}
assert!(weak.upgrade().is_none());
let spec = cache.get_or_parse(b"<IH", vm).unwrap();
let weak = PyRc::downgrade(&spec);
drop(spec);
cache.clear();
assert!(weak.upgrade().is_none());
});
}
}