esp-storage: Fix incorrect usage of MaybeUninit

esp-storage/CHANGELOG.md

@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fix incorrect usage of MaybeUninit (#1902)
+
 ### Removed
 
 ## 0.3.0 - 2023-08-16

esp-storage/src/buffer.rs

@@ -0,0 +1,61 @@
+use core::{mem::MaybeUninit, slice};
+
+#[cfg(feature = "nor-flash")]
+pub fn uninit_slice(bytes: &[u8]) -> &[MaybeUninit<u8>] {
+    unsafe { core::mem::transmute(bytes) }
+}
+
+#[cfg(feature = "nor-flash")]
+pub fn uninit_slice_mut(bytes: &mut [u8]) -> &mut [MaybeUninit<u8>] {
+    unsafe { core::mem::transmute(bytes) }
+}
+
+pub type FlashWordBuffer = FlashBuffer<4, 1>;
+pub type FlashSectorBuffer = FlashBuffer<4096, 1024>;
+
+#[repr(C)]
+pub union FlashBuffer<const N: usize, const M: usize> {
+    bytes: [MaybeUninit<u8>; M],
+    words: [MaybeUninit<u32>; N],
+}
+
+impl<const N: usize, const M: usize> FlashBuffer<N, M> {
+    pub const fn uninit() -> Self {
+        debug_assert!(N * 4 == M);
+        Self {
+            words: [MaybeUninit::uninit(); N],
+        }
+    }
+
+    pub fn as_bytes(&self) -> &[MaybeUninit<u8>] {
+        unsafe { self.bytes.as_ref() }
+    }
+
+    pub fn as_bytes_mut(&mut self) -> &mut [MaybeUninit<u8>] {
+        unsafe { self.bytes.as_mut() }
+    }
+
+    pub fn as_words(&self) -> &[MaybeUninit<u32>] {
+        unsafe { self.words.as_ref() }
+    }
+
+    pub fn as_words_mut(&mut self) -> &mut [MaybeUninit<u32>] {
+        unsafe { self.words.as_mut() }
+    }
+
+    pub unsafe fn assume_init_bytes(&self) -> &[u8] {
+        slice::from_raw_parts(self.bytes.as_ptr() as *const u8, self.bytes.len())
+    }
+
+    pub unsafe fn assume_init_bytes_mut(&mut self) -> &mut [u8] {
+        slice::from_raw_parts_mut(self.bytes.as_mut_ptr() as *mut u8, self.bytes.len())
+    }
+
+    pub unsafe fn assume_init_words(&self) -> &[u32] {
+        slice::from_raw_parts(self.words.as_ptr() as *const u32, self.words.len())
+    }
+
+    pub unsafe fn assume_init_words_mut(&mut self) -> &mut [u32] {
+        slice::from_raw_parts_mut(self.words.as_mut_ptr() as *mut u32, self.words.len())
+    }
+}

esp-storage/src/common.rs

@@ -1,27 +1,6 @@
-use core::ops::{Deref, DerefMut};
+use core::mem::MaybeUninit;
 
-use crate::chip_specific;
-
-#[repr(C, align(4))]
-pub struct FlashSectorBuffer {
-    // NOTE: Ensure that no unaligned fields are added above `data` to maintain its required
-    // alignment
-    data: [u8; FlashStorage::SECTOR_SIZE as usize],
-}
-
-impl Deref for FlashSectorBuffer {
-    type Target = [u8; FlashStorage::SECTOR_SIZE as usize];
-
-    fn deref(&self) -> &Self::Target {
-        &self.data
-    }
-}
-
-impl DerefMut for FlashSectorBuffer {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.data
-    }
-}
+use crate::{buffer::*, chip_specific};
 
 #[derive(Debug)]
 #[non_exhaustive]
