1use std::{future::Future, ops::Deref};
2
3use anchor_client::{
4 anchor_lang::system_program,
5 solana_sdk::{pubkey::Pubkey, signer::Signer},
6};
7use gmsol_solana_utils::transaction_builder::TransactionBuilder;
8use gmsol_store::{
9 accounts, instruction,
10 states::user::{ReferralCodeBytes, ReferralCodeV2, UserHeader},
11};
12
13use crate::utils::ZeroCopy;
14
15pub trait UserOps<C> {
17 fn prepare_user(&self, store: &Pubkey) -> crate::Result<TransactionBuilder<C>>;
19
20 fn initialize_referral_code(
22 &self,
23 store: &Pubkey,
24 code: ReferralCodeBytes,
25 ) -> crate::Result<TransactionBuilder<C>>;
26
27 fn set_referrer(
29 &self,
30 store: &Pubkey,
31 code: ReferralCodeBytes,
32 hint_referrer: Option<Pubkey>,
33 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
34
35 fn transfer_referral_code(
37 &self,
38 store: &Pubkey,
39 receiver: &Pubkey,
40 hint_code: Option<ReferralCodeBytes>,
41 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
42
43 fn cancel_referral_code_transfer(
45 &self,
46 store: &Pubkey,
47 hint_code: Option<ReferralCodeBytes>,
48 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
49
50 fn accept_referral_code(
52 &self,
53 store: &Pubkey,
54 code: ReferralCodeBytes,
55 hint_owner: Option<Pubkey>,
56 ) -> impl Future<Output = crate::Result<TransactionBuilder<C>>>;
57}
58
59impl<C: Deref<Target = impl Signer> + Clone> UserOps<C> for crate::Client<C> {
60 fn prepare_user(&self, store: &Pubkey) -> crate::Result<TransactionBuilder<C>> {
61 let owner = self.payer();
62 let user = self.find_user_address(store, &owner);
63 let rpc = self
64 .store_transaction()
65 .anchor_accounts(accounts::PrepareUser {
66 owner,
67 store: *store,
68 user,
69 system_program: system_program::ID,
70 })
71 .anchor_args(instruction::PrepareUser {});
72 Ok(rpc)
73 }
74
75 fn initialize_referral_code(
76 &self,
77 store: &Pubkey,
78 code: ReferralCodeBytes,
79 ) -> crate::Result<TransactionBuilder<C>> {
80 let owner = self.payer();
81 let referral_code = self.find_referral_code_address(store, code);
82 let user = self.find_user_address(store, &owner);
83 let rpc = self
84 .store_transaction()
85 .anchor_accounts(accounts::InitializeReferralCode {
86 owner,
87 store: *store,
88 referral_code,
89 user,
90 system_program: system_program::ID,
91 })
92 .anchor_args(instruction::InitializeReferralCode { code });
93 Ok(rpc)
94 }
95
96 async fn set_referrer(
97 &self,
98 store: &Pubkey,
99 code: ReferralCodeBytes,
100 hint_referrer_user: Option<Pubkey>,
101 ) -> crate::Result<TransactionBuilder<C>> {
102 let owner = self.payer();
103 let user = self.find_user_address(store, &owner);
104
105 let referral_code = self.find_referral_code_address(store, code);
106
107 let referrer_user = match hint_referrer_user {
108 Some(referrer) => referrer,
109 None => {
110 let code = self
111 .account::<ZeroCopy<ReferralCodeV2>>(&referral_code)
112 .await?
113 .ok_or(crate::Error::NotFound)?
114 .0;
115 let owner = code.owner;
116 self.find_user_address(store, &owner)
117 }
118 };
119
120 let rpc = self
121 .store_transaction()
122 .anchor_accounts(accounts::SetReferrer {
123 owner,
124 store: *store,
125 user,
126 referral_code,
127 referrer_user,
128 })
129 .anchor_args(instruction::SetReferrer { code });
130
131 Ok(rpc)
132 }
133
134 async fn transfer_referral_code(
135 &self,
136 store: &Pubkey,
137 receiver: &Pubkey,
138 hint_code: Option<ReferralCodeBytes>,
139 ) -> crate::Result<TransactionBuilder<C>> {
140 let owner = self.payer();
141 let user = self.find_user_address(store, &owner);
142 let receiver_user = self.find_user_address(store, receiver);
143
144 let referral_code = match hint_code {
145 Some(code) => self.find_referral_code_address(store, code),
146 None => {
147 let user = self
148 .account::<ZeroCopy<UserHeader>>(&user)
149 .await?
150 .ok_or(crate::Error::NotFound)?;
151 *user
152 .0
153 .referral()
154 .code()
155 .ok_or(crate::Error::invalid_argument("referral code is not set"))?
156 }
157 };
158
159 let rpc = self
160 .store_transaction()
161 .anchor_accounts(accounts::TransferReferralCode {
162 owner,
163 store: *store,
164 user,
165 referral_code,
166 receiver_user,
167 })
168 .anchor_args(instruction::TransferReferralCode {});
169
170 Ok(rpc)
171 }
172
173 async fn cancel_referral_code_transfer(
174 &self,
175 store: &Pubkey,
176 hint_code: Option<ReferralCodeBytes>,
177 ) -> crate::Result<TransactionBuilder<C>> {
178 let owner = self.payer();
179 let user = self.find_user_address(store, &owner);
180
181 let referral_code = match hint_code {
182 Some(code) => self.find_referral_code_address(store, code),
183 None => {
184 let user = self
185 .account::<ZeroCopy<UserHeader>>(&user)
186 .await?
187 .ok_or(crate::Error::NotFound)?;
188 *user
189 .0
190 .referral()
191 .code()
192 .ok_or(crate::Error::invalid_argument("referral code is not set"))?
193 }
194 };
195
196 let rpc = self
197 .store_transaction()
198 .anchor_accounts(accounts::CancelReferralCodeTransfer {
199 owner,
200 store: *store,
201 user,
202 referral_code,
203 })
204 .anchor_args(instruction::CancelReferralCodeTransfer {});
205
206 Ok(rpc)
207 }
208
209 async fn accept_referral_code(
210 &self,
211 store: &Pubkey,
212 code: ReferralCodeBytes,
213 hint_owner: Option<Pubkey>,
214 ) -> crate::Result<TransactionBuilder<C>> {
215 let next_owner = self.payer();
216 let receiver_user = self.find_user_address(store, &next_owner);
217 let referral_code = self.find_referral_code_address(store, code);
218
219 let owner = match hint_owner {
220 Some(owner) => owner,
221 None => {
222 let code = self
223 .account::<ZeroCopy<ReferralCodeV2>>(&referral_code)
224 .await?
225 .ok_or(crate::Error::NotFound)?
226 .0;
227 code.owner
228 }
229 };
230
231 let user = self.find_user_address(store, &owner);
232
233 let rpc = self
234 .store_transaction()
235 .anchor_accounts(accounts::AcceptReferralCode {
236 next_owner,
237 store: *store,
238 user,
239 referral_code,
240 receiver_user,
241 })
242 .anchor_args(instruction::AcceptReferralCode {});
243 Ok(rpc)
244 }
245}