feat: batch cold account loads in light client

Author: sergeytimoshin

sdk-libs/client/src/indexer/photon_indexer.rs

@@ -1142,17 +1142,16 @@ impl Indexer for PhotonIndexer {
                 .value
                 .iter()
                 .map(|x| {
-                    let mut proof_vec = x.proof.clone();
-                    if proof_vec.len() < STATE_MERKLE_TREE_CANOPY_DEPTH {
+                    if x.proof.len() < STATE_MERKLE_TREE_CANOPY_DEPTH {
                         return Err(IndexerError::InvalidParameters(format!(
                             "Merkle proof length ({}) is less than canopy depth ({})",
-                            proof_vec.len(),
+                            x.proof.len(),
                             STATE_MERKLE_TREE_CANOPY_DEPTH,
                         )));
                     }
-                    proof_vec.truncate(proof_vec.len() - STATE_MERKLE_TREE_CANOPY_DEPTH);
+                    let proof_len = x.proof.len() - STATE_MERKLE_TREE_CANOPY_DEPTH;
 
-                    let proof = proof_vec
+                    let proof = x.proof[..proof_len]
                         .iter()
                         .map(|s| Hash::from_base58(s))
                         .collect::<Result<Vec<[u8; 32]>, IndexerError>>()
@@ -1703,15 +1702,13 @@ impl Indexer for PhotonIndexer {
 
     async fn get_subtrees(
         &self,
-        _merkle_tree_pubkey: [u8; 32],
+        merkle_tree_pubkey: [u8; 32],
         _config: Option<IndexerRpcConfig>,
     ) -> Result<Response<Items<[u8; 32]>>, IndexerError> {
-        #[cfg(not(feature = "v2"))]
-        unimplemented!();
-        #[cfg(feature = "v2")]
-        {
-            todo!();
-        }
+        Err(IndexerError::NotImplemented(format!(
+            "PhotonIndexer::get_subtrees is not implemented for merkle tree {}",
+            solana_pubkey::Pubkey::new_from_array(merkle_tree_pubkey)
+        )))
     }
 }
 

sdk-libs/client/src/indexer/types/queue.rs

@@ -66,11 +66,14 @@ pub struct AddressQueueData {
 }
 
 impl AddressQueueData {
+    const ADDRESS_TREE_HEIGHT: usize = 40;
+
     /// Reconstruct a merkle proof for a given low_element_index from the deduplicated nodes.
     pub fn reconstruct_proof<const HEIGHT: usize>(
         &self,
         address_idx: usize,
     ) -> Result<[[u8; 32]; HEIGHT], IndexerError> {
+        self.validate_proof_height::<HEIGHT>()?;
         let leaf_index = *self.low_element_indices.get(address_idx).ok_or_else(|| {
             IndexerError::MissingResult {
                 context: "reconstruct_proof".to_string(),
@@ -126,6 +129,7 @@ impl AddressQueueData {
         &self,
         address_range: std::ops::Range<usize>,
     ) -> Result<Vec<[[u8; 32]; HEIGHT]>, IndexerError> {
+        self.validate_proof_height::<HEIGHT>()?;
         let node_lookup = self.build_node_lookup();
         let mut proofs = Vec::with_capacity(address_range.len());
 
@@ -140,23 +144,24 @@ impl AddressQueueData {
     pub fn reconstruct_all_proofs<const HEIGHT: usize>(
         &self,
     ) -> Result<Vec<[[u8; 32]; HEIGHT]>, IndexerError> {
+        self.validate_proof_height::<HEIGHT>()?;
         self.reconstruct_proofs::<HEIGHT>(0..self.addresses.len())
     }
 
     fn build_node_lookup(&self) -> HashMap<u64, usize> {
-        self.nodes
-            .iter()
-            .copied()
-            .enumerate()
-            .map(|(idx, node)| (node, idx))
-            .collect()
+        let mut lookup = HashMap::with_capacity(self.nodes.len());
+        for (idx, node) in self.nodes.iter().copied().enumerate() {
+            lookup.entry(node).or_insert(idx);
+        }
+        lookup
     }
 
     fn reconstruct_proof_with_lookup<const HEIGHT: usize>(
         &self,
         address_idx: usize,
         node_lookup: &HashMap<u64, usize>,
     ) -> Result<[[u8; 32]; HEIGHT], IndexerError> {
+        self.validate_proof_height::<HEIGHT>()?;
         let leaf_index = *self.low_element_indices.get(address_idx).ok_or_else(|| {
             IndexerError::MissingResult {
                 context: "reconstruct_proof".to_string(),
@@ -209,6 +214,18 @@ impl AddressQueueData {
     fn encode_node_index(level: usize, position: u64) -> u64 {
         ((level as u64) << 56) | position
     }
+
+    fn validate_proof_height<const HEIGHT: usize>(&self) -> Result<(), IndexerError> {
+        if HEIGHT == Self::ADDRESS_TREE_HEIGHT {
+            return Ok(());
+        }
+
+        Err(IndexerError::InvalidParameters(format!(
+            "address queue proofs require HEIGHT={} but got HEIGHT={}",
+            Self::ADDRESS_TREE_HEIGHT,
+            HEIGHT
+        )))
+    }
 }
 
 #[cfg(test)]

sdk-libs/client/src/interface/load_accounts.rs

@@ -53,6 +53,9 @@ pub enum LoadAccountsError {
     #[error("Cold PDA at index {index} (pubkey {pubkey}) missing data")]
     MissingPdaCompressed { index: usize, pubkey: Pubkey },
 
+    #[error("Cold PDA (pubkey {pubkey}) missing data")]
+    MissingPdaCompressedData { pubkey: Pubkey },
+
     #[error("Cold ATA at index {index} (pubkey {pubkey}) missing data")]
     MissingAtaCompressed { index: usize, pubkey: Pubkey },
 
@@ -67,6 +70,7 @@ pub enum LoadAccountsError {
 }
 
 const MAX_ATAS_PER_IX: usize = 8;
+const MAX_PDAS_PER_IX: usize = 8;
 
 /// Build load instructions for cold accounts. Returns empty vec if all hot.
 ///
@@ -113,14 +117,18 @@ where
         })
         .collect();
 
-    let pda_hashes = collect_pda_hashes(&cold_pdas)?;
+    let pda_groups = group_pda_specs(&cold_pdas, MAX_PDAS_PER_IX);
+    let pda_hashes = pda_groups
+        .iter()
+        .map(|group| collect_pda_hashes(group))
+        .collect::<Result<Vec<_>, _>>()?;
     let ata_hashes = collect_ata_hashes(&cold_atas)?;
     let mint_hashes = collect_mint_hashes(&cold_mints)?;
 
     let (pda_proofs, ata_proofs, mint_proofs) = futures::join!(
-        fetch_proofs(&pda_hashes, indexer),
+        fetch_proof_batches(&pda_hashes, indexer),
         fetch_proofs_batched(&ata_hashes, MAX_ATAS_PER_IX, indexer),
-        fetch_proofs(&mint_hashes, indexer),
+        fetch_individual_proofs(&mint_hashes, indexer),
     );
 
     let pda_proofs = pda_proofs?;
@@ -136,18 +144,17 @@ where
 
     // 2. DecompressAccountsIdempotent for all cold PDAs (including token PDAs).
     //    Token PDAs are created on-chain via CPI inside DecompressVariant.
-    for (spec, proof) in cold_pdas.iter().zip(pda_proofs) {
+    for (group, proof) in pda_groups.into_iter().zip(pda_proofs) {
         out.push(build_pda_load(
-            &[spec],
+            &group,
             proof,
             fee_payer,
             compression_config,
         )?);
     }
 
     // 3. ATA loads (CreateAssociatedTokenAccount + Transfer2) - requires mint to exist
-    let ata_chunks: Vec<_> = cold_atas.chunks(MAX_ATAS_PER_IX).collect();
-    for (chunk, proof) in ata_chunks.into_iter().zip(ata_proofs) {
+    for (chunk, proof) in cold_atas.chunks(MAX_ATAS_PER_IX).zip(ata_proofs) {
         out.extend(build_ata_load(chunk, proof, fee_payer)?);
     }
 
@@ -195,23 +202,77 @@ fn collect_mint_hashes(ifaces: &[&AccountInterface]) -> Result<Vec<[u8; 32]>, Lo
         .collect()
 }
 
-async fn fetch_proofs<I: Indexer>(
+/// Groups already-ordered PDA specs into contiguous runs of the same program id.
+///
+/// This preserves input order rather than globally regrouping by program. Callers that
+/// want maximal batching across interleaved program ids should sort before calling.
+fn group_pda_specs<'a, V>(
+    specs: &[&'a PdaSpec<V>],
+    max_per_group: usize,
+) -> Vec<Vec<&'a PdaSpec<V>>> {
+    assert!(max_per_group > 0, "max_per_group must be non-zero");
+    if specs.is_empty() {
+        return Vec::new();
+    }
+
+    let mut groups = Vec::new();
+    let mut current = Vec::with_capacity(max_per_group);
+    let mut current_program: Option<Pubkey> = None;
+
+    for spec in specs {
+        let program_id = spec.program_id();
+        let should_split = current_program
+            .map(|existing| existing != program_id || current.len() >= max_per_group)
+            .unwrap_or(false);
+
+        if should_split {
+            groups.push(current);
+            current = Vec::with_capacity(max_per_group);
+        }
+
+        current_program = Some(program_id);
+        current.push(*spec);
+    }
+
+    if !current.is_empty() {
+        groups.push(current);
+    }
+
+    groups
+}
+
+async fn fetch_individual_proofs<I: Indexer>(
     hashes: &[[u8; 32]],
     indexer: &I,
 ) -> Result<Vec<ValidityProofWithContext>, IndexerError> {
     if hashes.is_empty() {
         return Ok(vec![]);
     }
-    let mut proofs = Vec::with_capacity(hashes.len());
-    for hash in hashes {
-        proofs.push(
-            indexer
-                .get_validity_proof(vec![*hash], vec![], None)
-                .await?
-                .value,
-        );
+
+    futures::future::try_join_all(hashes.iter().map(|hash| async move {
+        indexer
+            .get_validity_proof(vec![*hash], vec![], None)
+            .await
+            .map(|response| response.value)
+    }))
+    .await
+}
+
+async fn fetch_proof_batches<I: Indexer>(
+    hash_batches: &[Vec<[u8; 32]>],
+    indexer: &I,
+) -> Result<Vec<ValidityProofWithContext>, IndexerError> {
+    if hash_batches.is_empty() {
+        return Ok(vec![]);
     }
-    Ok(proofs)
+
+    futures::future::try_join_all(hash_batches.iter().map(|hashes| async move {
+        indexer
+            .get_validity_proof(hashes.clone(), vec![], None)
+            .await
+            .map(|response| response.value)
+    }))
+    .await
 }
 
 async fn fetch_proofs_batched<I: Indexer>(
@@ -222,16 +283,13 @@ async fn fetch_proofs_batched<I: Indexer>(
     if hashes.is_empty() {
         return Ok(vec![]);
     }
-    let mut proofs = Vec::with_capacity(hashes.len().div_ceil(batch_size));
-    for chunk in hashes.chunks(batch_size) {
-        proofs.push(
-            indexer
-                .get_validity_proof(chunk.to_vec(), vec![], None)
-                .await?
-                .value,
-        );
-    }
-    Ok(proofs)
+
+    let hash_batches = hashes
+        .chunks(batch_size)
+        .map(|chunk| chunk.to_vec())
+        .collect::<Vec<_>>();
+
+    fetch_proof_batches(&hash_batches, indexer).await
 }
 
 fn build_pda_load<V>(
@@ -262,11 +320,16 @@ where
     let hot_addresses: Vec<Pubkey> = specs.iter().map(|s| s.address()).collect();
     let cold_accounts: Vec<(CompressedAccount, V)> = specs
         .iter()
-        .map(|s| {
-            let compressed = s.compressed().expect("cold spec must have data").clone();
-            (compressed, s.variant.clone())
+        .map(|s| -> Result<_, LoadAccountsError> {
+            let compressed =
+                s.compressed()
+                    .cloned()
+                    .ok_or(LoadAccountsError::MissingPdaCompressedData {
+                        pubkey: s.address(),
+                    })?;
+            Ok((compressed, s.variant.clone()))
         })
-        .collect();
+        .collect::<Result<Vec<_>, _>>()?;
 
     let program_id = specs.first().map(|s| s.program_id()).unwrap_or_default();