Manipulating slice through &mut parameter not optimized well

[hanna-kruppe]

When passing a &mut &[u8] to a function that trims elements from the start/end of the slice, the code contains pointless length checks and stack traffic. Writing the function slightly differently, or inlining it, does produce the code I'd expect.

This is a minimal example (playground of all following snippets):

pub fn trim_in_place(a: &mut &[u8]) {
    while a.first() == Some(&42) {
        *a = &a[1..];
    }
}

I expect this to compile to a tight loop that simply decrements the slice length and increments the data pointer in registers, until the length reaches 0 or a byte value of 42 is encountered.

Instead, the loop writes to the stack on every iteration and checks twice whether the length is zero, trying to panic for the second check (presumably comes from the a[1..] slicing).

.LBB0_5:
    movzx   edx, byte ptr [rax - 1]
    cmp edx, 42
    jne .LBB0_8 # exit function
    cmp rcx, -1
    je  .LBB0_9 # panic, never taken
    mov qword ptr [rdi], rax
    mov qword ptr [rdi + 8], rcx
    dec rcx
    inc rax
    cmp rcx, -1
    jne .LBB0_5

The .LBB0_9 branch cannot ever be taken, and is indeed optimized out when writing the function in differently. The stack traffic also disappears, though that probably has an entirely different cause. This variant compiles to better code:

pub fn trim_functionally(mut a: &[u8]) -> &[u8] {
    while a.first() == Some(&42) {
        a = &a[1..];
    }
    a
}

It does all the loop calculations in registers. Also, marking the first function as inline(always) and defining the second function in terms of it ({ trim_in_place(&mut a); a }) optimizes to the same code.

.LBB1_2:
    movzx   eax, byte ptr [rdi]
    cmp eax, 42
    jne .LBB1_3 # exit function
    inc rdi
    dec rsi
    jne .LBB1_2

Both variants check if the slice is empty before entering the loop and return immediately in that case.

Background

I encountered this problem in the context of my dec2flt work, in a larger and more complicated function that trims two slices on both ends, and the slices are stored in a struct. Applying inline(always) to that function gives a small but measurable performance improvement across the board, but increases code size by 10 KiB (Edit: I just didn't measure properly. Size is about the same.). I haven't checked the assembly code (far too much code) but I suspect this means this problem, and solution, carries over to real code.

Meta

Reproduces on nightly playground. Stable is worse for all variants, presumably because it doesn't have #26411 (yay, progress!). My local rustc has the same problem. The exact commit is meaningless because it only exists in my rust.git fork. It's based f9f580936d81ed527115ae86375f69bb77723b1c.

[eefriedman]

The first example is much harder for LLVM to optimize because each iteration of the loop stores to a value which would be globally visible if the loop panics. I think LLVM would figure it out with the right set of optimization passes, though.

The other examples are much easier to optimize because LLVM can trivially eliminate the stack traffic using scalarrepl.

[hanna-kruppe]

According to the LLVM IR, the &mut &[u8]parameter is noalias and it's not passed to the panic function as argument. Rust-semantics wise I also don't see how the slice would escape in case of a panic. I'm not an expert on either, but I don't see how LLVM could consider it to be escaping.

[eefriedman]

The slice "escapes" like so:

use std::sync::Mutex;
use std::sync::Arc;

pub fn trim_in_place(a: &mut &[u8]) {
    while a.first() == Some(&42) {
        *a = &a[2..];
    }
}

fn main() {
    static X: &'static [u8] = &[42, 0, 42];
    let m = Arc::new(Mutex::new(X));
    let m2 = m.clone();
    let _ = std::thread::spawn(move || {
        let mut r = m.lock().unwrap();
        trim_in_place(&mut *r)
    }).join();
    let r = match m2.lock() {
        Ok(r) => r,
        Err(r) => r.into_inner()
    };
    assert_eq!(*r, &[42]);
}

[dotdash]

The mentioned transformation that leads to better code is what @reinerp suggested here: #26494 (comment)

[dotdash]

I'm preparing a PR that will eliminate the extra null check. For the loop itself, I don't have a fix for rustc (yet?), but this optimizes better:

pub fn trim_in_place(a: &mut &[u8]) {
    let mut x = *a;
    while x.first() == Some(&42) {
        x = &x[1..];
    }
    *a = x;
}

You get the tight loop with that, and with the PR I'm preparing, the result is this:

_ZN13trim_in_place20h5d93414c92587a1aeaaE:
    .cfi_startproc
    movq    (%rdi), %rax
    movq    8(%rdi), %rdx
    xorl    %ecx, %ecx
    testq   %rdx, %rdx
    je  .LBB0_5
    xorl    %ecx, %ecx
    .align  16, 0x90
.LBB0_2:
    movzbl  (%rax), %esi
    cmpl    $42, %esi
    jne .LBB0_3
    incq    %rax
    decq    %rdx
    jne .LBB0_2
    jmp .LBB0_5
.LBB0_3:
    movq    %rdx, %rcx
.LBB0_5:
    movq    %rax, (%rdi)
    movq    %rcx, 8(%rdi)
    retq

Still slightly weird WRT its usage of rcx, but at least the loop looks good.

[Aatch]

@dotdash yeah, LLVM register usage can be a little weird sometimes. It's live-range analysis can produce some odd results at times, meaning that it won't re-use a register even it can.