firefox-main/third_party/rust/objc2-foundation/src/iter.rs

Source code

Revision control

Copy as Markdown

Other Tools

#![allow(dead_code)]
use core::ffi::c_ulong;
use core::ptr::{self, NonNull};
use objc2::rc::Retained;
use objc2::runtime::AnyObject;
use objc2::runtime::ProtocolObject;
use objc2::ClassType;
use objc2::Message;
use crate::{NSFastEnumeration, NSFastEnumerationState};
/// Swift and Objective-C both have a stack buffer size of 16, so we do that
/// as well.
///
/// TODO: Consider lowering this a bit, since the most common type of
/// enumeration (e.g. NSArray) doesn't use the buffer:
const BUF_SIZE: usize = 16;
/// Track mutation mistakes when debug assertions are enabled (they should
/// be made impossible by Rust at compile-time, so hence we don't need to
/// track them in release mode).
///
/// This is set to `None` initially, but later loaded to `Some(_)` after
/// the first enumeration.
type MutationState = Option<c_ulong>;
/// Helper type for doing fast enumeration.
///
/// See the following other implementations of this:
#[derive(Debug, PartialEq)]
struct FastEnumeratorHelper {
state: NSFastEnumerationState,
buf: [*mut AnyObject; BUF_SIZE],
// TODO: We could possibly optimize things a bit by doing
// `itemsPtr.add(1); items_count -= 1;` on every loop, instead of storing
// `current_item` - but it's not really defined whether we're allowed to
// mutate `itemsPtr`? All of the GNUStep code and the open sourced Apple
// code I've found would work under this scheme, and from a cursory GitHub
// search I only found the following that assume that `itemsPtr` is stable
// across invocations:
//
// Though they also assume that the given stack buffer has a stable
// address, which is something that would be prohibitively expensive to
// ensure, so don't think we should consider that one.
current_item: usize,
items_count: usize,
}
// SAFETY: Neither `FastEnumeratorHelper`, nor the inner enumeration state,
// need to be bound to any specific thread (at least, we can generally assume
// this for the enumerable types in Foundation - but `NSEnumerator` won't be
// `Send`, and neither will its iterator types).
unsafe impl Send for FastEnumeratorHelper {}
// SAFETY: `FastEnumeratorHelper` is only mutable behind `&mut`, and as such
// is safe to share with other threads.
unsafe impl Sync for FastEnumeratorHelper {}
#[cfg(feature = "unstable-mutation-return-null")]
extern "C" {
static kCFNull: NonNull<AnyObject>;
}
// Extracted as separate functions to make assembly nicer
#[cfg(not(feature = "unstable-mutation-return-null"))]
#[cfg_attr(debug_assertions, track_caller)]
fn items_ptr_null() -> ! {
panic!("`itemsPtr` was NULL, likely due to mutation during iteration");
}
#[cfg(not(feature = "unstable-mutation-return-null"))]
#[cfg_attr(debug_assertions, track_caller)]
fn mutation_detected() -> ! {
panic!("mutation detected during enumeration");
}
impl FastEnumeratorHelper {
#[inline]
fn new() -> Self {
Self {
state: NSFastEnumerationState {
state: 0,
itemsPtr: ptr::null_mut(),
mutationsPtr: ptr::null_mut(),
extra: [0; 5],
},
buf: [ptr::null_mut(); BUF_SIZE],
current_item: 0,
items_count: 0,
}
}
#[inline]
const fn remaining_items_at_least(&self) -> usize {
self.items_count - self.current_item
}
/// Load the next array of items into the enumeration state.
///
///
/// # Safety
///
/// The collection must be the same on each call.
#[inline]
unsafe fn load_next_items(&mut self, collection: &ProtocolObject<dyn NSFastEnumeration>) {
// SAFETY: The ptr comes from a slice, which is always non-null.
//
// The callee may choose to use this or not, see:
let buf_ptr = unsafe { NonNull::new_unchecked(self.buf.as_mut_ptr()) };
// SAFETY:
// - The state is mutable, and we don't touch the `state` and
// `extra` fields.
// - The buffer is mutable and its length is correctly set.
// - The collection and state are guaranteed by the caller to match.
self.items_count = unsafe {
collection.countByEnumeratingWithState_objects_count(
NonNull::from(&mut self.state),
buf_ptr,
self.buf.len(),
)
};
// For unwind safety, we only set this _after_ the message send has
// completed, otherwise, if it unwinded, upon iterating again we might
// invalidly assume that the buffer was ready.
self.current_item = 0;
}
/// Get the next item from the given collection.
///
/// We use a `ProtocolObject` instead of a generic, so that there is only
/// one instance of this function in the compilation unit (should improve
/// compilation speed).
///
/// If `mutations_state` is `Some`, mutation tracking is activated.
///
///
/// # Safety
///
/// The collection must be the same on each call.
#[inline]
#[track_caller]
unsafe fn next_from(
&mut self,
collection: &ProtocolObject<dyn NSFastEnumeration>,
mutations_state: Option<&mut MutationState>,
) -> Option<NonNull<AnyObject>> {
// If we've exhausted the current array of items.
if self.current_item >= self.items_count {
// Get the next array of items.
//
// SAFETY: Upheld by caller.
unsafe { self.load_next_items(collection) };
// If the next array was also empty.
if self.items_count == 0 {
// We are done enumerating.
return None;
}
if mutations_state.is_some() && self.state.itemsPtr.is_null() {
#[cfg(feature = "unstable-mutation-return-null")]
return Some(unsafe { kCFNull });
#[cfg(not(feature = "unstable-mutation-return-null"))]
items_ptr_null();
}
}
if let Some(mutations_state) = mutations_state {
// If the mutation ptr is not set, we do nothing.
if let Some(ptr) = NonNull::new(self.state.mutationsPtr) {
// SAFETY:
// - The pointer is not NULL.
//
// - The enumerator is expected to give back a dereferenceable
// pointer, that is alive for as long as the collection is
// alive.
//
// Note that iterating past the first returned `None` is not
// tested by most Objective-C implementations, so it may
// deallocate the mutations ptr in that case?
//
// - The enumeration should not be modifiable across threads,
// so neither will this pointer be accessed from different
// threads.
//
// Note that this assumption is relatively likely to be
// violated, but if that is the case, the program already
// has UB, so then it is better that we detect it.
//
// - The value is an integer, so is always initialized.
//
//
// We do an unaligned read here since we have no guarantees
// about this pointer, and efficiency doesn't really matter.
let new_state = unsafe { ptr.as_ptr().read_unaligned() };
match *mutations_state {
// On the first iteration, initialize the mutation state
None => {
*mutations_state = Some(new_state);
}
// On subsequent iterations, verify that the state hasn't
// changed.
Some(current_state) => {
if current_state != new_state {
#[cfg(feature = "unstable-mutation-return-null")]
return Some(unsafe { kCFNull });
#[cfg(not(feature = "unstable-mutation-return-null"))]
mutation_detected();
}
}
}
}
}
// Compute a pointer to the current item.
//
// SAFETY: The index is checked above to be in bounds of the returned
// array, and the item pointer is checked to not be NULL.
let ptr = unsafe { self.state.itemsPtr.add(self.current_item) };
// Move to the next item.
self.current_item += 1;
// And read the current item.
//
// SAFETY:
// - The pointer is not NULL, it is guaranteed to have been set
// by `countByEnumeratingWithState:objects:count:`.
// - The pointer is within the bounds of the returned array.
//
// Note: GNUStep may sometimes return unaligned pointers. TODO: Should
// we go back to using `read` on Apple systems for better performance?
let obj = unsafe { ptr.read_unaligned() };
// SAFETY: The returned array contains no NULL pointers.
Some(unsafe { NonNull::new_unchecked(obj) })
}
}
// Unfortunately, `NSFastEnumeration` doesn't provide a way for enumerated
// objects to clean up their work after having being enumerated over.
//
// impl Drop for FastEnumeratorHelper { ... }
/// Internal helper trait to figure out the output type of something that
/// implements `NSFastEnumeration`.
///
///
/// # Safety
///
/// The associated `Item` type must be the type that is used in Objective-C's
/// `for (type elem in collection) { stmts; }` enumeration, and the collection
/// itself must not contain any lifetime parameter.
pub(crate) unsafe trait FastEnumerationHelper: Message + NSFastEnumeration {
type Item: Message;
fn maybe_len(&self) -> Option<usize>;
}
// Iterator implementations we _can't_ do:
//
//
// # `ExactSizeIterator`
//
// This could probably be implemented correctly for the concrete class
// `NSArray`, but not for arbitrary `NSArray` subclasses, which are always be
// valid to convert to `NSArray`.
//
//
// # `FusedIterator`
//
// Although most fast enumerators are probably fused, this is not guaranteed
// by the documentation (actually, some may even use the last iteration as a
// cleanup step, and do it twice, which we should perhaps protect against?).
//
//
// # `DoubleEndedIterator`
//
// `NSArray` has `reverseObjectEnumerator`, but this is a separate object from
// `objectEnumerator`, and as such it is not possible to enumerate it in the
// backwards direction.
// Explicit lifetime bound on `C` to future-proof against possible soundness
// mistakes in the future.
//
// TODO: Consider adding `#[repr(C)]`, to allow LLVM to perform better
// zero-initialization.
#[derive(Debug, PartialEq)]
pub(crate) struct IterUnchecked<'a, C: ?Sized + 'a> {
helper: FastEnumeratorHelper,
/// 'a and C are covariant.
collection: &'a C,
#[cfg(debug_assertions)]
mutations_state: MutationState,
}
impl<'a, C: ?Sized + FastEnumerationHelper> IterUnchecked<'a, C> {
pub(crate) fn new(collection: &'a C) -> Self {
Self {
helper: FastEnumeratorHelper::new(),
collection,
#[cfg(debug_assertions)]
mutations_state: None,
}
}
}
impl<'a, C: FastEnumerationHelper> Iterator for IterUnchecked<'a, C> {
type Item = &'a C::Item;
#[inline]
#[track_caller]
fn next(&mut self) -> Option<&'a C::Item> {
#[cfg(debug_assertions)]
let mutations_state = Some(&mut self.mutations_state);
#[cfg(not(debug_assertions))]
let mutations_state = None;
// SAFETY: The collection is the same on each iteration.
//
// Passing the collection as immutable here is safe since it isn't
// mutated itself, and it is `UnsafeCell` anyhow.
let obj = unsafe {
self.helper
.next_from(ProtocolObject::from_ref(self.collection), mutations_state)?
};
// SAFETY: The lifetime is bound to the collection, and the type is
// correct.
//
// [Enumeration variables are externally retained][enum-retained], so
// it is okay to return a reference instead of retaining here.
//
Some(unsafe { obj.cast::<C::Item>().as_ref() })
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.helper.remaining_items_at_least(),
self.collection.maybe_len(),
)
}
}
#[derive(Debug, PartialEq)]
pub(crate) struct Iter<'a, C: ?Sized + 'a> {
helper: FastEnumeratorHelper,
/// 'a and C are covariant.
collection: &'a C,
mutations_state: MutationState,
}
impl<'a, C: ?Sized + FastEnumerationHelper> Iter<'a, C> {
pub(crate) fn new(collection: &'a C) -> Self {
Self {
helper: FastEnumeratorHelper::new(),
collection,
mutations_state: None,
}
}
}
impl<C: FastEnumerationHelper> Iterator for Iter<'_, C> {
type Item = Retained<C::Item>;
#[inline]
#[track_caller]
fn next(&mut self) -> Option<Retained<C::Item>> {
let obj = unsafe {
self.helper.next_from(
ProtocolObject::from_ref(self.collection),
Some(&mut self.mutations_state),
)?
};
Some(unsafe { obj.cast::<C::Item>().as_ref() }.retain())
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.helper.remaining_items_at_least(),
self.collection.maybe_len(),
)
}
}
// Note: I first considered having `IntoIter<T>` instead of `IntoIter<C>`, but
// that would be unsound for `NSDictionary<K, V>`, since if the `V` has a
// lifetime, that lifetime would have been erased (so we'd have to add an
// `V: 'static` bound, which is easy to forget).
//
// Also, this design allows us to have proper `Send + Sync` implementations,
// as well as allowing us to implement `FusedIterator` and `ExactSizeIterator`
// when possible.
#[derive(Debug, PartialEq)]
pub(crate) struct IntoIter<C: ?Sized> {
helper: FastEnumeratorHelper,
/// C is covariant.
collection: Retained<C>,
mutations_state: MutationState,
}
impl<C: ?Sized + FastEnumerationHelper> IntoIter<C> {
pub(crate) fn new(collection: Retained<C>) -> Self {
Self {
helper: FastEnumeratorHelper::new(),
collection,
mutations_state: None,
}
}
pub(crate) fn new_mutable<T>(collection: Retained<T>) -> Self
where
T: ClassType<Super = C>,
C: Sized,
{
Self {
helper: FastEnumeratorHelper::new(),
// SAFETY: Same as `Retained::into_super`, except we avoid the
// `'static` bounds, which aren't needed because the superclass
// carries the same generics.
collection: unsafe { Retained::cast_unchecked(collection) },
mutations_state: None,
}
}
}
impl<C: FastEnumerationHelper> Iterator for IntoIter<C> {
type Item = Retained<C::Item>;
#[inline]
#[track_caller]
fn next(&mut self) -> Option<Retained<C::Item>> {
let collection = ProtocolObject::from_ref(&*self.collection);
// SAFETY: The collection is the same on each iteration.
let obj = unsafe {
self.helper
.next_from(collection, Some(&mut self.mutations_state))?
};
// SAFETY: TODO
let obj = unsafe { Retained::retain(obj.cast::<C::Item>().as_ptr()) };
// SAFETY: The object was `NonNull`, and `retain` returns the same
// pointer.
Some(unsafe { obj.unwrap_unchecked() })
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.helper.remaining_items_at_least(),
self.collection.maybe_len(),
)
}
}
#[derive(Debug, PartialEq)]
pub(crate) struct IterUncheckedWithBackingEnum<'a, C: ?Sized + 'a, E: ?Sized + 'a> {
helper: FastEnumeratorHelper,
/// 'a and C are covariant.
collection: &'a C,
/// E is covariant.
enumerator: Retained<E>,
#[cfg(debug_assertions)]
mutations_state: MutationState,
}
impl<'a, C, E> IterUncheckedWithBackingEnum<'a, C, E>
where
C: ?Sized + FastEnumerationHelper,
E: ?Sized + FastEnumerationHelper,
{
pub(crate) unsafe fn new(collection: &'a C, enumerator: Retained<E>) -> Self {
Self {
helper: FastEnumeratorHelper::new(),
collection,
enumerator,
#[cfg(debug_assertions)]
mutations_state: None,
}
}
}
impl<'a, C, E> Iterator for IterUncheckedWithBackingEnum<'a, C, E>
where
C: ?Sized + FastEnumerationHelper,
E: FastEnumerationHelper,
{
type Item = &'a E::Item;
#[inline]
#[track_caller]
fn next(&mut self) -> Option<&'a E::Item> {
#[cfg(debug_assertions)]
let mutations_state = Some(&mut self.mutations_state);
#[cfg(not(debug_assertions))]
let mutations_state = None;
// SAFETY: The enumerator is the same on each iteration.
let obj = unsafe {
self.helper
.next_from(ProtocolObject::from_ref(&*self.enumerator), mutations_state)?
};
// SAFETY: TODO
Some(unsafe { obj.cast::<E::Item>().as_ref() })
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.helper.remaining_items_at_least(),
// Assume that the length of the collection matches the length of
// the enumerator.
self.collection.maybe_len(),
)
}
}
#[derive(Debug, PartialEq)]
pub(crate) struct IterWithBackingEnum<'a, C: ?Sized + 'a, E: ?Sized + 'a> {
helper: FastEnumeratorHelper,
/// 'a and C are covariant.
collection: &'a C,
/// E is covariant.
enumerator: Retained<E>,
mutations_state: MutationState,
}
impl<'a, C, E> IterWithBackingEnum<'a, C, E>
where
C: ?Sized + FastEnumerationHelper,
E: ?Sized + FastEnumerationHelper,
{
pub(crate) unsafe fn new(collection: &'a C, enumerator: Retained<E>) -> Self {
Self {
helper: FastEnumeratorHelper::new(),
collection,
enumerator,
mutations_state: None,
}
}
}
impl<C, E> Iterator for IterWithBackingEnum<'_, C, E>
where
C: ?Sized + FastEnumerationHelper,
E: FastEnumerationHelper,
{
type Item = Retained<E::Item>;
#[inline]
#[track_caller]
fn next(&mut self) -> Option<Retained<E::Item>> {
// SAFETY: The enumerator is the same on each iteration.
let obj = unsafe {
self.helper.next_from(
ProtocolObject::from_ref(&*self.enumerator),
Some(&mut self.mutations_state),
)?
};
// SAFETY: TODO
Some(unsafe { obj.cast::<E::Item>().as_ref() }.retain())
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
(
self.helper.remaining_items_at_least(),
// Assume that the length of the collection matches the length of
// the enumerator.
self.collection.maybe_len(),
)
}
}
#[doc(hidden)]
macro_rules! __impl_iter {
(
impl<$($lifetime:lifetime, )? $t1:ident: $bound1:ident $(+ $bound1_b:ident)? $(, $t2:ident: $bound2:ident $(+ $bound2_b:ident)?)?> Iterator<Item = $item:ty> for $for:ty { ... }
) => {
impl<$($lifetime, )? $t1: $bound1 $(+ $bound1_b)? $(, $t2: $bound2 $(+ $bound2_b)?)?> Iterator for $for {
type Item = $item;
#[inline]
#[track_caller]
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
}
}
#[doc(hidden)]
macro_rules! __impl_into_iter {
() => {};
(
$(#[$m:meta])*
impl<$param:ident: Message> IntoIterator for &$ty:ident<$param2:ident> {
type IntoIter = $iter:ident<'_, $param3:ident>;
}
$($rest:tt)*
) => {
$(#[$m])*
impl<'a, $param: Message> IntoIterator for &'a $ty<$param2> {
type Item = Retained<$param3>;
type IntoIter = $iter<'a, $param3>;
#[inline]
fn into_iter(self) -> Self::IntoIter {
$iter($crate::iter::Iter::new(&self))
}
}
__impl_into_iter! {
$($rest)*
}
};
(
$(#[$m:meta])*
impl<$param:ident: Message> IntoIterator for Retained<$ty:ident<$param2:ident>> {
#[uses($new_fn:ident)]
type IntoIter = $into_iter:ident<$param3:ident>;
}
$($rest:tt)*
) => {
$(#[$m])*
impl<$param: Message> objc2::rc::RetainedIntoIterator for $ty<$param2> {
type Item = Retained<$param3>;
type IntoIter = $into_iter<$param3>;
#[inline]
fn retained_into_iter(this: Retained<Self>) -> Self::IntoIter {
$into_iter($crate::iter::IntoIter::$new_fn(this))
}
}
__impl_into_iter! {
$($rest)*
}
};
}
#[cfg(test)]
#[cfg(feature = "NSArray")]
#[cfg(feature = "NSValue")]
mod tests {
use alloc::vec::Vec;
use core::mem::size_of;
use super::*;
use crate::{NSArray, NSNumber};
#[test]
#[cfg_attr(
any(not(target_pointer_width = "64"), debug_assertions),
ignore = "assertions assume pointer-width of 64, and the size only really matter in release mode"
)]
fn test_enumerator_helper() {
// We should attempt to reduce these if possible
assert_eq!(size_of::<NSFastEnumerationState>(), 64);
assert_eq!(size_of::<FastEnumeratorHelper>(), 208);
assert_eq!(size_of::<IterUnchecked<'_, NSArray<NSNumber>>>(), 216);
assert_eq!(size_of::<Iter<'_, NSArray<NSNumber>>>(), 232);
assert_eq!(size_of::<IntoIter<NSArray<NSNumber>>>(), 232);
}
#[test]
fn test_enumerator() {
let vec: Vec<_> = (0..4).map(NSNumber::new_usize).collect();
let array = NSArray::from_retained_slice(&vec);
let enumerator = array.iter();
assert_eq!(enumerator.count(), 4);
let enumerator = array.iter();
assert!(enumerator.enumerate().all(|(i, obj)| obj.as_usize() == i));
}
#[test]
fn test_into_enumerator() {
let vec: Vec<_> = (0..4).map(NSNumber::new_usize).collect();
let array = NSArray::from_retained_slice(&vec);
let enumerator = array.into_iter();
assert!(enumerator.enumerate().all(|(i, obj)| obj.as_usize() == i));
}
}