diff --git a/compiler/rustc_next_trait_solver/src/lib.rs b/compiler/rustc_next_trait_solver/src/lib.rs index 57bbf8772d322..2059d4702248f 100644 --- a/compiler/rustc_next_trait_solver/src/lib.rs +++ b/compiler/rustc_next_trait_solver/src/lib.rs @@ -17,3 +17,4 @@ pub mod normalize; pub mod placeholder; pub mod resolve; pub mod solve; +mod vec_extractor; diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs index 6273506b8adc7..83e4c36e86e2b 100644 --- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs +++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs @@ -38,6 +38,7 @@ use crate::solve::{ NestedNormalizationGoals, NoSolution, QueryInput, QueryResult, Response, SucceededInErased, VisibleForLeakCheck, inspect, }; +use crate::vec_extractor::Extractor; mod probe; mod solver_region_constraints; @@ -1001,7 +1002,13 @@ where ) -> Result, NoSolutionOrRerunNonErased> { // If this loop did not result in any progress, what's our final certainty. let mut unchanged_certainty = Some(Certainty::Yes); - for (source, goal, stalled_on) in mem::take(&mut self.nested_goals) { + + let mut nested_goals = mem::take(&mut self.nested_goals); + let mut extractor = Extractor::new(&mut nested_goals); + + while let Some(((source, goal, stalled_on), hole)) = + extractor.entry().map(|entry| entry.take()) + { // We never handle `NormalizesTo` as a nested goal debug_assert!(!matches!( goal.predicate.kind().skip_binder(), @@ -1015,7 +1022,7 @@ where match certainty { Certainty::Yes => {} Certainty::Maybe { .. } => { - self.nested_goals.push((source, goal, None)); + hole.fill((source, goal, None)); unchanged_certainty = unchanged_certainty.map(|c| c.and(certainty)); } } @@ -1023,7 +1030,16 @@ where } let GoalEvaluation { goal, certainty, has_changed, stalled_on } = - self.evaluate_goal(source, goal, stalled_on)?; + match self.evaluate_goal(source, goal, stalled_on) { + Ok(it) => it, + Err(err) => { + extractor.drop_rest(); + drop(extractor); + self.nested_goals = nested_goals; + return Err(err); + } + }; + if has_changed == HasChanged::Yes { unchanged_certainty = None; } @@ -1031,12 +1047,15 @@ where match certainty { Certainty::Yes => {} Certainty::Maybe { .. } => { - self.nested_goals.push((source, goal, stalled_on)); + hole.fill((source, goal, stalled_on)); unchanged_certainty = unchanged_certainty.map(|c| c.and(certainty)); } } } + drop(extractor); + self.nested_goals = nested_goals; + Ok(unchanged_certainty) } diff --git a/compiler/rustc_next_trait_solver/src/vec_extractor.rs b/compiler/rustc_next_trait_solver/src/vec_extractor.rs new file mode 100644 index 0000000000000..601154a579930 --- /dev/null +++ b/compiler/rustc_next_trait_solver/src/vec_extractor.rs @@ -0,0 +1,179 @@ +use std::marker::PhantomData; +use std::mem::ManuallyDrop; +use std::ptr; +use std::ptr::NonNull; + +use crate::vec_extractor::raw_slice::RawSliceIter; + +mod raw_slice; + +/// A cursor-like API for vec; generalization of `Vec::drain` and `Vec::extract_if`. +pub struct Extractor<'e, T: 'e> { + // Invariants: + // - `vec` points to a vector with length set to 0, and is valid for reads & + // writes + // - `rest` and `kept_end` point into `*vec`'s allocation + // - the range between `vec.as_ref().as_ptr()` and `kept_end` is intialized + // - `kept_end <= rest.as_ptr()` + // + // Essentially kept elements rest + // ________/ __________/ + // / \ / \ + // vec allocation: [ ] + // ^ ^\_____________/^ ^\____________/ + // Extractor { | | | | | | + // vec: --> Vec { | | extracted | | spare vec + // ptr -+ | elements | | capacity + // cap: N, | | | + // len: 0, | | | + // } | | | + // kept_end: ----------------+ | | + // | | + // rest: ------------------------------------+----------+ + // + // } + kept_end: NonNull, + rest: RawSliceIter<'e, T>, + vec: NonNull>, + _ghost: PhantomData<&'e mut Vec>, +} + +impl Drop for Extractor<'_, T> { + fn drop(&mut self) { + unsafe { + let vec = self.vec.as_mut(); + + let kept = self.kept_end.as_ptr().offset_from_unsigned(vec.as_ptr()); + let rest = self.rest.len(); + let len = kept + rest; + + ptr::copy(self.rest.as_slice().as_ptr().cast(), self.kept_end.as_ptr(), rest); + + vec.set_len(len); + } + } +} + +pub struct Entry<'e, 'a, T> { + // Invariants: + // - `val` points into `*extractor.vec`, is valid for reads & writes, + // does not overlap with `extractor.rest`, and is `>= extractor.kept_end` + extractor: &'a mut Extractor<'e, T>, + val: &'a mut T, +} + +impl Drop for Entry<'_, '_, T> { + fn drop(&mut self) { + unsafe { + let kept_end = self.extractor.kept_end; + ptr::copy(ptr::from_ref(self.val), kept_end.as_ptr(), 1); + self.extractor.kept_end = kept_end.add(1); + } + } +} + +pub struct Hole<'e, 'a, T> { + // Invariants: + // - `extractor.kept_end` does not overlap with `extractor.rest` + extractor: &'a mut Extractor<'e, T>, +} + +impl<'e, T> Extractor<'e, T> { + pub fn new(v: &'e mut Vec) -> Self { + let len = v.len(); + + unsafe { + // Safety: setting length to 0 is always sound + v.set_len(0); + + // Safety: this is equivalent to `v.as_slice()`, but before setting + // the length to 0 + // TODO + let slice = NonNull::slice_from_raw_parts(NonNull::new(v.as_mut_ptr()).unwrap(), len); + + // Safety: + // - `vec` points to a vector with length 0 and is valid + // - `rest` and `kept_end` point into `*vec`'s allocation + // - the range between `vec.as_ref().as_ptr()` and `kept_end` is + // empty (and this initialized) + // - `kept_end <= rest.as_ptr()` (they are equal) + Self { + kept_end: slice.cast(), + rest: RawSliceIter::new(slice), + vec: NonNull::from(v), + _ghost: PhantomData, + } + } + } + + pub fn entry(&mut self) -> Option> { + let val = self.rest.next()?; + + // Safety: `val` points into extractor's vec, since `rest` does. + // it is valid for reads & writes (as we got it from the vec), + // and does not overlap with `rest` (we just got it from there) + Some(Entry { extractor: self, val: unsafe { { val }.as_mut() } }) + } + + pub fn drop_rest(&mut self) { + unsafe { + std::ptr::drop_in_place(self.rest.as_slice().as_ptr()); + self.rest = <_>::default(); + } + } +} + +impl<'e, 'a, T> Entry<'e, 'a, T> { + pub fn take(self) -> (T, Hole<'e, 'a, T>) { + unsafe { + let (extractor, val) = self.into_raw_parts(); + let val = NonNull::from(val).read(); + (val, Hole { extractor }) + } + } + + pub fn keep(self) { /* drop is equivalent to keeping */ + } + + fn into_raw_parts(self) -> (&'a mut Extractor<'e, T>, &'a mut T) { + let this = ManuallyDrop::new(self); + + // Safety: moving out a field + unsafe { (ptr::from_ref(&this.extractor).read(), ptr::from_ref(&this.val).read()) } + } +} + +impl<'e, 'a, T> Hole<'e, 'a, T> { + pub fn fill(self, val: T) -> Entry<'e, 'a, T> { + unsafe { + let end = self.extractor.kept_end; + end.write(val); + + Entry { extractor: self.extractor, val: { end }.as_mut() } + } + } +} + +#[test] +fn smoke() { + let mut v = (0..100).into_iter().map(|x| (x, x.to_string())).collect::>(); + + let mut e = Extractor::new(&mut v); + + while let Some(entry) = e.entry() { + let (elem, hole) = entry.take(); + + if elem.0 > 50 { + break; + } + + if elem.0 % 3 == 0 { + hole.fill((1000, "meow".to_owned())); + } else if elem.0 % 3 == 1 { + hole.fill(elem); + } + } + + e.drop_rest(); + dbg!(v); +} diff --git a/compiler/rustc_next_trait_solver/src/vec_extractor/raw_slice.rs b/compiler/rustc_next_trait_solver/src/vec_extractor/raw_slice.rs new file mode 100644 index 0000000000000..20711cd2176f6 --- /dev/null +++ b/compiler/rustc_next_trait_solver/src/vec_extractor/raw_slice.rs @@ -0,0 +1,266 @@ +use std::iter::FusedIterator; +use std::marker::PhantomData; +use std::ptr::{self, NonNull}; + +/// A version of [`core::slice::Iter`] that works on raw pointers. +/// This should really be part of `core`... +#[must_use = "iterators are lazy and do nothing unless consumed"] +pub struct RawSliceIter<'a, T> { + /// The pointer to the next element to return, or the past-the-end location + /// if the iterator is empty. + /// + /// This address will be used for all ZST elements, never changed. + ptr: NonNull, + /// For non-ZSTs, the non-null pointer to the past-the-end element. + /// + /// For ZSTs, this is `ptr::without_provenance_mut(len)`. + end_or_len: *const T, + + /// + _of_the_opera: PhantomData<&'a [T]>, +} + +impl<'a, T> RawSliceIter<'a, T> { + /// Helper for checking if `T` is a `ZST`, similar to how `core` has `SizedTypeProperties`. + const ZST: bool = size_of::() == 0; + + /// Creates a new slice iter from a pointer to a slice + /// + /// # Safety + /// + /// The pointer and its [length][NonNull::len] must satisfy the safety guarantees of [`NonNull::add`]. + /// + /// In particular, if any of the following conditions are violated, the result is Undefined Behavior: + /// - The computed offset, `len * size_of::()` bytes, must not overflow `isize`. + /// - If the computed offset is non-zero, then self must be derived from a pointer to some allocation + /// and the entire memory range between self and the result must be in bounds of that allocation. + /// In particular, this range must not “wrap around” the edge of the address space. + /// + /// Allocations can never be larger than `isize::MAX` bytes, so if the computed offset stays in bounds + /// of the allocation, it is guaranteed to satisfy the first requirement. + /// + /// **Additionally** it is undefined behavior if any of the methods (even safe ones!) are called + /// after the allocation that `slice` points to has been deallocated. That is, the above safety + /// requirements needs to be upheld while `RawSliceIter` is in use. + /// + /// The lifetime of the `RawSliceIter` might help with that, although note that this method + /// does not bound the lifetime in anyway, so the caller is responsible to make sure it's correct. + #[inline] + pub const unsafe fn new(slice: NonNull<[T]>) -> Self { + let ptr: NonNull = slice.cast(); + let len = slice.len(); + + unsafe { + let end_or_len = + if Self::ZST { ptr::without_provenance(len) } else { ptr.as_ptr().add(len) }; + + Self { ptr, end_or_len, _of_the_opera: PhantomData } + } + } + + pub const fn from_ref(slice: &'a [T]) -> Self { + // This should just be `NonNull::from(slice)`, but it's not const stable :( + // + // Safety: + // + // - The pointer & length come from a slice, guaranteeing that everything is in bounds of + // an allocation + // - The input slice has the same lifetime as the output iterator, so it's guaranteed that + // the allocation will live long enough for any call to a method on an iterator + unsafe { Self::new(NonNull::new(ptr::from_ref(slice).cast_mut()).unwrap()) } + } + + pub const fn from_mut(slice: &'a mut [T]) -> Self { + // This should just be `NonNull::from(slice)`, but it's not const stable :( + // + // Safety: + // + // - The pointer & length come from a slice, guaranteeing that everything is in bounds of + // an allocation + // - The input slice has the same lifetime as the output iterator, so it's guaranteed that + // the allocation will live long enough for any call to a method on an iterator + unsafe { Self::new(NonNull::new(ptr::from_mut(slice)).unwrap()) } + } + + #[must_use] + #[inline] + pub fn as_slice(&self) -> NonNull<[T]> { + let len = unsafe { + if Self::ZST { + self.end_or_len.addr() + } else { + self.end_or_len.offset_from_unsigned(self.ptr.as_ptr()) + } + }; + + NonNull::slice_from_raw_parts(self.ptr, len) + } + + #[inline(always)] + fn zst( + &self, + zst: impl FnOnce(NonNull, usize) -> R, + non_zst: impl FnOnce(NonNull, NonNull) -> R, + ) -> R { + if Self::ZST { + zst(self.ptr, self.end_or_len.addr()) + } else { + non_zst(self.ptr, unsafe { NonNull::new_unchecked(self.end_or_len.cast_mut()) }) + } + } + + #[inline(always)] + fn zst_mut( + &mut self, + zst: impl FnOnce(&mut NonNull, &mut usize) -> R, + non_zst: impl FnOnce(&mut NonNull, &mut NonNull) -> R, + ) -> R { + if Self::ZST { + zst(&mut self.ptr, unsafe { &mut *ptr::from_mut(&mut self.end_or_len).cast::() }) + } else { + non_zst(&mut self.ptr, unsafe { + &mut *ptr::from_mut(&mut self.end_or_len).cast::>() + }) + } + } +} + +impl Clone for RawSliceIter<'_, T> { + #[inline] + fn clone(&self) -> Self { + RawSliceIter { ..*self } + } +} + +impl<'a, T> RawSliceIter<'a, T> { + /// Returns the last element and moves the end of the iterator backwards by 1. + /// + /// # Safety + /// + /// The iterator must not be empty + #[inline] + unsafe fn next_back_unchecked(&mut self) -> NonNull { + unsafe { self.pre_dec_end(1) } + } + + /// Helper function for moving the start of the iterator forwards by `offset` elements, + /// returning the old start. + /// Unsafe because the offset must not exceed `self.len()`. + #[inline(always)] + unsafe fn post_inc_start(&mut self, offset: usize) -> NonNull { + let old = self.ptr; + + // SAFETY: the caller guarantees that `offset` doesn't exceed `self.len()`, + // so this new pointer is inside `self` and thus guaranteed to be non-null. + unsafe { + self.zst_mut( + |_, len| *len = len.unchecked_sub(offset), + |start, _| *start = start.add(offset), + ); + } + + old + } + + /// Helper function for moving the end of the iterator backwards by `offset` elements, + /// returning the new end. + /// Unsafe because the offset must not exceed `self.len()`. + #[inline(always)] + unsafe fn pre_dec_end(&mut self, offset: usize) -> NonNull { + self.zst_mut( + |ptr, len| unsafe { + *len = len.unchecked_sub(offset); + *ptr + }, + |_, end| unsafe { + *end = end.sub(offset); + *end + }, + ) + } + + #[inline(always)] + pub fn is_empty(&self) -> bool { + self.zst(|_, len| len == 0, |start, end| start == end) + } +} + +impl ExactSizeIterator for RawSliceIter<'_, T> { + #[inline(always)] + fn len(&self) -> usize { + self.zst(|_, len| len, |start, end| unsafe { end.offset_from_unsigned(start) }) + } +} + +impl<'a, T> Iterator for RawSliceIter<'a, T> { + type Item = NonNull; + + #[inline] + fn next(&mut self) -> Option> { + // intentionally not using the helpers because this is + // one of the most mono'd things in the library. + + let ptr = self.ptr; + let end_or_len = self.end_or_len; + + // SAFETY: See inner comments. (For some reason having multiple + // block breaks inlining this -- if you can fix that please do!) + unsafe { + if Self::ZST { + let len = end_or_len.addr(); + if len == 0 { + return None; + } + // SAFETY: just checked that it's not zero, so subtracting one + // cannot wrap. (Ideally this would be `checked_sub`, which + // does the same thing internally, but as of 2025-02 that + // doesn't optimize quite as small in MIR.) + self.end_or_len = ptr::without_provenance(len.unchecked_sub(1)); + } else { + // SAFETY: by type invariant, the `end_or_len` field is always + // non-null for a non-ZST pointee. (This transmute ensures we + // get `!nonnull` metadata on the load of the field.) + if ptr == core::mem::transmute::<*const T, NonNull>(end_or_len) { + return None; + } + // SAFETY: since it's not empty, per the check above, moving + // forward one keeps us inside the slice, and this is valid. + self.ptr = ptr.add(1); + } + + Some(ptr) + } + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let exact = self.len(); + (exact, Some(exact)) + } + + #[inline] + fn count(self) -> usize { + self.len() + } +} + +impl<'a, T> DoubleEndedIterator for RawSliceIter<'a, T> { + #[inline] + fn next_back(&mut self) -> Option> { + // SAFETY: The call to `next_back_unchecked` + // is safe since we check if the iterator is empty first. + #[expect(unstable_name_collisions)] + unsafe { + if self.is_empty() { None } else { Some(self.next_back_unchecked()) } + } + } +} + +impl FusedIterator for RawSliceIter<'_, T> {} + +impl Default for RawSliceIter<'_, T> { + /// Creates an empty raw slice iterator. + fn default() -> Self { + Self::from_ref(&[]) + } +}