From c26fa1d0e9038946a6ce5f42d5672dea52d3ab3e Mon Sep 17 00:00:00 2001
From: Sam Kim <skim13@cs.stanford.edu>
Date: Thu, 16 Dec 2021 09:15:29 -0500
Subject: [PATCH] refactor: create pod struct for ZeroBalanceProof

---
 zk-token-sdk/src/instruction/close_account.rs |  7 ++++---
 zk-token-sdk/src/zk_token_elgamal/convert.rs  | 16 +++++++++++++++-
 zk-token-sdk/src/zk_token_elgamal/pod.rs      | 10 ++++++++++
 3 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/zk-token-sdk/src/instruction/close_account.rs b/zk-token-sdk/src/instruction/close_account.rs
index 66f21aca0..e1dbc4d80 100644
--- a/zk-token-sdk/src/instruction/close_account.rs
+++ b/zk-token-sdk/src/instruction/close_account.rs
@@ -11,6 +11,7 @@ use {
         },
         errors::ProofError,
         instruction::Verifiable,
+        sigma_proofs::zero_balance_proof::ZeroBalanceProof,
         transcript::TranscriptProtocol,
     },
     curve25519_dalek::{
@@ -71,9 +72,7 @@ impl Verifiable for CloseAccountData {
 #[repr(C)]
 #[allow(non_snake_case)]
 pub struct CloseAccountProof {
-    pub Y_P: pod::CompressedRistretto, // 32 bytes
-    pub Y_D: pod::CompressedRistretto, // 32 bytes
-    pub z: pod::Scalar,                // 32 bytes
+    pub proof: pod::ZeroBalanceProof,
 }
 
 #[allow(non_snake_case)]
@@ -89,6 +88,8 @@ impl CloseAccountProof {
         // add a domain separator to record the start of the protocol
         transcript.close_account_proof_domain_sep();
 
+        
+
         // extract the relevant scalar and Ristretto points from the input
         let P = source_keypair.public.get_point();
         let s = source_keypair.secret.get_scalar();
diff --git a/zk-token-sdk/src/zk_token_elgamal/convert.rs b/zk-token-sdk/src/zk_token_elgamal/convert.rs
index b216cab1d..5483d100a 100644
--- a/zk-token-sdk/src/zk_token_elgamal/convert.rs
+++ b/zk-token-sdk/src/zk_token_elgamal/convert.rs
@@ -20,7 +20,7 @@ mod target_arch {
                 elgamal::{ElGamalCiphertext, ElGamalPubkey},
                 pedersen::{PedersenCommitment, PedersenDecryptHandle},
             },
-            sigma_proofs::{equality_proof::EqualityProof, validity_proof::ValidityProof},
+            sigma_proofs::{equality_proof::EqualityProof, validity_proof::ValidityProof, zero_balance_proof::ZeroBalanceProof},
             errors::ProofError,
             range_proof::RangeProof,
         },
@@ -169,6 +169,20 @@ mod target_arch {
         }
     }
 
+    impl From<ZeroBalanceProof> for pod::ZeroBalanceProof {
+        fn from(proof: ZeroBalanceProof) -> Self {
+            Self(proof.to_bytes())
+        }
+    }
+
+    impl TryFrom<pod::ZeroBalanceProof> for ZeroBalanceProof {
+        type Error = ProofError;
+
+        fn try_from(pod: pod::ZeroBalanceProof) -> Result<Self, Self::Error> {
+            Self::from_bytes(&pod.0)
+        }
+    }
+
     impl TryFrom<RangeProof> for pod::RangeProof64 {
         type Error = ProofError;
 
diff --git a/zk-token-sdk/src/zk_token_elgamal/pod.rs b/zk-token-sdk/src/zk_token_elgamal/pod.rs
index 06b708df6..3514d25d3 100644
--- a/zk-token-sdk/src/zk_token_elgamal/pod.rs
+++ b/zk-token-sdk/src/zk_token_elgamal/pod.rs
@@ -69,6 +69,16 @@ pub struct ValidityProof(pub [u8; 160]);
 unsafe impl Zeroable for ValidityProof {}
 unsafe impl Pod for ValidityProof {}
 
+/// Serialization of zero balance proofs
+#[derive(Clone, Copy)]
+#[repr(transparent)]
+pub struct ZeroBalanceProof(pub [u8; 160]);
+
+// `ZeroBalanceProof` is a Pod and Zeroable.
+// Add the marker traits manually because `bytemuck` only adds them for some `u8` arrays
+unsafe impl Zeroable for ZeroBalanceProof {}
+unsafe impl Pod for ZeroBalanceProof {}
+
 /// Serialization of range proofs for 64-bit numbers (for `Withdraw` instruction)
 #[derive(Clone, Copy)]
 #[repr(transparent)]