@@ -71,8 +50,9 @@ impl FlashStorage {
         #[cfg(any(feature = "esp32", feature = "esp32s2"))]
         const ADDR: u32 = 0x1000;
 
-        let mut buffer = [0u8; 8];
-        storage.internal_read(ADDR, &mut buffer).ok();
+        let mut buffer = FlashWordBuffer::uninit();
+        storage.internal_read(ADDR, buffer.as_bytes_mut()).unwrap();
+        let buffer = unsafe { buffer.assume_init_bytes() };
         let mb = match buffer[3] & 0xf0 {
             0x00 => 1,
             0x10 => 2,
@@ -115,11 +95,11 @@ impl FlashStorage {
     pub(crate) fn internal_read(
         &mut self,
         offset: u32,
-        bytes: &mut [u8],
+        bytes: &mut [MaybeUninit<u8>],
     ) -> Result<(), FlashStorageError> {
         check_rc(chip_specific::esp_rom_spiflash_read(
             offset,
-            bytes.as_ptr() as *mut u32,
+            bytes.as_mut_ptr() as *mut u32,
             bytes.len() as u32,
         ))
     }

esp-storage/src/lib.rs

@@ -30,7 +30,7 @@ mod chip_specific;
 mod common;
 
 #[cfg(any(feature = "storage", feature = "nor-flash"))]
-use common::FlashSectorBuffer;
+mod buffer;
 #[cfg(any(feature = "storage", feature = "nor-flash"))]
 pub use common::{FlashStorage, FlashStorageError};
 

esp-storage/src/nor_flash.rs

@@ -1,5 +1,3 @@
-use core::mem::MaybeUninit;
-
 use embedded_storage::nor_flash::{
     ErrorType,
     NorFlash,
@@ -8,31 +6,13 @@ use embedded_storage::nor_flash::{
     ReadNorFlash,
 };
 
-use crate::{FlashSectorBuffer, FlashStorage, FlashStorageError};
-
-#[cfg(feature = "bytewise-read")]
-#[repr(C, align(4))]
-struct FlashWordBuffer {
-    // NOTE: Ensure that no unaligned fields are added above `data` to maintain its required
-    // alignment
-    data: [u8; FlashStorage::WORD_SIZE as usize],
-}
-
 #[cfg(feature = "bytewise-read")]
-impl core::ops::Deref for FlashWordBuffer {
-    type Target = [u8; FlashStorage::WORD_SIZE as usize];
-
-    fn deref(&self) -> &Self::Target {
-        &self.data
-    }
-}
-
-#[cfg(feature = "bytewise-read")]
-impl core::ops::DerefMut for FlashWordBuffer {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.data
-    }
-}
+use crate::buffer::FlashWordBuffer;
+use crate::{
+    buffer::{uninit_slice, uninit_slice_mut, FlashSectorBuffer},
+    FlashStorage,
+    FlashStorageError,
+};
 
 impl FlashStorage {
     #[inline(always)]
@@ -71,13 +51,13 @@ impl ReadNorFlash for FlashStorage {
         let (offset, bytes) = {
             let byte_offset = (offset % Self::WORD_SIZE) as usize;
             if byte_offset > 0 {
-                let mut word_buffer = MaybeUninit::<FlashWordBuffer>::uninit();
-                let word_buffer = unsafe { word_buffer.assume_init_mut() };
+                let mut word_buffer = FlashWordBuffer::uninit();
 
                 let offset = offset - byte_offset as u32;
-                let length = bytes.len().min(word_buffer.len() - byte_offset);
+                let length = bytes.len().min(Self::WORD_SIZE as usize - byte_offset);
 
-                self.internal_read(offset, &mut word_buffer[..])?;
+                self.internal_read(offset, word_buffer.as_bytes_mut())?;
+                let word_buffer = unsafe { word_buffer.assume_init_bytes_mut() };
                 bytes[..length].copy_from_slice(&word_buffer[byte_offset..][..length]);
 
                 (offset + Self::WORD_SIZE, &mut bytes[length..])
@@ -94,38 +74,37 @@ impl ReadNorFlash for FlashStorage {
             {
                 // Chunk already is word aligned so we can read directly to it
                 #[cfg(not(feature = "bytewise-read"))]
-                self.internal_read(offset, chunk)?;
+                self.internal_read(offset, uninit_slice_mut(chunk))?;
 
                 #[cfg(feature = "bytewise-read")]
                 {
                     let length = chunk.len();
                     let byte_length = length % Self::WORD_SIZE as usize;
                     let length = length - byte_length;
 
-                    self.internal_read(offset, &mut chunk[..length])?;
+                    self.internal_read(offset, &mut uninit_slice_mut(chunk)[..length])?;
 
                     // Read not aligned rest of data
                     if byte_length > 0 {
-                        let mut word_buffer = MaybeUninit::<FlashWordBuffer>::uninit();
-                        let word_buffer = unsafe { word_buffer.assume_init_mut() };
+                        let mut word_buffer = FlashWordBuffer::uninit();
 
-                        self.internal_read(offset + length as u32, &mut word_buffer[..])?;
+                        self.internal_read(offset + length as u32, word_buffer.as_bytes_mut())?;
+                        let word_buffer = unsafe { word_buffer.assume_init_bytes_mut() };
                         chunk[length..].copy_from_slice(&word_buffer[..byte_length]);
                     }
                 }
             }
         } else {
             // Bytes buffer isn't word-aligned so we might read only via aligned buffer
-            let mut buffer = MaybeUninit::<FlashSectorBuffer>::uninit();
-            let buffer = unsafe { buffer.assume_init_mut() };
+            let mut buffer = FlashSectorBuffer::uninit();
 
             for (offset, chunk) in (offset..)
                 .step_by(Self::SECTOR_SIZE as _)
                 .zip(bytes.chunks_mut(Self::SECTOR_SIZE as _))
             {
                 // Read to temporary buffer first (chunk length is aligned)
                 #[cfg(not(feature = "bytewise-read"))]
-                self.internal_read(offset, &mut buffer[..chunk.len()])?;
+                self.internal_read(offset, &mut buffer.as_bytes_mut()[..chunk.len()])?;
 
                 // Read to temporary buffer first (chunk length is not aligned)
                 #[cfg(feature = "bytewise-read")]
@@ -138,9 +117,11 @@ impl ReadNorFlash for FlashStorage {
                         length
                     };
 
-                    self.internal_read(offset, &mut buffer[..length])?;
+                    self.internal_read(offset, &mut buffer.as_bytes_mut()[..length])?;
                 }
 
+                let buffer = unsafe { buffer.assume_init_bytes() };
+
                 // Copy to bytes buffer
                 chunk.copy_from_slice(&buffer[..chunk.len()]);
             }
@@ -173,17 +154,16 @@ impl NorFlash for FlashStorage {
             }
         } else {
             // Bytes buffer isn't word-aligned so we might write only via aligned buffer
-            let mut buffer = MaybeUninit::<FlashSectorBuffer>::uninit();
-            let buffer = unsafe { buffer.assume_init_mut() };
+            let mut buffer = FlashSectorBuffer::uninit();
 
             for (offset, chunk) in (offset..)
                 .step_by(Self::SECTOR_SIZE as _)
                 .zip(bytes.chunks(Self::SECTOR_SIZE as _))
             {
                 // Copy to temporary buffer first
-                buffer[..chunk.len()].copy_from_slice(chunk);
+                buffer.as_bytes_mut()[..chunk.len()].copy_from_slice(uninit_slice(chunk));
                 // Write from temporary buffer
-                self.internal_write(offset, &buffer[..chunk.len()])?;
+                self.internal_write(offset, chunk)?;
             }
         }
 

esp-storage/src/storage.rs

@@ -1,8 +1,8 @@
-use core::mem::MaybeUninit;
+use core::mem::{self, MaybeUninit};
 
 use embedded_storage::{ReadStorage, Storage};
 
-use crate::{FlashSectorBuffer, FlashStorage, FlashStorageError};
+use crate::{buffer::FlashSectorBuffer, FlashStorage, FlashStorageError};
 
 impl ReadStorage for FlashStorage {
     type Error = FlashStorageError;
@@ -14,8 +14,7 @@ impl ReadStorage for FlashStorage {
         let mut aligned_offset = offset - data_offset;
 
         // Bypass clearing sector buffer for performance reasons
-        let mut sector_data = MaybeUninit::<FlashSectorBuffer>::uninit();
-        let sector_data = unsafe { sector_data.assume_init_mut() };
+        let mut sector_data = FlashSectorBuffer::uninit();
 
         while !bytes.is_empty() {
             let len = bytes.len().min((Self::SECTOR_SIZE - data_offset) as _);
@@ -24,8 +23,13 @@ impl ReadStorage for FlashStorage {
                 & !(Self::WORD_SIZE - 1) as usize;
 
             // Read only needed data words
-            self.internal_read(aligned_offset, &mut sector_data[..aligned_end])?;
+            self.internal_read(
+                aligned_offset,
+                &mut sector_data.as_bytes_mut()[..aligned_end],
+            )?;
 
+            let sector_data = &sector_data.as_bytes()[..aligned_end];
+            let sector_data = unsafe { mem::transmute::<&[MaybeUninit<u8>], &[u8]>(sector_data) };
             bytes[..len].copy_from_slice(&sector_data[data_offset as usize..][..len]);
 
             aligned_offset += Self::SECTOR_SIZE;
@@ -52,17 +56,17 @@ impl Storage for FlashStorage {
         let mut aligned_offset = offset - data_offset;
 
         // Bypass clearing sector buffer for performance reasons
-        let mut sector_data = MaybeUninit::<FlashSectorBuffer>::uninit();
-        let sector_data = unsafe { sector_data.assume_init_mut() };
+        let mut sector_data = FlashSectorBuffer::uninit();
 
         while !bytes.is_empty() {
-            self.internal_read(aligned_offset, &mut sector_data[..])?;
+            self.internal_read(aligned_offset, sector_data.as_bytes_mut())?;
+            let sector_data = unsafe { sector_data.assume_init_bytes_mut() };
 
             let len = bytes.len().min((Self::SECTOR_SIZE - data_offset) as _);
 
             sector_data[data_offset as usize..][..len].copy_from_slice(&bytes[..len]);
             self.internal_erase(aligned_offset / Self::SECTOR_SIZE)?;
-            self.internal_write(aligned_offset, &sector_data[..])?;
+            self.internal_write(aligned_offset, sector_data)?;
 
             aligned_offset += Self::SECTOR_SIZE;
             data_offset = 0;