feat(zair-sdk): support all gaptree modes in the integration api

crates/zair-nonmembership/src/sparse/sapling.rs

@@ -263,13 +263,17 @@ impl NonMembershipTree {
     }
 
     /// Returns the root of the tree.
-    ///
-    /// Returns `None` if the tree is empty.
     #[must_use]
     pub const fn root(&self) -> NonMembershipNode {
         self.cached_root
     }
 
+    /// Returns the root of the tree as bytes.
+    #[must_use]
+    pub const fn root_bytes(&self) -> [u8; 32] {
+        self.root().to_bytes()
+    }
+
     /// Returns the number of leaves in the tree.
     #[must_use]
     pub const fn leaf_count(&self) -> usize {
@@ -298,6 +302,21 @@ impl NonMembershipTree {
             .map_err(|e| MerklePathError::WitnessError(format!("{e:?}")))
     }
 
+    /// Generate a witness (authentication path) for a marked leaf.
+    ///
+    /// # Arguments
+    /// * `position` - The position of the leaf (must have been marked during construction).
+    ///
+    /// # Returns
+    /// A vector of sibling hashes represented as bytes from leaf to root.
+    ///
+    /// # Errors
+    /// Returns an error if the position is not marked or witness generation fails.
+    pub fn witness_bytes(&self, position: Position) -> Result<Vec<[u8; 32]>, MerklePathError> {
+        self.witness(position)
+            .map(|path| path.into_iter().map(|node| node.to_bytes()).collect())
+    }
+
     /// Verify that a leaf at the given position produces the expected root.
     ///
     /// # Arguments

crates/zair-sdk/src/api/claims.rs

@@ -7,7 +7,8 @@ use tracing::{debug, info};
 use zair_core::base::{Nullifier, Pool, SanitiseNullifiers};
 use zair_core::schema::proof_inputs::{ClaimInput, PublicInputs};
 use zair_nonmembership::{
-    MerklePathError, NonMembershipTree, OrchardNonMembershipTree, TreePosition,
+    MerklePathError, NonMembershipTree, OrchardGapTree, OrchardNonMembershipTree, SaplingGapTree,
+    TreePosition, map_orchard_user_positions, map_sapling_user_positions,
 };
 use zair_scan::ViewingKeys;
 
@@ -50,6 +51,10 @@ pub struct LoadedPoolData {
 
 /// Pool-specific non-membership tree variants.
 pub enum PoolMerkleTree {
+    /// Sapling pool gap tree (dense representation from precomputed file).
+    Sapling(SaplingGapTree),
+    /// Orchard pool gap tree (dense representation from precomputed file).
+    Orchard(OrchardGapTree),
     /// Sapling pool non-membership tree using sparse representation.
     SaplingSparse(NonMembershipTree),
     /// Orchard pool non-membership tree using sparse representation.
@@ -61,7 +66,9 @@ impl PoolMerkleTree {
     #[must_use]
     pub fn root_bytes(&self) -> [u8; 32] {
         match self {
-            Self::SaplingSparse(tree) => tree.root().to_bytes(),
+            Self::Sapling(tree) => tree.root_bytes(),
+            Self::Orchard(tree) => tree.root_bytes(),
+            Self::SaplingSparse(tree) => tree.root_bytes(),
             Self::OrchardSparse(tree) => tree.root_bytes(),
         }
     }
@@ -73,9 +80,9 @@ impl PoolMerkleTree {
     /// Returns an error if the witness cannot be generated for the given position.
     pub fn witness_bytes(&self, position: u64) -> Result<Vec<[u8; 32]>, MerklePathError> {
         match self {
-            Self::SaplingSparse(tree) => tree
-                .witness(position.into())
-                .map(|path| path.into_iter().map(|node| node.to_bytes()).collect()),
+            Self::Sapling(tree) => tree.witness_bytes(position),
+            Self::Orchard(tree) => tree.witness_bytes(position),
+            Self::SaplingSparse(tree) => tree.witness_bytes(position.into()),
             Self::OrchardSparse(tree) => tree.witness_bytes(position.into()),
         }
     }
@@ -142,6 +149,103 @@ pub async fn build_pool_merkle_tree_from_memory(
     })
 }
 
+/// Build the non-membership merkle tree for a pool from precomputed gap-tree bytes.
+///
+/// Corresponds to `GapTreeMode::None` - requires precomputed gap-tree file bytes.
+///
+/// # Errors
+///
+/// Returns an error if gap-tree bytes are invalid or user nullifier mapping fails.
+pub async fn build_pool_merkle_tree_from_file(
+    chain_nullifiers: SanitiseNullifiers,
+    user_nullifiers: SanitiseNullifiers,
+    gap_tree_bytes: &[u8],
+    pool: Pool,
+) -> Result<LoadedPoolData, ClaimsError> {
+    let use_orchard_tree = pool == Pool::Orchard;
+
+    info!(
+        count = chain_nullifiers.len(),
+        %pool,
+        "Loaded chain nullifiers"
+    );
+
+    info!(%pool, "Loading gap tree from precomputed bytes...");
+    let user_positions = if use_orchard_tree {
+        map_orchard_user_positions(&chain_nullifiers, &user_nullifiers)
+            .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?
+    } else {
+        map_sapling_user_positions(&chain_nullifiers, &user_nullifiers)
+            .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?
+    };
+
+    let tree = if use_orchard_tree {
+        PoolMerkleTree::Orchard(
+            OrchardGapTree::from_bytes(gap_tree_bytes)
+                .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?,
+        )
+    } else {
+        PoolMerkleTree::Sapling(
+            SaplingGapTree::from_bytes(gap_tree_bytes)
+                .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?,
+        )
+    };
+
+    info!(%pool, "Gap tree loaded");
+    Ok(LoadedPoolData {
+        tree,
+        user_nullifiers: user_positions,
+    })
+}
+
+/// Build gap tree from nullifier bytes and return serialized bytes (for persistence).
+///
+/// Corresponds to `GapTreeMode::Rebuild` - builds gap tree and returns serialized bytes
+/// that can be persisted to a file.
+///
+/// # Errors
+///
+/// Returns an error if tree construction fails.
+pub async fn build_gap_tree_and_serialize(
+    snapshot_nullifiers_bytes: &[u8],
+    pool: Pool,
+) -> Result<Vec<u8>, ClaimsError> {
+    let use_orchard_tree = pool == Pool::Orchard;
+    let chain_nullifiers_vec = zair_scan::read_nullifiers(snapshot_nullifiers_bytes)
+        .await
+        .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?;
+    let chain_nullifiers = SanitiseNullifiers::new(chain_nullifiers_vec);
+
+    info!(
+        count = chain_nullifiers.len(),
+        %pool,
+        "Building gap tree from snapshot nullifiers..."
+    );
+
+    let serialized = tokio::task::spawn_blocking(move || {
+        if use_orchard_tree {
+            let tree = OrchardGapTree::from_nullifiers_with_progress(
+                &chain_nullifiers,
+                |_current, _total| {},
+            )
+            .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?;
+            Ok::<_, ClaimsError>(tree.to_bytes())
+        } else {
+            let tree = SaplingGapTree::from_nullifiers_with_progress(
+                &chain_nullifiers,
+                |_current, _total| {},
+            )
+            .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))?;
+            Ok(tree.to_bytes())
+        }
+    })
+    .await
+    .map_err(|e| ClaimsError::MerkleTreeBuild(pool.to_string(), e.to_string()))??;
+
+    info!(%pool, "Gap tree built and serialized");
+    Ok(serialized)
+}
+
 /// Generate airdrop claims for the user's notes.
 ///
 /// This generic function works with any metadata type implementing `NoteMetadata`,

crates/zair-sdk/src/api/mod.rs

@@ -25,6 +25,7 @@ use crate::api::key::KeyError;
 use crate::api::prove::ProveError;
 use crate::api::scan::ScanError;
 use crate::api::sign::SignError;
+pub use crate::commands::GapTreeMode;
 
 /// Errors that can occur in the API pipeline.
 #[derive(Debug, Error)]
@@ -67,6 +68,9 @@ pub enum ApiError {
 /// * `orchard_params_bytes` - Halo2 params bytes (None if no Orchard claims)
 /// * `birthday_height` - earliest block for note scanning
 /// * `config` - airdrop configuration
+/// * `gap_tree_mode` - gap tree handling mode (None, Rebuild, or Sparse)
+/// * `sapling_gap_tree_bytes` - precomputed Sapling gap tree bytes (for None/Rebuild mode)
+/// * `orchard_gap_tree_bytes` - precomputed Orchard gap tree bytes (for None/Rebuild mode)
 /// * `message_hashes` - pre-computed message hashes to sign; per‑proof hashes (keyed by airdrop
 ///   nullifier) take precedence over the shared fallback
 ///
@@ -85,6 +89,9 @@ pub async fn run(
     orchard_params_bytes: Option<&[u8]>,
     birthday_height: u64,
     config: &AirdropConfiguration,
+    gap_tree_mode: GapTreeMode,
+    sapling_gap_tree_bytes: Option<&[u8]>,
+    orchard_gap_tree_bytes: Option<&[u8]>,
     message_hashes: Option<ResolvedMessageHashes>,
 ) -> Result<ClaimSubmission, ApiError> {
     let ufvk = key::derive_ufvk_from_seed(to_zcash_network(config.network), account_id, seed)
@@ -97,6 +104,9 @@ pub async fn run(
         &ufvk,
         birthday_height,
         config,
+        gap_tree_mode,
+        sapling_gap_tree_bytes,
+        orchard_gap_tree_bytes,
     )
     .await
     .map_err(ApiError::Scan)?;

crates/zair-sdk/src/api/scan.rs

@@ -16,7 +16,7 @@ use zcash_keys::keys::UnifiedFullViewingKey;
 
 use crate::api::claims::{
     GapTreeMode, OrchardPool, PoolClaimResult, PoolProcessor, SaplingPool,
-    build_pool_merkle_tree_from_memory, generate_claims,
+    build_pool_merkle_tree_from_file, build_pool_merkle_tree_from_memory, generate_claims,
 };
 use crate::common::{resolve_lightwalletd_url, to_zcash_network};
 
@@ -55,13 +55,14 @@ pub enum ScanError {
 }
 
 #[instrument(level = "debug", skip_all, fields(pool = %P::POOL))]
-pub(crate) async fn process_pool_claims_from_memory<P: PoolProcessor>(
+pub(crate) async fn process_pool_claims<P: PoolProcessor>(
     pool_enabled_in_config: bool,
     visitor: &AccountNotesVisitor,
     viewing_keys: &ViewingKeys,
     airdrop_config: &AirdropConfiguration,
     snapshot_nullifiers: Option<SanitiseNullifiers>,
     gap_tree_mode: GapTreeMode,
+    gap_tree_bytes: Option<&[u8]>,
 ) -> Result<PoolClaimResult<P::PrivateInputs>, ScanError> {
     if !pool_enabled_in_config {
         return Ok(PoolClaimResult::empty());
@@ -86,14 +87,33 @@ pub(crate) async fn process_pool_claims_from_memory<P: PoolProcessor>(
     };
 
     let user_nullifiers = SanitiseNullifiers::new(notes.keys().copied().collect());
-    let pool_data = build_pool_merkle_tree_from_memory(
-        snapshot_nullifiers,
-        user_nullifiers,
-        P::POOL,
-        gap_tree_mode,
-    )
-    .await
-    .map_err(|e| ScanError::ClaimGeneration(P::POOL.to_string(), e.to_string()))?;
+
+    let pool_data = match gap_tree_mode {
+        GapTreeMode::Sparse => build_pool_merkle_tree_from_memory(
+            snapshot_nullifiers,
+            user_nullifiers,
+            P::POOL,
+            gap_tree_mode,
+        )
+        .await
+        .map_err(|e| ScanError::ClaimGeneration(P::POOL.to_string(), e.to_string()))?,
+        GapTreeMode::None | GapTreeMode::Rebuild => {
+            let Some(gap_tree_bytes) = gap_tree_bytes else {
+                return Err(ScanError::ClaimGeneration(
+                    P::POOL.to_string(),
+                    format!("GapTreeMode::{gap_tree_mode:?} requires gap_tree_bytes"),
+                ));
+            };
+            build_pool_merkle_tree_from_file(
+                snapshot_nullifiers,
+                user_nullifiers,
+                gap_tree_bytes,
+                P::POOL,
+            )
+            .await
+            .map_err(|e| ScanError::ClaimGeneration(P::POOL.to_string(), e.to_string()))?
+        }
+    };
 
     let anchor = pool_data.tree.root_bytes();
     let Some(expected_root) = P::expected_root(airdrop_config) else {
@@ -128,6 +148,9 @@ pub(crate) async fn process_pool_claims_from_memory<P: PoolProcessor>(
 /// * `ufvk_encoded` - encoded Unified Full Viewing Key string
 /// * `birthday_height` - earliest block for note scanning
 /// * `config` - airdrop configuration
+/// * `gap_tree_mode` - gap tree handling mode (None, Rebuild, or Sparse)
+/// * `sapling_gap_tree_bytes` - precomputed Sapling gap tree bytes (for None/Rebuild mode)
+/// * `orchard_gap_tree_bytes` - precomputed Orchard gap tree bytes (for None/Rebuild mode)
 ///
 /// # Returns
 ///
@@ -144,6 +167,9 @@ pub async fn airdrop_claim_from_config(
     ufvk_encoded: &str,
     birthday_height: u64,
     config: &AirdropConfiguration,
+    gap_tree_mode: GapTreeMode,
+    sapling_gap_tree_bytes: Option<&[u8]>,
+    orchard_gap_tree_bytes: Option<&[u8]>,
 ) -> Result<AirdropClaimInputs, ScanError> {
     let network = to_zcash_network(config.network);
     let lightwalletd_url = resolve_lightwalletd_url(network, lightwalletd_url.as_deref());
@@ -199,23 +225,25 @@ pub async fn airdrop_claim_from_config(
         None
     };
 
-    let sapling_result = process_pool_claims_from_memory::<SaplingPool>(
+    let sapling_result = process_pool_claims::<SaplingPool>(
         config.sapling.is_some(),
         &visitor,
         &viewing_keys,
         config,
         sapling_nullifiers.map(SanitiseNullifiers::new),
-        GapTreeMode::Sparse,
+        gap_tree_mode,
+        sapling_gap_tree_bytes,
     )
     .await?;
 
-    let orchard_result = process_pool_claims_from_memory::<OrchardPool>(
+    let orchard_result = process_pool_claims::<OrchardPool>(
         config.orchard.is_some(),
         &visitor,
         &viewing_keys,
         config,
         orchard_nullifiers.map(SanitiseNullifiers::new),
-        GapTreeMode::Sparse,
+        gap_tree_mode,
+        orchard_gap_tree_bytes,
     )
     .await?;