Compiler Explorer

Source code

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

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

#[no_mangle]
pub fn chun_swap_many_strings(a: &mut [String; 50], b: &mut [String; 50]){
    swap_chunked(a, b);
}

#[no_mangle]
pub fn chun_swap_strings(a: &mut String, b: &mut String){
    const {
        assert!(size_of::<String>() == 24);
        assert!(align_of::<String>() == 8);
    }
    swap_chunked(a, b);
}

pub struct Another{
    pub a: u16,
    pub b: u32,
    pub c: u16,
    pub d: u16
}

#[no_mangle]
pub fn chun_swap_another(a: &mut Another, b: &mut Another){
    const {
        assert!(size_of::<Another>() == 12);
        assert!(align_of::<Another>() == 4);
    }
    swap_chunked(a, b);
}

#[repr(C, packed)]
pub struct LargePacked {
    a: bool,
    b: u64,
    c: u32,
    d: u64,
    e: [bool;2],
    f: u16,
    g: (usize, usize),
}

#[no_mangle]
pub fn chun_swap_large_packed(a: &mut LargePacked, b: &mut LargePacked){
    const {
        assert!(size_of::<LargePacked>() == 41);
        assert!(align_of::<LargePacked>() == 1);
    }
    swap_chunked(a, b);
}

#[repr(C, packed)]
pub struct SmallPacked {
    a: bool,
    b: u64,
}

#[no_mangle]
pub fn chun_swap_small_packed(a: &mut SmallPacked, b: &mut SmallPacked){
    const {
        assert!(size_of::<SmallPacked>() == 9);
        assert!(align_of::<SmallPacked>() == 1);
    }
    swap_chunked(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 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);
            }
        }
    }

#[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);
        }
    }

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;
        }
    }
    while 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;
        }
    }

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);
    }
}