safe_token_2022/
pod.rs

1//! Safecoin program utilities for Plain Old Data types
2use {
3    bytemuck::{Pod, Zeroable},
4    solana_program::{program_error::ProgramError, program_option::COption, pubkey::Pubkey},
5    safe_zk_token_sdk::zk_token_elgamal::pod,
6    std::convert::TryFrom,
7};
8
9/// A Pubkey that encodes `None` as all `0`, meant to be usable as a Pod type,
10/// similar to all NonZero* number types from the bytemuck library.
11#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
12#[repr(transparent)]
13pub struct OptionalNonZeroPubkey(Pubkey);
14impl TryFrom<Option<Pubkey>> for OptionalNonZeroPubkey {
15    type Error = ProgramError;
16    fn try_from(p: Option<Pubkey>) -> Result<Self, Self::Error> {
17        match p {
18            None => Ok(Self(Pubkey::default())),
19            Some(pubkey) => {
20                if pubkey == Pubkey::default() {
21                    Err(ProgramError::InvalidArgument)
22                } else {
23                    Ok(Self(pubkey))
24                }
25            }
26        }
27    }
28}
29impl TryFrom<COption<Pubkey>> for OptionalNonZeroPubkey {
30    type Error = ProgramError;
31    fn try_from(p: COption<Pubkey>) -> Result<Self, Self::Error> {
32        match p {
33            COption::None => Ok(Self(Pubkey::default())),
34            COption::Some(pubkey) => {
35                if pubkey == Pubkey::default() {
36                    Err(ProgramError::InvalidArgument)
37                } else {
38                    Ok(Self(pubkey))
39                }
40            }
41        }
42    }
43}
44impl std::fmt::Display for OptionalNonZeroPubkey {
45    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
46        write!(f, "{:?}", &self)
47    }
48}
49impl From<OptionalNonZeroPubkey> for Option<Pubkey> {
50    fn from(p: OptionalNonZeroPubkey) -> Self {
51        if p.0 == Pubkey::default() {
52            None
53        } else {
54            Some(p.0)
55        }
56    }
57}
58impl From<OptionalNonZeroPubkey> for COption<Pubkey> {
59    fn from(p: OptionalNonZeroPubkey) -> Self {
60        if p.0 == Pubkey::default() {
61            COption::None
62        } else {
63            COption::Some(p.0)
64        }
65    }
66}
67
68/// ElGamal public key used for encryption
69pub type EncryptionPubkey = pod::ElGamalPubkey;
70/// An EncryptionPubkey that encodes `None` as all `0`, meant to be usable as a Pod type.
71#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
72#[repr(transparent)]
73pub struct OptionalNonZeroEncryptionPubkey(EncryptionPubkey);
74impl OptionalNonZeroEncryptionPubkey {
75    /// Checks equality between an OptionalNonZeroEncryptionPubkey and an EncryptionPubkey when
76    /// interpreted as bytes.
77    pub fn equals(&self, other: &EncryptionPubkey) -> bool {
78        &self.0 == other
79    }
80}
81impl TryFrom<Option<EncryptionPubkey>> for OptionalNonZeroEncryptionPubkey {
82    type Error = ProgramError;
83    fn try_from(p: Option<EncryptionPubkey>) -> Result<Self, Self::Error> {
84        match p {
85            None => Ok(Self(EncryptionPubkey::default())),
86            Some(encryption_pubkey) => {
87                if encryption_pubkey == EncryptionPubkey::default() {
88                    Err(ProgramError::InvalidArgument)
89                } else {
90                    Ok(Self(encryption_pubkey))
91                }
92            }
93        }
94    }
95}
96impl std::fmt::Display for OptionalNonZeroEncryptionPubkey {
97    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
98        write!(f, "{:?}", &self)
99    }
100}
101impl From<OptionalNonZeroEncryptionPubkey> for Option<EncryptionPubkey> {
102    fn from(p: OptionalNonZeroEncryptionPubkey) -> Self {
103        if p.0 == EncryptionPubkey::default() {
104            None
105        } else {
106            Some(p.0)
107        }
108    }
109}
110
111/// The standard `bool` is not a `Pod`, define a replacement that is
112#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
113#[repr(transparent)]
114pub struct PodBool(u8);
115impl From<bool> for PodBool {
116    fn from(b: bool) -> Self {
117        Self(if b { 1 } else { 0 })
118    }
119}
120impl From<&PodBool> for bool {
121    fn from(b: &PodBool) -> Self {
122        b.0 != 0
123    }
124}
125
126impl From<PodBool> for bool {
127    fn from(b: PodBool) -> Self {
128        b.0 != 0
129    }
130}
131
132/// Simple macro for implementing conversion functions between Pod* ints and standard ints.
133///
134/// The standard int types can cause alignment issues when placed in a `Pod`,
135/// so these replacements are usable in all `Pod`s.
136macro_rules! impl_int_conversion {
137    ($P:ty, $I:ty) => {
138        impl From<$I> for $P {
139            fn from(n: $I) -> Self {
140                Self(n.to_le_bytes())
141            }
142        }
143        impl From<$P> for $I {
144            fn from(pod: $P) -> Self {
145                Self::from_le_bytes(pod.0)
146            }
147        }
148    };
149}
150
151/// `u16` type that can be used in `Pod`s
152#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
153#[repr(transparent)]
154pub struct PodU16([u8; 2]);
155impl_int_conversion!(PodU16, u16);
156
157/// `i16` type that can be used in `Pod`s
158#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
159#[repr(transparent)]
160pub struct PodI16([u8; 2]);
161impl_int_conversion!(PodI16, i16);
162
163/// `u64` type that can be used in `Pod`s
164#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
165#[repr(transparent)]
166pub struct PodU64([u8; 8]);
167impl_int_conversion!(PodU64, u64);
168
169/// `i64` type that can be used in `Pod`s
170#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
171#[repr(transparent)]
172pub struct PodI64([u8; 8]);
173impl_int_conversion!(PodI64, i64);
174
175/// On-chain size of a `Pod` type
176pub fn pod_get_packed_len<T: Pod>() -> usize {
177    std::mem::size_of::<T>()
178}
179
180/// Convert a `Pod` into a slice (zero copy)
181pub fn pod_bytes_of<T: Pod>(t: &T) -> &[u8] {
182    bytemuck::bytes_of(t)
183}
184
185/// Convert a slice into a `Pod` (zero copy)
186pub fn pod_from_bytes<T: Pod>(bytes: &[u8]) -> Result<&T, ProgramError> {
187    bytemuck::try_from_bytes(bytes).map_err(|_| ProgramError::InvalidArgument)
188}
189
190/// Maybe convert a slice into a `Pod` (zero copy)
191///
192/// Returns `None` if the slice is empty, but `Err` if all other lengths but `get_packed_len()`
193/// This function exists primarily because `Option<T>` is not a `Pod`.
194pub fn pod_maybe_from_bytes<T: Pod>(bytes: &[u8]) -> Result<Option<&T>, ProgramError> {
195    if bytes.is_empty() {
196        Ok(None)
197    } else {
198        bytemuck::try_from_bytes(bytes)
199            .map(Some)
200            .map_err(|_| ProgramError::InvalidArgument)
201    }
202}
203
204/// Convert a slice into a mutable `Pod` (zero copy)
205pub fn pod_from_bytes_mut<T: Pod>(bytes: &mut [u8]) -> Result<&mut T, ProgramError> {
206    bytemuck::try_from_bytes_mut(bytes).map_err(|_| ProgramError::InvalidArgument)
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212
213    #[test]
214    fn test_pod_bool() {
215        assert!(pod_from_bytes::<PodBool>(&[]).is_err());
216        assert!(pod_from_bytes::<PodBool>(&[0, 0]).is_err());
217
218        for i in 0..=u8::MAX {
219            assert_eq!(i != 0, bool::from(pod_from_bytes::<PodBool>(&[i]).unwrap()));
220        }
221    }
222
223    #[test]
224    fn test_pod_u64() {
225        assert!(pod_from_bytes::<PodU64>(&[]).is_err());
226        assert_eq!(
227            1u64,
228            u64::from(*pod_from_bytes::<PodU64>(&[1, 0, 0, 0, 0, 0, 0, 0]).unwrap())
229        );
230    }
231
232    #[test]
233    fn test_pod_option() {
234        assert_eq!(
235            Some(Pubkey::new_from_array([1; 32])),
236            Option::<Pubkey>::from(*pod_from_bytes::<OptionalNonZeroPubkey>(&[1; 32]).unwrap())
237        );
238        assert_eq!(
239            None,
240            Option::<Pubkey>::from(*pod_from_bytes::<OptionalNonZeroPubkey>(&[0; 32]).unwrap())
241        );
242        assert!(pod_from_bytes::<OptionalNonZeroPubkey>(&[]).is_err());
243        assert!(pod_from_bytes::<OptionalNonZeroPubkey>(&[0; 1]).is_err());
244        assert!(pod_from_bytes::<OptionalNonZeroPubkey>(&[1; 1]).is_err());
245    }
246}
safe_token_2022/pod.rs

safe_token_2022/
pod.rs
