Skip to content

TypedArena: Eliminate intermediate copy if Iterator size is known #77945

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

TypedArena: Eliminate intermediate copy if Iterator size is known #77945

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 61 additions & 3 deletions compiler/rustc_arena/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,8 @@
#![feature(dropck_eyepatch)]
#![feature(new_uninit)]
#![feature(maybe_uninit_slice)]
#![feature(min_specialization)]
#![feature(trusted_len)]
#![cfg_attr(test, feature(test))]

use rustc_data_structures::cold_path;
Expand All @@ -22,6 +24,7 @@ use smallvec::SmallVec;
use std::alloc::Layout;
use std::cell::{Cell, RefCell};
use std::cmp;
use std::iter::TrustedLen;
use std::marker::{PhantomData, Send};
use std::mem::{self, MaybeUninit};
use std::ptr;
Expand Down Expand Up @@ -109,6 +112,53 @@ impl<T> Default for TypedArena<T> {
}
}

trait IterExt<I, T> {
fn write_to_arena(iter: I, arena: &TypedArena<T>) -> &mut [T];
}

impl<T, I> IterExt<I, T> for I
where
I: Iterator<Item = T>,
{
#[inline]
default fn write_to_arena(iter: I, arena: &TypedArena<T>) -> &mut [T] {
arena.alloc_from_iter_gen(iter)
}
}

impl<T, I> IterExt<I, T> for I
where
I: Iterator<Item = T> + TrustedLen,
{
#[inline]
fn write_to_arena(mut iter: I, arena: &TypedArena<T>) -> &mut [T] {
let size_hint = iter.size_hint();

match size_hint {
(0, Some(_)) => &mut [],
(len, Some(_)) => {
// no need to check min == max because of TrustedLen

// SAFETY: TrustedLen implementors must ensure they return exactly as many
// elements as they tell via `size_hint()`.
unsafe {
// We know the exact number of elements the iterator will produce here
let start_addr = arena.alloc_raw_slice(len);
let mut mem = start_addr;

while let Some(value) = iter.next() {
Copy link
Contributor

@Marwes Marwes Oct 16, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't this unsound if next panics? I assume the arena would try to call drop on uninitialized memory in that case (*).

(*) When I did this same optimization I added a drop guard which fills the slice with default values (**).
(**) Default can also panic but in that case the process abort, though I know in my case that Default won't panic for any of the types that use it.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Huh, thanks for noticing that. I guess this sort of explains why this optimization wasn't done before 😅

I'm trying to make this as general as possible, so no T: Default here :) However, this means the only way I see to solve this issue is to catch the unwind and undo the failed partial allocation by hand. Which could get messy.

I wonder if it's worth doing, considering the otherwise minimal impact of this PR.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The problem with that is that there may be more things allocated after this allocation which has already succeded and handed out a reference. So there isn't a way to undo this allocation as far as I can see (without adding another Vec to keep track of failed allocations I guess, I suppose that Vec wouldn't be used 99% of the time so it wouldn't add any overhead other than an increased arena size and complexity).

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't it the idea of TrustedLen that the iterator will always have that many elements? If next panics, it would violate that contract.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure which benefit TrustedLen would give people at all if next is allowed to panic.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@bugadani oh right. That's super weird though because then what's the benefit of TrustedLen over ExactSizeIterator?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a last attempt to salvage this PR, the dropping arena could get a Vec<Vec<T>> member that is used instead of the chunks when given an iterator. Then we can simply put the temporary vec that the pre-PR iterator code creates into that list and don't have to copy stuff into the chunk. We already have to pay the cost of deallocating the Vec anyways, all we do is move it into the future. This bloats up the structures of the native allocator and destroys cache locality, but might be worth a try.

Copy link
Contributor Author

@bugadani bugadani Oct 16, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@est31 I guess ExactSizeIterator isn't unsafe and so it's technically allowed to misbehave? I'm not sure.

Regarding the memory question:

TypedArena works in pages. When a page can't hold the requested data, a new page is created. The allocator then uses pointers in the last page. This complicates things a bit, but I don't think by a huge deal.

If an allocation panics, we can drop the memory that we allocated, as well as the memory that the inner operations allocated. We know about these regions, because we have them saved on our stack, or we can if currently something is missing.

If a nested tree of of alloc_from_iter calls fail for whatever reason, the panic will bubble up, and can be caught by each level of alloc call, basically unwinding their work.

The only question is, how do I bubble up the innermost panic? :)

Since Iterator isn't UnwindSafe, this looks like a lot more complicated than what should be done.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One thing that could work, with only a little overhead, is to allocate a new chunk for any allocation calls inside the iterator. I believe this would work without TrustedLen also.

Pseudo rust

let slice = self.alloc_raw_slice();
let chunk_index = self.current_chunk;
// Move to the next chunk, if no next chunk exists the next allocation will force a new chunk to be created
self.current_chunk += 1;
for elem in iterator {
    slice.fill(elem);
}
// Move back to the current chunk so we can try filling the rest of it
self.current_chunk = chunk_index;

// Search for the first chunk
fn get_current_chunk(&self, needed: usize) -> &mut Vec {
    for chunk in &mut self.chunks[self.current_chunk..] {
        if chunk.capacity()  - chunk.len() >= needed {
             self.current_chunk = index;
             return chunk
        }
    }
    self.allocate_chunk()
}

Copy link
Contributor Author

@bugadani bugadani Oct 16, 2020 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Huh, that sounds like an interesting idea! There is a theoretical possibility of wasting a ton of memory, but let's hope it's not an issue in practice :)

ptr::write(mem, value);
mem = mem.add(1);
}

return slice::from_raw_parts_mut(start_addr, len);
}
}
_ => panic!("Trying to allocate from an iterator with more than usize::MAX elements"),
}
}
}

impl<T> TypedArena<T> {
/// Allocates an object in the `TypedArena`, returning a reference to it.
#[inline]
Expand Down Expand Up @@ -183,10 +233,10 @@ impl<T> TypedArena<T> {
}
}

/// General case implementation for `alloc_from_iter()`.
#[inline]
pub fn alloc_from_iter<I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
assert!(mem::size_of::<T>() != 0);
let mut vec: SmallVec<[_; 8]> = iter.into_iter().collect();
fn alloc_from_iter_gen<I: Iterator<Item = T>>(&self, iter: I) -> &mut [T] {
let mut vec: SmallVec<[_; 8]> = iter.collect();
if vec.is_empty() {
return &mut [];
}
Expand All @@ -201,6 +251,14 @@ impl<T> TypedArena<T> {
}
}

#[inline]
pub fn alloc_from_iter<I: IntoIterator<Item = T>>(&self, iter: I) -> &mut [T] {
let iter = iter.into_iter();
assert!(mem::size_of::<T>() != 0);

I::IntoIter::write_to_arena(iter, self)
}

/// Grows the arena.
#[inline(never)]
#[cold]
Expand Down