From dc0de77b3f8da2983a0a14f4b695e740c1ffaf8f Mon Sep 17 00:00:00 2001
From: Nikita Baksalyar <nikita.baksalyar@gmail.com>
Date: Fri, 6 Dec 2024 23:33:48 +0000
Subject: [PATCH] Remove unsafe code from macro expansions (#818)

* Remove unsafe code from macro expansions

* Use a static match to convert between enum variants and `u8`

* Use a static match in `abi_decode_raw`

* chore: clean up

* fix

---------

Co-authored-by: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
---
 .../sol-macro-expander/src/expand/contract.rs |  8 ++---
 crates/sol-macro-expander/src/expand/enum.rs  | 35 ++++++++-----------
 2 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/crates/sol-macro-expander/src/expand/contract.rs b/crates/sol-macro-expander/src/expand/contract.rs
index e32578e57..6edd55756 100644
--- a/crates/sol-macro-expander/src/expand/contract.rs
+++ b/crates/sol-macro-expander/src/expand/contract.rs
@@ -488,7 +488,7 @@ pub(super) fn expand(cx: &mut ExpCtxt<'_>, contract: &ItemContract) -> Result<To
         #mod_descr_doc
         #(#mod_attrs)*
         #mod_iface_doc
-        #[allow(non_camel_case_types, non_snake_case, clippy::pub_underscore_fields, clippy::style)]
+        #[allow(non_camel_case_types, non_snake_case, clippy::pub_underscore_fields, clippy::style, clippy::empty_structs_with_brackets)]
         pub mod #name {
             use super::*;
             use #alloy_sol_types as alloy_sol_types;
@@ -715,7 +715,7 @@ impl CallLikeExpander<'_> {
                 }
 
                 #[inline]
-                #[allow(unsafe_code, non_snake_case)]
+                #[allow(non_snake_case)]
                 fn abi_decode_raw(
                     selector: [u8; 4],
                     data: &[u8],
@@ -737,8 +737,8 @@ impl CallLikeExpander<'_> {
                             selector,
                         ));
                     };
-                    // SAFETY: `idx` is a valid index into `DECODE_SHIMS`.
-                    (unsafe { DECODE_SHIMS.get_unchecked(idx) })(data, validate)
+                    // `SELECTORS` and `DECODE_SHIMS` have the same length and are sorted in the same order.
+                    DECODE_SHIMS[idx](data, validate)
                 }
 
                 #[inline]
diff --git a/crates/sol-macro-expander/src/expand/enum.rs b/crates/sol-macro-expander/src/expand/enum.rs
index 4e90b1456..4b694a13c 100644
--- a/crates/sol-macro-expander/src/expand/enum.rs
+++ b/crates/sol-macro-expander/src/expand/enum.rs
@@ -70,6 +70,12 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, enumm: &ItemEnum) -> Result<TokenStream>
     let uint8 = quote!(alloy_sol_types::sol_data::Uint<8>);
     let uint8_st = quote!(<#uint8 as alloy_sol_types::SolType>);
 
+    let index_to_variant = variants.iter().enumerate().map(|(idx, variant)| {
+        let ident = &variant.ident;
+        let idx = idx as u8;
+        quote! { #idx => ::core::result::Result::Ok(Self::#ident), }
+    });
+
     let doc = docs.then(|| mk_doc(format!("```solidity\n{enumm}\n```")));
     let tokens = quote! {
         #(#attrs)*
@@ -98,15 +104,13 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, enumm: &ItemEnum) -> Result<TokenStream>
             impl ::core::convert::TryFrom<u8> for #name {
                 type Error = alloy_sol_types::Error;
 
-                #[allow(unsafe_code)]
                 #[inline]
-                fn try_from(v: u8) -> alloy_sol_types::Result<Self> {
-                    if v <= #max {
-                        ::core::result::Result::Ok(unsafe { ::core::mem::transmute(v) })
-                    } else {
-                        ::core::result::Result::Err(alloy_sol_types::Error::InvalidEnumValue {
+                fn try_from(value: u8) -> alloy_sol_types::Result<Self> {
+                    match value {
+                        #(#index_to_variant)*
+                        value => ::core::result::Result::Err(alloy_sol_types::Error::InvalidEnumValue {
                             name: #name_s,
-                            value: v,
+                            value,
                             max: #max,
                         })
                     }
@@ -127,7 +131,7 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, enumm: &ItemEnum) -> Result<TokenStream>
 
                 #[inline]
                 fn stv_eip712_data_word(&self) -> alloy_sol_types::Word {
-                    #uint8_st::eip712_data_word(self.as_u8())
+                    #uint8_st::eip712_data_word(&(*self as u8))
                 }
 
                 #[inline]
@@ -170,17 +174,17 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, enumm: &ItemEnum) -> Result<TokenStream>
             impl alloy_sol_types::EventTopic for #name {
                 #[inline]
                 fn topic_preimage_length(rust: &Self::RustType) -> usize {
-                    <#uint8 as alloy_sol_types::EventTopic>::topic_preimage_length(rust.as_u8())
+                    <#uint8 as alloy_sol_types::EventTopic>::topic_preimage_length(&(*rust as u8))
                 }
 
                 #[inline]
                 fn encode_topic_preimage(rust: &Self::RustType, out: &mut alloy_sol_types::private::Vec<u8>) {
-                    <#uint8 as alloy_sol_types::EventTopic>::encode_topic_preimage(rust.as_u8(), out);
+                    <#uint8 as alloy_sol_types::EventTopic>::encode_topic_preimage(&(*rust as u8), out);
                 }
 
                 #[inline]
                 fn encode_topic(rust: &Self::RustType) -> alloy_sol_types::abi::token::WordToken {
-                    <#uint8 as alloy_sol_types::EventTopic>::encode_topic(rust.as_u8())
+                    <#uint8 as alloy_sol_types::EventTopic>::encode_topic(&(*rust as u8))
                 }
             }
 
@@ -188,15 +192,6 @@ pub(super) fn expand(cx: &ExpCtxt<'_>, enumm: &ItemEnum) -> Result<TokenStream>
             impl alloy_sol_types::SolEnum for #name {
                 const COUNT: usize = #count;
             }
-
-            #[automatically_derived]
-            impl #name {
-                #[allow(unsafe_code, clippy::inline_always)]
-                #[inline(always)]
-                fn as_u8(&self) -> &u8 {
-                    unsafe { &*(self as *const Self as *const u8) }
-                }
-            }
         };
     };
     Ok(tokens)