Fast idiomatic Floyd-Warshall algorithm in Rust

At first blush, one would hope this would be enough:

fn floyd_warshall(dist: &mut [Vec<i32>]) {
    let n = dist.len();
    for i in 0..n {
        assert!(i < dist.len());
        for j in 0..n {
            assert!(j < dist.len());
            assert!(i < dist[j].len());
            let v2 = dist[j][i];
            for k in 0..n {
                assert!(k < dist[i].len());
                assert!(k < dist[j].len());
                dist[j][k] = min(dist[j][k], v2 + dist[i][k]);
            }
        }
    }
}

Adding asserts is a known trick to convince the Rust optimizer that variables are indeed in bounds. However, it doesn't work here. What we need to do is somehow make it even more obvious to the Rust compiler that these loops are in bounds without resorting to esoteric code.

To accomplish that, I moved to a 2D array as suggested by David Eisenstat:

fn floyd_warshall<const N:usize>(mut dist: Box<[[i32; N]; N]>) -> Box<[[i32; N]; N]> {
    for i in 0..N {
        for j in 0..N {
            for k in 0..N {
                dist[j][k] = min(dist[j][k], dist[j][i] + dist[i][k]);
            }
        }
    }
    dist
}

This uses constant generics, a relatively new feature of Rust, to specify the size of a given 2d array on the heap. On its own, this change does well on my machine (100ms faster than usafe, and ~20ms behind unsafe). If, additionally, you move the v2 calculation outside the k-loop like this:

fn floyd_warshall<const N:usize>(mut dist: Box<[[i32; N]; N]>) -> Box<[[i32; N]; N]> {
    for i in 0..N {
        for j in 0..N {
            let v2 = dist[j][i];
            for k in 0..N {
                dist[j][k] = min(dist[j][k], v2 + dist[i][k]);
            }
        }
    }
    dist
}

The improvement is substantial (from ~300ms to ~100ms on my machine). The same optimization works with floyd_warshall_unsafe bringing it to ~100ms on average on my machine. When inspecting the assembly (with #[inline(never)] on floyd_warshall) it doesn't look like bounds checks occur for either, and both look vectorized to some extent. Although, I am no expert at reading assembly.

Because this is such a hot loop (with up to three bounds checks) I'm not surprised performance suffers so much. Unfortunately, the usage of the indices in this case is complicated enough to prevent the assert trick from giving you an easy fix. There are other known cases where an assert check is desired to improve the performance but the compiler is unable to use the information sufficiently. Here is one such example.

Here is the playground with my changes.

After some experiments, based on ideas suggested in Andrew's answer, and comments in related issue I found solution, which:

• still uses the same interface (e.g. &mut [Vec<i32>] as argument)
• does't use unsafe
• 3-4x faster than unsafe version
• quite ugly :(

Code looks like this:

fn floyd_warshall_fast(dist: &mut [Vec<i32>]) {
    let n = dist.len();
    for i in 0..n {
        for j in 0..n {
            if i == j {
                continue;
            }
            let (dist_j, dist_i) = if j < i {
                let (lo, hi) = dist.split_at_mut(i);
                (&mut lo[j][..n], &mut hi[0][..n])
            } else {
                let (lo, hi) = dist.split_at_mut(j);
                (&mut hi[0][..n], &mut lo[i][..n])
            };
            let dist_ji = dist_j[i];
            for k in 0..n {
                dist_j[k] = min(dist_j[k], dist_ji + dist_i[k]);
            }
        }
    }
}

There are couple of ideas inside:

• We compute dist_ji once as it doesn't change inside the most inner cycle, and compiler doesn't need to think about it.
• We "prove" that dist[i] and dist[j] are actually two different vectors. This is done by this ugly split_at_mut thing and i == j special case (would really love to know an easier solution). After that we can treat dist[i] and dist[j] absolutely separately, and for example the compiler can vectorize this loop, because it knows that data doesn't overlap.
• Last trick is to "prove" to the compiler that both dist[i] and dist[j] have at least n elements. This is done by [..n] when computing dist[i] and dist[j] (e.g. we use &mut lo[j][..n] instead of just &mut lo[j]). After that, compiler understands that the inner loop never uses out of bounds values, and removes checks.

Interesting that only when all three optimizations are used, it gives a big speed up. If we use only any two of them, the compiler can't optimize it.