diff --git a/crates/common/src/lock/cell_lock.rs b/crates/common/src/lock/cell_lock.rs index 0e045c5950b..9732b973aeb 100644 --- a/crates/common/src/lock/cell_lock.rs +++ b/crates/common/src/lock/cell_lock.rs @@ -89,7 +89,7 @@ unsafe impl RawRwLock for RawCellRwLock { #[inline] unsafe fn unlock_shared(&self) { - self.state.set(self.state.get() - ONE_READER) + self.state.update(|x| x - ONE_READER) } #[inline] diff --git a/crates/compiler-core/src/marshal.rs b/crates/compiler-core/src/marshal.rs index decb25d5283..fa568226aa1 100644 --- a/crates/compiler-core/src/marshal.rs +++ b/crates/compiler-core/src/marshal.rs @@ -158,13 +158,17 @@ impl Read for &[u8] { fn read_slice(&mut self, n: u32) -> Result<&[u8]> { self.read_slice_borrow(n) } + + fn read_array(&mut self) -> Result<&[u8; N]> { + let (chunk, rest) = self.split_first_chunk::().ok_or(MarshalError::Eof)?; + *self = rest; + Ok(chunk) + } } impl<'a> ReadBorrowed<'a> for &'a [u8] { fn read_slice_borrow(&mut self, n: u32) -> Result<&'a [u8]> { - let data = self.get(..n as usize).ok_or(MarshalError::Eof)?; - *self = &self[n as usize..]; - Ok(data) + self.split_off(..n as usize).ok_or(MarshalError::Eof) } } diff --git a/crates/vm/src/codecs.rs b/crates/vm/src/codecs.rs index 3241dee4981..cdae4c2ba13 100644 --- a/crates/vm/src/codecs.rs +++ b/crates/vm/src/codecs.rs @@ -476,7 +476,7 @@ impl<'a> DecodeErrorHandler> for SurrogatePass { let p = &s[byte_range.start..]; fn slice(p: &[u8]) -> Option<[u8; N]> { - p.get(..N).map(|x| x.try_into().unwrap()) + p.first_chunk().copied() } let c = match standard_encoding { diff --git a/crates/vm/src/stdlib/ctypes/array.rs b/crates/vm/src/stdlib/ctypes/array.rs index eea1fd765d5..9032bf01f0b 100644 --- a/crates/vm/src/stdlib/ctypes/array.rs +++ b/crates/vm/src/stdlib/ctypes/array.rs @@ -712,27 +712,22 @@ impl PyCArray { } Some("f") => { // c_float - if offset + 4 <= buffer.len() { - let bytes: [u8; 4] = buffer[offset..offset + 4].try_into().unwrap(); - let val = f32::from_ne_bytes(bytes); - Ok(vm.ctx.new_float(val as f64).into()) - } else { - Ok(vm.ctx.new_float(0.0).into()) - } + let val = buffer[offset..] + .first_chunk::<4>() + .copied() + .map_or(0.0, f32::from_ne_bytes); + Ok(vm.ctx.new_float(val as f64).into()) } Some("d") | Some("g") => { // c_double / c_longdouble - read f64 from first 8 bytes - if offset + 8 <= buffer.len() { - let bytes: [u8; 8] = buffer[offset..offset + 8].try_into().unwrap(); - let val = f64::from_ne_bytes(bytes); - Ok(vm.ctx.new_float(val).into()) - } else { - Ok(vm.ctx.new_float(0.0).into()) - } + let val = buffer[offset..] + .first_chunk::<8>() + .copied() + .map_or(0.0, f64::from_ne_bytes); + Ok(vm.ctx.new_float(val).into()) } _ => { - if offset + element_size <= buffer.len() { - let bytes = &buffer[offset..offset + element_size]; + if let Some(bytes) = buffer[offset..].get(..element_size) { Ok(Self::bytes_to_int(bytes, element_size, type_code, vm)) } else { Ok(vm.ctx.new_int(0).into()) diff --git a/crates/vm/src/stdlib/ctypes/base.rs b/crates/vm/src/stdlib/ctypes/base.rs index 58d9466adb2..55cf358dce3 100644 --- a/crates/vm/src/stdlib/ctypes/base.rs +++ b/crates/vm/src/stdlib/ctypes/base.rs @@ -1836,71 +1836,53 @@ pub(super) fn buffer_to_ffi_value(type_code: &str, buffer: &[u8]) -> FfiArgValue FfiArgValue::U8(v) } "h" => { - let v = if buffer.len() >= 2 { - i16::from_ne_bytes(buffer[..2].try_into().unwrap()) - } else { - 0 - }; + let v = buffer.first_chunk().copied().map_or(0, i16::from_ne_bytes); FfiArgValue::I16(v) } "H" => { - let v = if buffer.len() >= 2 { - u16::from_ne_bytes(buffer[..2].try_into().unwrap()) - } else { - 0 - }; + let v = buffer.first_chunk().copied().map_or(0, u16::from_ne_bytes); FfiArgValue::U16(v) } "i" => { - let v = if buffer.len() >= 4 { - i32::from_ne_bytes(buffer[..4].try_into().unwrap()) - } else { - 0 - }; + let v = buffer.first_chunk().copied().map_or(0, i32::from_ne_bytes); FfiArgValue::I32(v) } "I" => { - let v = if buffer.len() >= 4 { - u32::from_ne_bytes(buffer[..4].try_into().unwrap()) - } else { - 0 - }; + let v = buffer.first_chunk().copied().map_or(0, u32::from_ne_bytes); FfiArgValue::U32(v) } "l" | "q" => { - let v = if buffer.len() >= 8 { - i64::from_ne_bytes(buffer[..8].try_into().unwrap()) - } else if buffer.len() >= 4 { - i32::from_ne_bytes(buffer[..4].try_into().unwrap()) as i64 + let v = if let Some(&bytes) = buffer.first_chunk::<8>() { + i64::from_ne_bytes(bytes) + } else if let Some(&bytes) = buffer.first_chunk::<4>() { + i32::from_ne_bytes(bytes).into() } else { 0 }; FfiArgValue::I64(v) } "L" | "Q" => { - let v = if buffer.len() >= 8 { - u64::from_ne_bytes(buffer[..8].try_into().unwrap()) - } else if buffer.len() >= 4 { - u32::from_ne_bytes(buffer[..4].try_into().unwrap()) as u64 + let v = if let Some(&bytes) = buffer.first_chunk::<8>() { + u64::from_ne_bytes(bytes) + } else if let Some(&bytes) = buffer.first_chunk::<4>() { + u32::from_ne_bytes(bytes).into() } else { 0 }; FfiArgValue::U64(v) } "f" => { - let v = if buffer.len() >= 4 { - f32::from_ne_bytes(buffer[..4].try_into().unwrap()) - } else { - 0.0 - }; + let v = buffer + .first_chunk::<4>() + .copied() + .map_or(0.0, f32::from_ne_bytes); FfiArgValue::F32(v) } "d" | "g" => { - let v = if buffer.len() >= 8 { - f64::from_ne_bytes(buffer[..8].try_into().unwrap()) - } else { - 0.0 - }; + let v = buffer + .first_chunk::<8>() + .copied() + .map_or(0.0, f64::from_ne_bytes); FfiArgValue::F64(v) } "z" | "Z" | "P" | "O" => FfiArgValue::Pointer(read_ptr_from_buffer(buffer)), @@ -1910,11 +1892,7 @@ pub(super) fn buffer_to_ffi_value(type_code: &str, buffer: &[u8]) -> FfiArgValue } "u" => { // wchar_t - 4 bytes on most platforms - let v = if buffer.len() >= 4 { - u32::from_ne_bytes(buffer[..4].try_into().unwrap()) - } else { - 0 - }; + let v = buffer.first_chunk().copied().map_or(0, u32::from_ne_bytes); FfiArgValue::U32(v) } _ => FfiArgValue::Pointer(0), @@ -2135,11 +2113,10 @@ pub(super) fn get_usize_attr( #[inline] pub(super) fn read_ptr_from_buffer(buffer: &[u8]) -> usize { const PTR_SIZE: usize = core::mem::size_of::(); - if buffer.len() >= PTR_SIZE { - usize::from_ne_bytes(buffer[..PTR_SIZE].try_into().unwrap()) - } else { - 0 - } + buffer + .first_chunk::() + .copied() + .map_or(0, usize::from_ne_bytes) } /// Check if a type is a "simple instance" (direct subclass of a simple type) diff --git a/crates/vm/src/stdlib/ctypes/function.rs b/crates/vm/src/stdlib/ctypes/function.rs index 295e6fd137d..3ea166f9871 100644 --- a/crates/vm/src/stdlib/ctypes/function.rs +++ b/crates/vm/src/stdlib/ctypes/function.rs @@ -567,10 +567,8 @@ fn extract_ptr_from_arg(arg: &PyObject, vm: &VirtualMachine) -> PyResult } if let Some(simple) = arg.downcast_ref::() { let buffer = simple.0.buffer.read(); - if buffer.len() >= core::mem::size_of::() { - return Ok(usize::from_ne_bytes( - buffer[..core::mem::size_of::()].try_into().unwrap(), - )); + if let Some(&bytes) = buffer.first_chunk::<{ size_of::() }>() { + return Ok(usize::from_ne_bytes(bytes)); } } if let Some(cdata) = arg.downcast_ref::() { diff --git a/crates/vm/src/stdlib/ctypes/simple.rs b/crates/vm/src/stdlib/ctypes/simple.rs index b2ae0f7cc5b..410628b5039 100644 --- a/crates/vm/src/stdlib/ctypes/simple.rs +++ b/crates/vm/src/stdlib/ctypes/simple.rs @@ -419,13 +419,10 @@ impl PyCSimpleType { if let Some(funcptr) = value.downcast_ref::() { let ptr_val = { let buffer = funcptr._base.buffer.read(); - if buffer.len() >= core::mem::size_of::() { - usize::from_ne_bytes( - buffer[..core::mem::size_of::()].try_into().unwrap(), - ) - } else { - 0 - } + buffer + .first_chunk::<{ size_of::() }>() + .copied() + .map_or(0, usize::from_ne_bytes) }; return Ok(CArgObject { tag: b'P', @@ -442,13 +439,10 @@ impl PyCSimpleType { if matches!(value_type_code.as_deref(), Some("z") | Some("Z")) { let ptr_val = { let buffer = simple.0.buffer.read(); - if buffer.len() >= core::mem::size_of::() { - usize::from_ne_bytes( - buffer[..core::mem::size_of::()].try_into().unwrap(), - ) - } else { - 0 - } + buffer + .first_chunk::<{ size_of::() }>() + .copied() + .map_or(0, usize::from_ne_bytes) }; return Ok(CArgObject { tag: b'Z', @@ -1360,68 +1354,47 @@ impl PyCSimple { let buffer = self.0.buffer.read(); let bytes: &[u8] = &buffer; - if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::u8().as_raw_ptr()) { - if !bytes.is_empty() { - return Some(FfiArgValue::U8(bytes[0])); - } + let ret = if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::u8().as_raw_ptr()) { + let byte = *bytes.first()?; + FfiArgValue::U8(byte) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::i8().as_raw_ptr()) { - if !bytes.is_empty() { - return Some(FfiArgValue::I8(bytes[0] as i8)); - } + let byte = *bytes.first()?; + FfiArgValue::I8(byte as i8) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::u16().as_raw_ptr()) { - if bytes.len() >= 2 { - return Some(FfiArgValue::U16(u16::from_ne_bytes([bytes[0], bytes[1]]))); - } + let bytes = *bytes.first_chunk::<2>()?; + FfiArgValue::U16(u16::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::i16().as_raw_ptr()) { - if bytes.len() >= 2 { - return Some(FfiArgValue::I16(i16::from_ne_bytes([bytes[0], bytes[1]]))); - } + let bytes = *bytes.first_chunk::<2>()?; + FfiArgValue::I16(i16::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::u32().as_raw_ptr()) { - if bytes.len() >= 4 { - return Some(FfiArgValue::U32(u32::from_ne_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], - ]))); - } + let bytes = *bytes.first_chunk::<4>()?; + FfiArgValue::U32(u32::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::i32().as_raw_ptr()) { - if bytes.len() >= 4 { - return Some(FfiArgValue::I32(i32::from_ne_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], - ]))); - } + let bytes = *bytes.first_chunk::<4>()?; + FfiArgValue::I32(i32::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::u64().as_raw_ptr()) { - if bytes.len() >= 8 { - return Some(FfiArgValue::U64(u64::from_ne_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], - ]))); - } + let bytes = *bytes.first_chunk::<8>()?; + FfiArgValue::U64(u64::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::i64().as_raw_ptr()) { - if bytes.len() >= 8 { - return Some(FfiArgValue::I64(i64::from_ne_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], - ]))); - } + let bytes = *bytes.first_chunk::<8>()?; + FfiArgValue::I64(i64::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::f32().as_raw_ptr()) { - if bytes.len() >= 4 { - return Some(FfiArgValue::F32(f32::from_ne_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], - ]))); - } + let bytes = *bytes.first_chunk::<4>()?; + FfiArgValue::F32(f32::from_ne_bytes(bytes)) } else if core::ptr::eq(ty.as_raw_ptr(), libffi::middle::Type::f64().as_raw_ptr()) { - if bytes.len() >= 8 { - return Some(FfiArgValue::F64(f64::from_ne_bytes([ - bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], - ]))); - } + let bytes = *bytes.first_chunk::<8>()?; + FfiArgValue::F64(f64::from_ne_bytes(bytes)) } else if core::ptr::eq( ty.as_raw_ptr(), libffi::middle::Type::pointer().as_raw_ptr(), - ) && bytes.len() >= core::mem::size_of::() - { - let val = - usize::from_ne_bytes(bytes[..core::mem::size_of::()].try_into().unwrap()); - return Some(FfiArgValue::Pointer(val)); - } - None + ) { + let bytes = *buffer.first_chunk::<{ size_of::() }>()?; + let val = usize::from_ne_bytes(bytes); + FfiArgValue::Pointer(val) + } else { + return None; + }; + Some(ret) } } diff --git a/crates/vm/src/stdlib/io.rs b/crates/vm/src/stdlib/io.rs index b98de6a87a5..b270fa2529b 100644 --- a/crates/vm/src/stdlib/io.rs +++ b/crates/vm/src/stdlib/io.rs @@ -2506,15 +2506,11 @@ mod _io { return None; } buf.resize(Self::BYTE_LEN, 0); - let buf: &[u8; Self::BYTE_LEN] = buf.as_slice().try_into().unwrap(); + let buf: &[u8; Self::BYTE_LEN] = buf.as_array()?; macro_rules! get_field { - ($t:ty, $off:ident) => {{ - <$t>::from_ne_bytes( - buf[Self::$off..][..core::mem::size_of::<$t>()] - .try_into() - .unwrap(), - ) - }}; + ($t:ty, $off:ident) => { + <$t>::from_ne_bytes(*buf[Self::$off..].first_chunk().unwrap()) + }; } Some(Self { start_pos: get_field!(Offset, START_POS_OFF), diff --git a/crates/vm/src/stdlib/winreg.rs b/crates/vm/src/stdlib/winreg.rs index 53a4fd0d556..ec8cec4e337 100644 --- a/crates/vm/src/stdlib/winreg.rs +++ b/crates/vm/src/stdlib/winreg.rs @@ -905,20 +905,18 @@ mod winreg { match typ { REG_DWORD => { // If there isn’t enough data, return 0. - if ret_data.len() < std::mem::size_of::() { - Ok(vm.ctx.new_int(0).into()) - } else { - let val = u32::from_ne_bytes(ret_data[..4].try_into().unwrap()); - Ok(vm.ctx.new_int(val).into()) - } + let val = ret_data + .first_chunk::<4>() + .copied() + .map_or(0, u32::from_ne_bytes); + Ok(vm.ctx.new_int(val).into()) } REG_QWORD => { - if ret_data.len() < std::mem::size_of::() { - Ok(vm.ctx.new_int(0).into()) - } else { - let val = u64::from_ne_bytes(ret_data[..8].try_into().unwrap()); - Ok(vm.ctx.new_int(val).into()) - } + let val = ret_data + .first_chunk::<8>() + .copied() + .map_or(0, u64::from_ne_bytes); + Ok(vm.ctx.new_int(val).into()) } REG_SZ | REG_EXPAND_SZ => { let u16_slice = bytes_as_wide_slice(ret_data); diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 3a534302c31..1af3ea4555c 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -745,10 +745,9 @@ impl VirtualMachine { return Err(self.new_recursion_error(_where.to_string())); } - self.recursion_depth.set(self.recursion_depth.get() + 1); - let result = f(); - self.recursion_depth.set(self.recursion_depth.get() - 1); - result + self.recursion_depth.update(|d| d + 1); + scopeguard::defer! { self.recursion_depth.update(|d| d - 1) } + f() } pub fn with_frame PyResult>( diff --git a/crates/wtf8/src/core_str_count.rs b/crates/wtf8/src/core_str_count.rs index f02f0a5708d..8f9d5585bc7 100644 --- a/crates/wtf8/src/core_str_count.rs +++ b/crates/wtf8/src/core_str_count.rs @@ -60,7 +60,7 @@ fn do_count_chars(s: &Wtf8) -> usize { // a subset of the sum of this chunk, like a `[u8; size_of::()]`. let mut counts = 0; - let (unrolled_chunks, remainder) = slice_as_chunks::<_, UNROLL_INNER>(chunk); + let (unrolled_chunks, remainder) = chunk.as_chunks::(); for unrolled in unrolled_chunks { for &word in unrolled { // Because `CHUNK_SIZE` is < 256, this addition can't cause the @@ -137,26 +137,6 @@ const fn usize_repeat_u16(x: u16) -> usize { } r } - -fn slice_as_chunks(slice: &[T]) -> (&[[T; N]], &[T]) { - assert!(N != 0, "chunk size must be non-zero"); - let len_rounded_down = slice.len() / N * N; - // SAFETY: The rounded-down value is always the same or smaller than the - // original length, and thus must be in-bounds of the slice. - let (multiple_of_n, remainder) = unsafe { slice.split_at_unchecked(len_rounded_down) }; - // SAFETY: We already panicked for zero, and ensured by construction - // that the length of the subslice is a multiple of N. - let array_slice = unsafe { slice_as_chunks_unchecked(multiple_of_n) }; - (array_slice, remainder) -} - -unsafe fn slice_as_chunks_unchecked(slice: &[T]) -> &[[T; N]] { - let new_len = slice.len() / N; - // SAFETY: We cast a slice of `new_len * N` elements into - // a slice of `new_len` many `N` elements chunks. - unsafe { std::slice::from_raw_parts(slice.as_ptr().cast(), new_len) } -} - const fn unlikely(x: bool) -> bool { x }