1use {
2 crate::errors::TokenProofGenerationError,
3 solana_zk_sdk::{
4 encryption::{
5 elgamal::{ElGamal, ElGamalCiphertext, ElGamalKeypair},
6 pedersen::Pedersen,
7 },
8 zk_elgamal_proof_program::proof_data::{
9 BatchedRangeProofU64Data, CiphertextCommitmentEqualityProofData,
10 },
11 },
12};
13
14const REMAINING_BALANCE_BIT_LENGTH: usize = 64;
15
16pub struct WithdrawProofData {
18 pub equality_proof_data: CiphertextCommitmentEqualityProofData,
19 pub range_proof_data: BatchedRangeProofU64Data,
20}
21
22pub fn withdraw_proof_data(
23 current_available_balance: &ElGamalCiphertext,
24 current_balance: u64,
25 withdraw_amount: u64,
26 elgamal_keypair: &ElGamalKeypair,
27) -> Result<WithdrawProofData, TokenProofGenerationError> {
28 let remaining_balance = current_balance
30 .checked_sub(withdraw_amount)
31 .ok_or(TokenProofGenerationError::NotEnoughFunds)?;
32
33 let (remaining_balance_commitment, remaining_balance_opening) =
35 Pedersen::new(remaining_balance);
36
37 #[allow(clippy::arithmetic_side_effects)]
39 let remaining_balance_ciphertext = current_available_balance - ElGamal::encode(withdraw_amount);
40
41 let equality_proof_data = CiphertextCommitmentEqualityProofData::new(
43 elgamal_keypair,
44 &remaining_balance_ciphertext,
45 &remaining_balance_commitment,
46 &remaining_balance_opening,
47 remaining_balance,
48 )
49 .map_err(TokenProofGenerationError::from)?;
50
51 let range_proof_data = BatchedRangeProofU64Data::new(
52 vec![&remaining_balance_commitment],
53 vec![remaining_balance],
54 vec![REMAINING_BALANCE_BIT_LENGTH],
55 vec![&remaining_balance_opening],
56 )
57 .map_err(TokenProofGenerationError::from)?;
58
59 Ok(WithdrawProofData {
60 equality_proof_data,
61 range_proof_data,
62 })
63}