如何最优并行化修改同一Rust向量多个切片的代码？

我们需要原地翻倍向量中多个切片的每个元素，切片由(start, end)位置对列表定义。下面这段代码按常规思路编写，但因为在Rayon的并行for_each中对向量进行可变借用，无法通过编译：

use rayon::prelude::*;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 3), (4, 7), (8, 10)];

    slice_pairs.into_par_iter().for_each(|(start, end)| {
        let slice = &mut data[start..end];
        for elem in slice.iter_mut() {
            *elem *= 2;
        }
    });

    println!("{:?}", data);
}

问题的核心是Rust借用检查器无法确认并行操作的切片是否重叠，因此拒绝编译。但如果我们能确保所有切片完全不重叠，就可以安全地并行处理。

下面是一段用unsafe实现的代码，它通过将向量基指针转成i64再转回的方式绕过了借用检查，但这种实现不够优雅：

use rayon::prelude::*;
use std::mem;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 4), (4, 7), (7, 10)];

    let ptr_outer = data.as_mut_ptr();
    let ptr_int : i64 = unsafe { mem::transmute(ptr_outer) };

    slice_pairs.into_par_iter().for_each(|(start, end)| {
        unsafe {
            let ptr : *mut i32 = mem::transmute(ptr_int);
            let slice = std::slice::from_raw_parts_mut(ptr.add(start), end - start);

            for elem in slice.iter_mut() {
                *elem *= 2;
            }
        }
    });

    println!("{:?}", data);
}

更优的实现方案 #

方案一：安全API实现（推荐） #

如果能预先拆分出所有非重叠的可变切片，就可以完全借助Rust的安全API完成并行处理，不需要任何unsafe代码。借用检查器能直接验证切片的独占性：

use rayon::prelude::*;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 4), (4, 7), (7, 10)];

    // 预先拆分出所有非重叠的可变切片
    let mut slices = Vec::new();
    let mut remaining = &mut data[..];
    for &(start, end) in &slice_pairs {
        // 计算当前切片在剩余区间中的拆分位置
        let offset = start - (remaining.as_ptr() as usize - data.as_ptr() as usize);
        let (current_slice, rest) = remaining.split_at_mut(end - start);
        slices.push(current_slice);
        remaining = rest;
    }

    // 并行处理每个切片
    slices.into_par_iter().for_each(|slice| {
        slice.iter_mut().for_each(|elem| *elem *= 2);
    });

    println!("{:?}", data);
}

如果切片不是连续的但确实不重叠，可以先对切片按start排序，再用同样的方式拆分，确保每次拆分的切片都是独占的。

方案二：优化后的Unsafe实现 #

如果切片是任意非重叠的，无法预先用安全API拆分，可以使用更简洁的unsafe代码，去掉不必要的指针类型转换，同时添加前置检查避免未定义行为：

use rayon::prelude::*;

fn main() {
    let mut data = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
    let slice_pairs = vec![(0, 4), (4, 7), (7, 10)];

    let base_ptr = data.as_mut_ptr();
    let data_len = data.len();

    // 可选但推荐：验证所有切片合法且不重叠，避免UB
    slice_pairs.iter().for_each(|&(start, end)| {
        assert!(start <= end, "无效的切片区间");
        assert!(end <= data_len, "切片超出向量范围");
    });
    let mut sorted_pairs = slice_pairs.clone();
    sorted_pairs.sort_by_key(|&(s, _)| s);
    sorted_pairs.windows(2).for_each(|window| {
        let (_, prev_end) = window[0];
        let (curr_start, _) = window[1];
        assert!(curr_start >= prev_end, "切片存在重叠");
    });

    slice_pairs.into_par_iter().for_each(|(start, end)| {
        unsafe {
            // 直接基于基指针构造切片，无需多余转换
            let slice = std::slice::from_raw_parts_mut(base_ptr.add(start), end - start);
            slice.iter_mut().for_each(|elem| *elem *= 2);
        }
    });

    println!("{:?}", data);
}

这段代码保留了unsafe的灵活性，但去掉了原代码中冗余的mem::transmute操作，同时通过断言确保切片的合法性和非重叠性，最大限度降低了未定义行为的风险。

内容的提问来源于stack exchange，提问作者Yossi Kreinin