wasmi_ir2: make Decode trait safe (#1662)

* make Decode trait and impls inherently safe

An unsafe and more efficient implement of Decode and Decoder will have to wrap these safe APIs and avoid checks via APIs such as unwrap_unchecked etc. if needed.

* add module docs to decode module

Author: Robbepop

crates/ir2/build/display/decode.rs

@@ -191,7 +191,7 @@ impl<const N: usize> Display for DisplayDecode<&'_ GenericOp<N>> {
                 .iter()
                 .map(|field| field.ident)
                 .map(SnakeCase)
-                .map(|ident| (indent.inc_by(3), ident, ": Decode::decode(decoder)"))
+                .map(|ident| (indent.inc_by(3), ident, ": Decode::decode(decoder)?"))
                 .map(DisplayConcat),
         );
         write!(
@@ -201,10 +201,10 @@ impl<const N: usize> Display for DisplayDecode<&'_ GenericOp<N>> {
                         {fields}\n\
             {indent}}}\n\
             {indent}impl Decode for {camel_ident} {{\n\
-            {indent}    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {{\n\
-            {indent}        Self {{\n\
+            {indent}    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {{\n\
+            {indent}        Ok(Self {{\n\
                                 {constructors}\n\
-            {indent}        }}\n\
+            {indent}        }})\n\
             {indent}    }}\n\
             {indent}}}\n\
             "

crates/ir2/build/mod.rs

@@ -121,7 +121,7 @@ fn generate_encode_rs(config: &Config, isa: &Isa, contents: &mut String) -> Resu
 
 fn generate_decode_rs(config: &Config, isa: &Isa, contents: &mut String) -> Result<(), Error> {
     let expected_size = match config.simd {
-        true => 50_000,
+        true => 55_000,
         false => 35_000,
     };
     write_to_buffer(contents, expected_size, |buffer| {

crates/ir2/src/decode/mod.rs

@@ -1,5 +1,7 @@
 #![allow(non_camel_case_types)]
 
+//! Definitions and utilities to decode encoded byte streams.
+
 mod op;
 
 use self::op::{
@@ -42,56 +44,76 @@ use crate::{
     Slot,
     SlotSpan,
 };
-use core::{mem, num::NonZero};
+use core::{
+    error::Error as CoreError,
+    fmt,
+    mem::{self, MaybeUninit},
+    num::NonZero,
+};
 
 /// Types that can be used to decode types implementing [`Decode`].
 pub trait Decoder {
     /// Reads enough bytes from `self` to populate `buffer`.
-    fn read_bytes(&mut self, buffer: &mut [u8]);
+    fn read_bytes(&mut self, buffer: &mut [u8]) -> Result<(), DecodeError>;
+}
+
+/// An error that may be returned when decoding items.
+#[derive(Debug, Copy, Clone)]
+pub enum DecodeError {
+    /// The decoder ran out of bytes.
+    OutOfBytes,
+    /// The decoder found an invalid bit pattern.
+    InvalidBitPattern,
+}
+
+impl CoreError for DecodeError {}
+impl fmt::Display for DecodeError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let s = match self {
+            DecodeError::OutOfBytes => "ran out of bytes to decode",
+            DecodeError::InvalidBitPattern => "encountered invalid bit pattern",
+        };
+        f.write_str(s)
+    }
 }
 
 /// Types that can be decoded using a type that implements [`Decoder`].
-pub trait Decode {
+pub trait Decode: Sized {
     /// Decodes `Self` via `decoder`.
-    ///
-    /// # Safety
-    ///
-    /// It is the callers responsibility to ensure that the decoder
-    /// decodes items in the order they have been encoded and on valid
-    /// positions within the decode stream.
-    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self;
+    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError>;
 }
 
 impl Decode for BoundedSlotSpan {
-    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
-        let span = SlotSpan::decode(decoder);
-        let len = u16::decode(decoder);
-        Self::new(span, len)
+    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+        let span = SlotSpan::decode(decoder)?;
+        let len = u16::decode(decoder)?;
+        Ok(Self::new(span, len))
     }
 }
 
 impl<const N: u16> Decode for FixedSlotSpan<N> {
-    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
-        Self::new_unchecked(SlotSpan::decode(decoder))
+    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+        let span = SlotSpan::decode(decoder)?;
+        Self::new(span).map_err(|_| DecodeError::InvalidBitPattern)
     }
 }
 
 impl Decode for BranchTableTarget {
-    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
-        let results = SlotSpan::decode(decoder);
-        let offset = BranchOffset::decode(decoder);
-        Self::new(results, offset)
+    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+        let results = SlotSpan::decode(decoder)?;
+        let offset = BranchOffset::decode(decoder)?;
+        Ok(Self::new(results, offset))
     }
 }
 
 macro_rules! impl_decode_for_primitive {
     ( $($ty:ty),* $(,)? ) => {
         $(
             impl Decode for $ty {
-                unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
+                fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
                     let mut bytes = [0_u8; mem::size_of::<$ty>()];
-                    decoder.read_bytes(&mut bytes);
-                    Self::from_ne_bytes(bytes)
+                    decoder.read_bytes(&mut bytes)?;
+                    Ok(Self::from_ne_bytes(bytes))
                 }
             }
         )*
@@ -105,8 +127,8 @@ macro_rules! impl_decode_using {
     ( $($ty:ty as $as:ty = $e:expr),* $(,)? ) => {
         $(
             impl Decode for $ty {
-                unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
-                    $e(<$as as Decode>::decode(decoder))
+                fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+                    Ok($e(<$as as Decode>::decode(decoder)?))
                 }
             }
         )*
@@ -127,32 +149,66 @@ impl_decode_using! {
     Data as u32 = Into::into,
     Elem as u32 = Into::into,
 
-    Address as u64 = |address| unsafe { Address::try_from(address).unwrap_unchecked() },
     Sign<f32> as bool = Sign::new,
     Sign<f64> as bool = Sign::new,
     SlotSpan as Slot = SlotSpan::new,
-    NonZero<i32> as i32 = |value| unsafe { NonZero::new_unchecked(value) },
-    NonZero<i64> as i64 = |value| unsafe { NonZero::new_unchecked(value) },
-    NonZero<u32> as u32 = |value| unsafe { NonZero::new_unchecked(value) },
-    NonZero<u64> as u64 = |value| unsafe { NonZero::new_unchecked(value) },
-    TrapCode as u8 = |code: u8| -> TrapCode {
-        TrapCode::try_from(code).unwrap_unchecked()
+}
+
+macro_rules! impl_decode_fallible_using {
+    ( $($ty:ty as $as:ty = $e:expr),* $(,)? ) => {
+        $(
+            impl Decode for $ty {
+                fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+                    $e(<$as as Decode>::decode(decoder)?)
+                }
+            }
+        )*
+    };
+}
+impl_decode_fallible_using! {
+    Address as u64 = |address| {
+        Address::try_from(address).map_err(|_| DecodeError::InvalidBitPattern)
+    },
+    NonZero<i32> as i32 = |value| {
+        NonZero::new(value).ok_or(DecodeError::InvalidBitPattern)
+    },
+    NonZero<i64> as i64 = |value| {
+        NonZero::new(value).ok_or(DecodeError::InvalidBitPattern)
+    },
+    NonZero<u32> as u32 = |value| {
+        NonZero::new(value).ok_or(DecodeError::InvalidBitPattern)
+    },
+    NonZero<u64> as u64 = |value| {
+        NonZero::new(value).ok_or(DecodeError::InvalidBitPattern)
+    },
+    TrapCode as u8 = |code: u8| {
+        TrapCode::try_from(code).map_err(|_| DecodeError::InvalidBitPattern)
     },
-    OpCode as u16 = |code: u16| -> OpCode {
-        OpCode::try_from(code).unwrap_unchecked()
+    OpCode as u16 = |code: u16| {
+        OpCode::try_from(code).map_err(|_| DecodeError::InvalidBitPattern)
     }
 }
 
 impl<const N: usize, T: Decode> Decode for [T; N] {
-    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
-        core::array::from_fn(|_| <T as Decode>::decode(decoder))
+    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+        let mut array = <MaybeUninit<[T; N]>>::uninit();
+        // Safety: we are going to decode and initialize all array items and won't read any.
+        let items = unsafe { &mut *array.as_mut_ptr() };
+        for item in items {
+            *item = <T as Decode>::decode(decoder)?;
+        }
+        // Safety: we have decoded and thus initialized all array items.
+        let array = unsafe { array.assume_init() };
+        Ok(array)
     }
 }
 
 #[cfg(feature = "simd")]
 impl<const N: u8> Decode for ImmLaneIdx<N> {
-    unsafe fn decode<D: Decoder>(decoder: &mut D) -> Self {
-        ImmLaneIdx::try_from(u8::decode(decoder)).unwrap_unchecked()
+    fn decode<D: Decoder>(decoder: &mut D) -> Result<Self, DecodeError> {
+        let byte = u8::decode(decoder)?;
+        let lane = ImmLaneIdx::try_from(byte).map_err(|_| DecodeError::InvalidBitPattern)?;
+        Ok(lane)
     }
 }