light_system_program/sdk/
address.rs

1use std::collections::HashMap;
2
3use anchor_lang::{err, solana_program::pubkey::Pubkey, Result};
4use light_utils::hash_to_bn254_field_size_be;
5
6use crate::{errors::SystemProgramError, NewAddressParams, NewAddressParamsPacked};
7pub fn derive_address(merkle_tree_pubkey: &Pubkey, seed: &[u8; 32]) -> Result<[u8; 32]> {
8    let hash = match hash_to_bn254_field_size_be(
9        [merkle_tree_pubkey.to_bytes(), *seed].concat().as_slice(),
10    ) {
11        Some(hash) => Ok::<[u8; 32], SystemProgramError>(hash.0),
12        None => return err!(SystemProgramError::DeriveAddressError),
13    }?;
14
15    Ok(hash)
16}
17
18pub fn add_and_get_remaining_account_indices(
19    pubkeys: &[Pubkey],
20    remaining_accounts: &mut HashMap<Pubkey, usize>,
21) -> Vec<u8> {
22    let mut vec = Vec::new();
23    let mut next_index: usize = remaining_accounts.len();
24    for pubkey in pubkeys.iter() {
25        match remaining_accounts.get(pubkey) {
26            Some(_) => {}
27            None => {
28                remaining_accounts.insert(*pubkey, next_index);
29                next_index += 1;
30            }
31        };
32        vec.push(*remaining_accounts.get(pubkey).unwrap() as u8);
33    }
34    vec
35}
36// Helper function to pack new address params for instruction data in rust clients
37pub fn pack_new_address_params(
38    new_address_params: &[NewAddressParams],
39    remaining_accounts: &mut HashMap<Pubkey, usize>,
40) -> Vec<NewAddressParamsPacked> {
41    let mut new_address_params_packed = new_address_params
42        .iter()
43        .map(|x| NewAddressParamsPacked {
44            seed: x.seed,
45            address_merkle_tree_root_index: x.address_merkle_tree_root_index,
46            address_merkle_tree_account_index: 0, // will be assigned later
47            address_queue_account_index: 0,       // will be assigned later
48        })
49        .collect::<Vec<NewAddressParamsPacked>>();
50    let mut next_index: usize = remaining_accounts.len();
51    for (i, params) in new_address_params.iter().enumerate() {
52        match remaining_accounts.get(&params.address_merkle_tree_pubkey) {
53            Some(_) => {}
54            None => {
55                remaining_accounts.insert(params.address_merkle_tree_pubkey, next_index);
56                next_index += 1;
57            }
58        };
59        new_address_params_packed[i].address_merkle_tree_account_index = *remaining_accounts
60            .get(&params.address_merkle_tree_pubkey)
61            .unwrap()
62            as u8;
63    }
64
65    for (i, params) in new_address_params.iter().enumerate() {
66        match remaining_accounts.get(&params.address_queue_pubkey) {
67            Some(_) => {}
68            None => {
69                remaining_accounts.insert(params.address_queue_pubkey, next_index);
70                next_index += 1;
71            }
72        };
73        new_address_params_packed[i].address_queue_account_index = *remaining_accounts
74            .get(&params.address_queue_pubkey)
75            .unwrap() as u8;
76    }
77    new_address_params_packed
78}
79
80#[cfg(test)]
81mod tests {
82    use solana_sdk::{signature::Keypair, signer::Signer};
83
84    use super::*;
85
86    #[test]
87    fn test_derive_address_with_valid_input() {
88        let merkle_tree_pubkey = Keypair::new().pubkey();
89        let seeds = [1u8; 32];
90        let result = derive_address(&merkle_tree_pubkey, &seeds);
91        let result_2 = derive_address(&merkle_tree_pubkey, &seeds);
92        assert_eq!(result, result_2);
93    }
94
95    #[test]
96    fn test_derive_address_no_collision_same_seeds_diff_pubkey() {
97        let merkle_tree_pubkey = Keypair::new().pubkey();
98        let merkle_tree_pubkey_2 = Keypair::new().pubkey();
99        let seed = [2u8; 32];
100
101        let result = derive_address(&merkle_tree_pubkey, &seed);
102        let result_2 = derive_address(&merkle_tree_pubkey_2, &seed);
103        assert_ne!(result, result_2);
104    }
105
106    #[test]
107    fn test_add_and_get_remaining_account_indices_empty() {
108        let pubkeys = vec![];
109        let mut remaining_accounts = HashMap::new();
110        let result = add_and_get_remaining_account_indices(&pubkeys, &mut remaining_accounts);
111        assert!(result.is_empty());
112    }
113
114    #[test]
115    fn test_add_and_get_remaining_account_indices_single() {
116        let pubkey = Keypair::new().pubkey();
117        let pubkeys = vec![pubkey];
118        let mut remaining_accounts = HashMap::new();
119        let result = add_and_get_remaining_account_indices(&pubkeys, &mut remaining_accounts);
120        assert_eq!(result, vec![0]);
121        assert_eq!(remaining_accounts.get(&pubkey), Some(&0));
122    }
123
124    #[test]
125    fn test_add_and_get_remaining_account_indices_multiple() {
126        let pubkey1 = Keypair::new().pubkey();
127        let pubkey2 = Keypair::new().pubkey();
128        let pubkeys = vec![pubkey1, pubkey2];
129        let mut remaining_accounts = HashMap::new();
130        let result = add_and_get_remaining_account_indices(&pubkeys, &mut remaining_accounts);
131        assert_eq!(result, vec![0, 1]);
132        assert_eq!(remaining_accounts.get(&pubkey1), Some(&0));
133        assert_eq!(remaining_accounts.get(&pubkey2), Some(&1));
134    }
135
136    #[test]
137    fn test_add_and_get_remaining_account_indices_duplicates() {
138        let pubkey = Keypair::new().pubkey();
139        let pubkeys = vec![pubkey, pubkey];
140        let mut remaining_accounts = HashMap::new();
141        let result = add_and_get_remaining_account_indices(&pubkeys, &mut remaining_accounts);
142        assert_eq!(result, vec![0, 0]);
143        assert_eq!(remaining_accounts.get(&pubkey), Some(&0));
144        assert_eq!(remaining_accounts.len(), 1);
145    }
146
147    #[test]
148    fn test_add_and_get_remaining_account_indices_multiple_duplicates() {
149        let pubkey1 = Keypair::new().pubkey();
150        let pubkey2 = Keypair::new().pubkey();
151        let pubkey3 = Keypair::new().pubkey();
152        let pubkeys = vec![pubkey1, pubkey2, pubkey1, pubkey3, pubkey2, pubkey1];
153        let mut remaining_accounts = HashMap::new();
154        let result = add_and_get_remaining_account_indices(&pubkeys, &mut remaining_accounts);
155        assert_eq!(result, vec![0, 1, 0, 2, 1, 0]);
156        assert_eq!(remaining_accounts.get(&pubkey1), Some(&0));
157        assert_eq!(remaining_accounts.get(&pubkey2), Some(&1));
158        assert_eq!(remaining_accounts.get(&pubkey3), Some(&2));
159        assert_eq!(remaining_accounts.len(), 3);
160    }
161}
light_system_program/sdk/address.rs

light_system_program/sdk/
address.rs
