Compiler Explorer

Source code

#![feature(inline_const)]
#![feature(const_mut_refs)]
#![feature(const_ptr_write)]

use std::mem::*;
use std::ptr;

pub struct Msg(String);

#[no_mangle]
pub fn swap_with_sep(a: &mut Msg, b: &mut Msg) {
    swap_chunked(a, b);
}

#[no_mangle]
pub fn swap_with_arrs(a: &mut Msg, b: &mut Msg) {
    swap_with_arrays(a, b);
}

#[inline]
const fn swap_chunked<T: Sized>(x: &mut T, y: &mut T) {
    const XMM_BYTES: usize = 128 / 8;
    const YMM_BYTES: usize = 256 / 8;

#[inline]
    const unsafe fn swap_simd<const SIZE_BYTES: usize>(
        x: *mut MaybeUninit<u64>,
        y: *mut MaybeUninit<u64>,
    ) {
        const {
            assert!(
                SIZE_BYTES == XMM_BYTES || SIZE_BYTES == YMM_BYTES,
                "Must have valid SIMD register size",
            );
        }
        unsafe {
            const {
                assert!(XMM_BYTES == size_of::<u64>() * 2, "Check number of temporaries below.");
                assert!(YMM_BYTES == size_of::<u64>() * 4, "Check number of temporaries below.");
            }
            let x0: MaybeUninit<u64> = ptr::read_unaligned(x);
            let x1: MaybeUninit<u64> = ptr::read_unaligned(x.add(1));
            let x2: MaybeUninit<u64>;
            let x3: MaybeUninit<u64>;
            if const { SIZE_BYTES == YMM_BYTES } {
                x2 = ptr::read_unaligned(x.add(2));
                x3 = ptr::read_unaligned(x.add(3));
            } else {
                x2 = MaybeUninit::uninit();
                x3 = MaybeUninit::uninit();
            }

ptr::copy_nonoverlapping::<MaybeUninit<u8>>(y.cast(), x.cast(), SIZE_BYTES);

ptr::write_unaligned(y, x0);
            ptr::write_unaligned(y.add(1), x1);
            if const { SIZE_BYTES == YMM_BYTES } {
                ptr::write_unaligned(y.add(2), x2);
                ptr::write_unaligned(y.add(3), x3);
            }
        }
    }

unsafe {
            let x = x.add(offset_before_tail);
            let y = y.add(offset_before_tail);
            let tmp_x: MaybeUninit<ChunkType> = ptr::read_unaligned(x.cast());
            let tmp_y: MaybeUninit<ChunkType> = ptr::read_unaligned(y.cast());
            ptr::write_unaligned(x.cast(), tmp_y);
            ptr::write_unaligned(y.cast(), tmp_x);
        }
    }

const {
        assert!(size_of::<T>() <= usize::MAX / 2, "We assume that overflows cannot happen.");
    }

let x: *mut MaybeUninit<u8> = (x as *mut T).cast();
    let y: *mut MaybeUninit<u8> = (y as *mut T).cast();

let mut byte_offset = 0;
    while byte_offset + YMM_BYTES <= const { size_of::<T>() } {
        unsafe {
            let x = x.add(byte_offset);
            let y = y.add(byte_offset);
            swap_simd::<YMM_BYTES>(x.cast(), y.cast());
            byte_offset += YMM_BYTES;
        }
    }
    if byte_offset + XMM_BYTES <= const { size_of::<T>() } {
        unsafe {
            let x = x.add(byte_offset);
            let y = y.add(byte_offset);
            swap_simd::<XMM_BYTES>(x.cast(), y.cast());
            byte_offset += XMM_BYTES;
        }
    }
    debug_assert!(byte_offset + XMM_BYTES > size_of::<T>());
    unsafe {
        const {
            assert!(size_of::<T>() % XMM_BYTES < 2 * size_of::<u64>());
        }
        swap_tail::<u64, T>(x, y);
        swap_tail::<u32, T>(x, y);
        swap_tail::<u16, T>(x, y);
        swap_tail::<u8, T>(x, y);
    }
}

#[inline]
const fn swap_with_arrays<T: Sized>(x: &mut T, y: &mut T) {
    const XMM_BYTES: usize = 128 / 8;
    const YMM_BYTES: usize = 256 / 8;

const {
                assert!(XMM_BYTES == size_of::<u64>() * 2, "Check number of temporaries below.");
                assert!(YMM_BYTES == size_of::<u64>() * 4, "Check number of temporaries below.");
            }

let mut tmp_x: [MaybeUninit<u8>; 16] = MaybeUninit::uninit().assume_init();
            let mut tmp_y: [MaybeUninit<u8>; 16] = MaybeUninit::uninit().assume_init();
            let mut tmp_x = tmp_x.as_mut_ptr();
            let mut tmp_y = tmp_y.as_mut_ptr();

ptr::copy_nonoverlapping(x.cast(), tmp_x, SIZE_BYTES);
            ptr::copy_nonoverlapping(y.cast(), tmp_y, SIZE_BYTES);
            ptr::copy_nonoverlapping(tmp_x, y.cast(), SIZE_BYTES);
            ptr::copy_nonoverlapping(tmp_y, x.cast(), SIZE_BYTES);
        }
    }

#[inline]
    const unsafe fn swap_tail<ChunkType: Copy, T>(
        x: *mut MaybeUninit<u8>,
        y: *mut MaybeUninit<u8>,
    ) {
        let remaining_bytes = const { size_of::<T>() % (2 * size_of::<ChunkType>()) };
        if remaining_bytes < const { size_of::<ChunkType>() } {
            return;
        }
        let offset_before_tail = size_of::<T>() - remaining_bytes;
        unsafe {
            let x = x.add(offset_before_tail);
            let y = y.add(offset_before_tail);
            let tmp_x: MaybeUninit<ChunkType> = ptr::read_unaligned(x.cast());
            let tmp_y: MaybeUninit<ChunkType> = ptr::read_unaligned(y.cast());
            ptr::write_unaligned(x.cast(), tmp_y);
            ptr::write_unaligned(y.cast(), tmp_x);
        }
    }

const {
        assert!(size_of::<T>() <= usize::MAX / 2, "We assume that overflows cannot happen.");
    }

let x: *mut MaybeUninit<u8> = (x as *mut T).cast();
    let y: *mut MaybeUninit<u8> = (y as *mut T).cast();