Skip to content

Added [T; N]::zip() #79451

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

Merged
merged 5 commits into from
Dec 22, 2020
Merged

Added [T; N]::zip() #79451

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3b8617b
Added [T; N]::zip()
usbalbin Nov 26, 2020
2f35fb1
Remove redundant tests
usbalbin Nov 27, 2020
be2c8f2
Update zip for better codegen, see discussion
usbalbin Dec 16, 2020
baa5e47
Update doc comment
usbalbin Dec 16, 2020
8b37259
Added reference to tracking issue
usbalbin Dec 16, 2020
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
28 changes: 28 additions & 0 deletions library/core/src/array/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -463,6 +463,34 @@ impl<T, const N: usize> [T; N] {
unsafe { crate::mem::transmute_copy::<_, [U; N]>(&dst) }
}

/// 'Zips up' two arrays into a single array of pairs.
/// `zip()` returns a new array where every element is a tuple where the first element comes from the first array, and the second element comes from the second array.
/// In other words, it zips two arrays together, into a single one.
///
/// # Examples
///
/// ```
/// #![feature(array_zip)]
/// let x = [1, 2, 3];
/// let y = [4, 5, 6];
/// let z = x.zip(y);
/// assert_eq!(z, [(1, 4), (2, 5), (3, 6)]);
/// ```
#[unstable(feature = "array_zip", issue = "none")]
pub fn zip<U>(self, rhs: [U; N]) -> [(T, U); N] {
use crate::mem::MaybeUninit;

let mut dst = MaybeUninit::uninit_array::<N>();
for ((lhs, rhs), dst) in IntoIter::new(self).zip(IntoIter::new(rhs)).zip(&mut dst) {
dst.write((lhs, rhs));
}
// FIXME: Convert to crate::mem::transmute once it works with generics.
// unsafe { crate::mem::transmute::<[MaybeUninit<U>; N], [U; N]>(dst) }
// SAFETY: At this point we've properly initialized the whole array
// and we just need to cast it to the correct type.
unsafe { crate::mem::transmute_copy::<_, [(T, U); N]>(&dst) }
Copy link
Contributor

@cynecx cynecx Nov 27, 2020 •
edited
Loading

Choose a reason for hiding this comment

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

This implementation (x86-64 codegen with -O) uses more stack space and includes branches/a loop compared to this one (which is branch-free and utilizes less stack, at least for smaller sizes):

let mut dst = MaybeUninit::<[(T, U); N]>::uninit();
let ptr = dst.as_mut_ptr() as *mut (T, U);
for (idx, (lhs, rhs)) in IntoIter::new(lhs).zip(IntoIter::new(rhs)).enumerate() {
    unsafe { ptr.add(idx).write((lhs, rhs)) }
}
unsafe { dst.assume_init() }

https://godbolt.org/z/T4aT78

Copy link
Contributor

@cynecx cynecx Nov 27, 2020 •
edited
Loading

Choose a reason for hiding this comment

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

However, that codegen is still not optimal (the memcpy is basically redundant and N/RVO could definitely help here).

Copy link
Contributor

@cynecx cynecx Nov 27, 2020 •
edited
Loading

Choose a reason for hiding this comment

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

Are [MaybeUninit<(T, U)>; N] and [(T, U); N] layout compatible/equivalent? Otherwise that transmute_copy would be UB.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I am quite new to working with unsafe but the doc states

[...] MaybeUninit will always guarantee that it has the same size, alignment, and ABI as T [...]

and there are examples with [Vec<u32>; N]) and [u8; N]. However I am not sure if there is anything special with [(T, U); N] that would make it not work.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Looking at map() from #75212, which I based my function on, uses unsafe { crate::mem::transmute_copy::<_, [U; N]>(&dst) } for any U. That would make the same for [(T, U); N] (not the same U) fine, right?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Sorry for the delay. I added a zippy2 to @cynecx's godbolt link here. However I am not too familiar with judging asm quality :)

Copy link
Contributor

@JulianKnodt JulianKnodt Dec 14, 2020 •
edited
Loading

Choose a reason for hiding this comment

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

ah ty for the update! The asm just helps me see what differences there are between the two versions, mainly looking for bounds checks which shouldn't be necessary here. The version you have has a loop in the LBB1 block, whereas cynecx's has no loop and is unrolled as he mentions in his top comment. If you increase the size of the arrays, you'll start to see more differences between the two.

The one thing is, if you change

for ((lhs, rhs), dst) in IntoIter::new(lhs).zip(IntoIter::new(rhs)).zip(&mut dst) {
  dst.write((lhs, rhs));
}

into

for (i, (lhs, rhs)) in IntoIter::new(lhs).zip(IntoIter::new(rhs)).enumerate() {
  dst[i].write((lhs, rhs));
}

both versions output the same asm. So the main difference between the two is here if you choose to zip or enumerate, rather than whether you use a pointer or a MaybeUninit

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Then would that be the preferred implementation (zip3 on godbolt)? Seems to me like best of both worlds, no extra unsafe yet the same asm?

Copy link
Contributor

@JulianKnodt JulianKnodt Dec 15, 2020 •
edited
Loading

Choose a reason for hiding this comment

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

To me that seems good, but I think that's ultimately up to you & the reviewer! I was just curious how each impl compared to the others.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Do you all think zip3 is the way to go? Or do you have any other suggestions? :)

}

/// Returns a slice containing the entire array. Equivalent to `&s[..]`.
#[unstable(feature = "array_methods", issue = "76118")]
pub fn as_slice(&self) -> &[T] {
Expand Down
8 changes: 8 additions & 0 deletions library/core/tests/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,14 @@ fn array_map() {
assert_eq!(b, [1, 2, 3]);
}

#[test]
fn array_zip() {
let a = [1, 2, 3];
let b = [4, 5, 6];
let c = a.zip(b);
assert_eq!(c, [(1, 4), (2, 5), (3, 6)]);
}

// See note on above test for why `should_panic` is used.
#[test]
#[should_panic(expected = "test succeeded")]
Expand Down
1 change: 1 addition & 0 deletions library/core/tests/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#![feature(array_from_ref)]
#![feature(array_methods)]
#![feature(array_map)]
#![feature(array_zip)]
#![feature(array_windows)]
#![feature(bool_to_option)]
#![feature(bound_cloned)]
Expand Down