Potential Undefined Behavior from Aliasing Violation

[lewismosciski]

Hello again,

The method as_mut_slice takes &self but returns &mut [f32]. This signature mismatch allows safe code to create multiple mutable references to the same data, leading to Undefined Behavior.

NNL/src/device/cpu.rs

Lines 256 to 260 in 6ae0f59

	pub fn as_mut_slice(&self) -> &mut [f32] {
	unsafe {
	std::slice::from_raw_parts_mut(
	self.ptr.as_mut_ptr(),
	self.size / std::mem::size_of::<f32>(),

POC:

use nnl::device::cpu::CpuMemory;

fn main() {
    let mem = CpuMemory::new(128).unwrap();

    // The incorrect `&self` signature allows creating two mutable references.
    let slice1 = mem.as_mut_slice();
    let slice2 = mem.as_mut_slice();

    // This is instant UB.
    slice1[0] = 10.0;
    slice2[0] = 20.0; 
    println!("Value: {}", slice1[0]);
}

Test with miri:

ccuu@ccuu:~/rust/rtest/nnl_test$ cargo +nightly miri run
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.28s
     Running `/home/ccuu/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/bin/cargo-miri runner target/miri/x86_64-unknown-linux-gnu/debug/nnl_test`
error: Undefined Behavior: attempting a write access using <330> at alloc220[0x0], but that tag does not exist in the borrow stack for this location
  --> src/main.rs:11:5
   |
11 |     slice1[0] = 10.0;
   |     ^^^^^^^^^^^^^^^^ this error occurs as part of an access at alloc220[0x0..0x4]
   |
   = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
   = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
help: <330> was created by a Unique retag at offsets [0x0..0x80]
  --> src/main.rs:7:18
   |
 7 |     let slice1 = mem.as_mut_slice();
   |                  ^^^^^^^^^^^^^^^^^^
help: <330> was later invalidated at offsets [0x0..0x80] by a Unique retag
  --> src/main.rs:8:18
   |
 8 |     let slice2 = mem.as_mut_slice();
   |                  ^^^^^^^^^^^^^^^^^^
   = note: BACKTRACE (of the first span):
   = note: inside `main` at src/main.rs:11:5: 11:21

The suggested fix is to change the signature to &mut self.

- pub fn as_mut_slice(&self) -> &mut [f32] {
+ pub fn as_mut_slice(&mut self) -> &mut [f32] {