From 7e2a370219609795bd76f8ca636fa75bfef9d90a Mon Sep 17 00:00:00 2001
From: Christian Kamm <mail@ckamm.de>
Date: Mon, 5 Dec 2022 13:54:19 +0100
Subject: [PATCH] Add account(zero_copy(safe_bytemuck_derives)) feature.

The default account(zero_copy) feature unsafe impls the bytemuck traits

The new one derives it instead, which runs desired sanity checks like
"struct has no padding".
---
 lang/attribute/account/src/lib.rs | 34 +++++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/lang/attribute/account/src/lib.rs b/lang/attribute/account/src/lib.rs
index 273be17322..4b3ea0d600 100644
--- a/lang/attribute/account/src/lib.rs
+++ b/lang/attribute/account/src/lib.rs
@@ -66,6 +66,7 @@ pub fn account(
 ) -> proc_macro::TokenStream {
     let mut namespace = "".to_string();
     let mut is_zero_copy = false;
+    let mut safe_bytemuck = false;
     let args_str = args.to_string();
     let args: Vec<&str> = args_str.split(',').collect();
     if args.len() > 2 {
@@ -80,6 +81,10 @@ pub fn account(
             .collect();
         if ns == "zero_copy" {
             is_zero_copy = true;
+            safe_bytemuck = false;
+        } else if ns == "zero_copy(safe_bytemuck_derives)" {
+            is_zero_copy = true;
+            safe_bytemuck = true;
         } else {
             namespace = ns;
         }
@@ -123,16 +128,37 @@ pub fn account(
         }
     };
 
+    let unsafe_bytemuck_impl = {
+        if !safe_bytemuck {
+            quote! {
+                #[automatically_derived]
+                unsafe impl #impl_gen anchor_lang::__private::bytemuck::Pod for #account_name #type_gen #where_clause {}
+                #[automatically_derived]
+                unsafe impl #impl_gen anchor_lang::__private::bytemuck::Zeroable for #account_name #type_gen #where_clause {}
+            }
+        } else {
+            quote! {}
+        }
+    };
+
+    let safe_bytemuck_derives = {
+        if safe_bytemuck {
+            quote! {
+                #[derive(bytemuck::Pod, bytemuck::Zeroable)]
+            }
+        } else {
+            quote! {}
+        }
+    };
+
     proc_macro::TokenStream::from({
         if is_zero_copy {
             quote! {
                 #[zero_copy]
+                #safe_bytemuck_derives
                 #account_strct
 
-                #[automatically_derived]
-                unsafe impl #impl_gen anchor_lang::__private::bytemuck::Pod for #account_name #type_gen #where_clause {}
-                #[automatically_derived]
-                unsafe impl #impl_gen anchor_lang::__private::bytemuck::Zeroable for #account_name #type_gen #where_clause {}
+                #unsafe_bytemuck_impl
 
                 #[automatically_derived]
                 impl #impl_gen anchor_lang::ZeroCopy for #account_name #type_gen #where_clause {}