Storage subset merkle tree onchain

programs/merkle-distributor/Cargo.toml

@@ -18,7 +18,7 @@ default = []
 localnet = []
 
 [dependencies]
-anchor-lang = "0.28.0"
+anchor-lang = { version = "0.28.0", features = ["init-if-needed"] }
 anchor-spl = "0.28.0"
 bytemuck = "1.14.0"
 jito-merkle-verify = { path = "../../verify" }

programs/merkle-distributor/src/error.rs

@@ -59,4 +59,10 @@ pub enum ErrorCode {
     InvalidLocker,
     #[msg("Escrow is not max lock")]
     EscrowIsNotMaxLock,
+    #[msg("Invalid remaining accounts")]
+    InvalidRemainingAccounts,
+    #[msg("Invalid account")]
+    InvalidAccount,
+    #[msg("Canopy root miss match with real root")]
+    CanopyRootMissMatch,
 }

programs/merkle-distributor/src/instructions/admin/create_canopy_tree.rs

@@ -0,0 +1,51 @@
+use crate::{
+    error::ErrorCode,
+    state::{canopy_tree::CanopyTree, merkle_distributor::MerkleDistributor},
+};
+use anchor_lang::prelude::*;
+
+#[derive(Accounts)]
+#[instruction(depth: u8)]
+pub struct CreateCanopyTree<'info> {
+    /// [CanopyTree]
+    #[account(
+        init,
+        seeds = [
+            b"CanopyTree".as_ref(),
+            distributor.key().to_bytes().as_ref(),
+        ],
+        bump,
+        space = CanopyTree::space(depth as usize),
+        payer = payer
+    )]
+    pub canopy_tree: Account<'info, CanopyTree>,
+
+    /// The [MerkleDistributor].
+    pub distributor: AccountLoader<'info, MerkleDistributor>,
+
+    /// Payer wallet, responsible for creating the distributor and paying for the transaction.
+    #[account(mut)]
+    pub payer: Signer<'info>,
+
+    /// The [System] program.
+    pub system_program: Program<'info, System>,
+}
+
+pub fn handle_create_canopy_tree(
+    ctx: Context<CreateCanopyTree>,
+    depth: u8,
+    root: [u8; 32],
+    canopy_nodes: Vec<[u8; 32]>,
+) -> Result<()> {
+    let canopy_tree = &mut ctx.accounts.canopy_tree;
+
+    let verify_canopy_root = canopy_tree.verify_canopy_root(root, canopy_nodes.clone());
+    require!(verify_canopy_root, ErrorCode::CanopyRootMissMatch);
+
+    canopy_tree.root = root;
+    canopy_tree.depth = depth;
+    canopy_tree.nodes = canopy_nodes;
+    canopy_tree.distributor = ctx.accounts.distributor.key();
+
+    Ok(())
+}

programs/merkle-distributor/src/instructions/admin/fund_distributor_root.rs

@@ -0,0 +1,73 @@
+use anchor_lang::{
+    accounts::{account::Account, program::Program},
+    context::Context,
+    prelude::*,
+    Accounts, Result, ToAccountInfo,
+};
+use anchor_spl::{
+    associated_token::AssociatedToken,
+    token::{self, Mint, Token, TokenAccount},
+};
+
+use crate::state::distributor_root::DistributorRoot;
+
+/// Accounts required for distributing tokens from the parent vault to distributor vaults.
+#[derive(Accounts)]
+pub struct FundDistributorRoot<'info> {
+    /// The [DistributorRoot]
+    #[account(mut, has_one = mint)]
+    pub distributor_root: AccountLoader<'info, DistributorRoot>,
+
+    /// Distributor root vault
+    #[account(
+        init_if_needed,
+        associated_token::mint = mint,
+        associated_token::authority = distributor_root,
+        payer = payer
+    )]
+    pub distributor_root_vault: Account<'info, TokenAccount>,
+
+    /// The mint to distribute.
+    pub mint: Account<'info, Mint>,
+
+    /// Payer.
+    #[account(mut)]
+    pub payer: Signer<'info>,
+
+    /// Payer Token Account.
+    #[account(mut)]
+    pub payer_token: Account<'info, TokenAccount>,
+
+    /// The [System] program.
+    pub system_program: Program<'info, System>,
+
+    /// The [Token] program.
+    pub token_program: Program<'info, Token>,
+
+    // Associated token program.
+    pub associated_token_program: Program<'info, AssociatedToken>,
+}
+
+pub fn handle_fund_distributor_root(
+    ctx: Context<FundDistributorRoot>,
+    max_amount: u64,
+) -> Result<()> {
+    let fund_amount = {
+        let mut distributor_root = ctx.accounts.distributor_root.load_mut()?;
+        distributor_root.get_and_set_fund_amount(max_amount)?
+    };
+
+    token::transfer(
+        CpiContext::new(
+            ctx.accounts.token_program.to_account_info(),
+            token::Transfer {
+                from: ctx.accounts.payer_token.to_account_info(),
+                to: ctx.accounts.distributor_root_vault.to_account_info(),
+                authority: ctx.accounts.payer.to_account_info(),
+            },
+        ),
+        fund_amount,
+    )?;
+
+    Ok(())
+}

programs/merkle-distributor/src/instructions/admin/fund_merkle_distributor_from_root.rs

@@ -0,0 +1,78 @@
+use anchor_lang::{
+    accounts::{account::Account, program::Program},
+    context::{Context, CpiContext},
+    prelude::*,
+    Accounts, Result, ToAccountInfo,
+};
+use anchor_spl::token::{self, Token, TokenAccount};
+
+use crate::state::{distributor_root::DistributorRoot, merkle_distributor::MerkleDistributor};
+
+/// Accounts required for distributing tokens from the parent vault to distributor vaults.
+#[derive(Accounts)]
+pub struct FundMerkleDisitributorFromRoot<'info> {
+    /// The [DistributorRoot].
+    pub distributor_root: AccountLoader<'info, DistributorRoot>,
+
+    /// Distributor root vault containing the tokens to distribute to distributor vault.
+    #[account(
+        mut,
+        associated_token::mint = distributor_root.load()?.mint,
+        associated_token::authority = distributor_root.key(),
+        address = distributor_root.load()?.distributor_root_vault,
+    )]
+    pub distributor_root_vault: Account<'info, TokenAccount>,
+
+    /// The [MerkleDistributor].
+    #[account(mut, constraint = distributor.load()?.distributor_root == distributor_root.key())]
+    pub distributor: AccountLoader<'info, MerkleDistributor>,
+
+    /// Distributor vault
+    #[account(
+        mut,
+        associated_token::mint = distributor.load()?.mint,
+        associated_token::authority = distributor.key(),
+    )]
+    pub distributor_vault: Account<'info, TokenAccount>,
+
+    /// SPL [Token] program.
+    pub token_program: Program<'info, Token>,
+}
+
+/// Handles the distribution of tokens from the parent vault to multiple distributor vaults.
+pub fn handle_fund_merkle_distributor_from_root<'info>(
+    ctx: Context<'_, '_, '_, 'info, FundMerkleDisitributorFromRoot<'info>>,
+) -> Result<()> {
+    let distributor_root = ctx.accounts.distributor_root.load()?;
+    let signer = distributor_root.signer();
+    let seeds = signer.seeds();
+
+    let mut distributor_state = ctx.accounts.distributor.load_mut()?;
+
+    // Check distributor has been funded token
+    if distributor_state.funded_amount == 0 {
+        let fund_amount = distributor_state.max_total_claim;
+        token::transfer(
+            CpiContext::new(
+                ctx.accounts.token_program.to_account_info(),
+                token::Transfer {
+                    from: ctx.accounts.distributor_root_vault.to_account_info(),
+                    to: ctx.accounts.distributor_vault.to_account_info(),
+                    authority: ctx.accounts.distributor_root.to_account_info(),
+                },
+            )
+            .with_signer(&[&seeds[..]]),
+            fund_amount,
+        )?;
+
+        distributor_state.accumulate_funded_amount(fund_amount)?;
+
+        msg!(
+            "Funded {} tokens to distributor version {}.",
+            fund_amount,
+            distributor_state.version
+        );
+    }
+
+    Ok(())
+}

programs/merkle-distributor/src/instructions/admin/mod.rs

@@ -5,6 +5,10 @@ pub mod new_distributor;
 pub mod set_activation_point;
 pub mod set_admin;
 pub mod set_clawback_receiver;
+pub mod new_distributor_root;
+pub mod fund_merkle_distributor_from_root;
+pub mod fund_distributor_root;
+pub mod create_canopy_tree;
 
 pub use clawback::*;
 pub use close_claim_status::*;
@@ -15,3 +19,7 @@ pub use set_admin::*;
 pub use set_clawback_receiver::*;
 pub mod set_operator;
 pub use set_operator::*;
+pub use new_distributor_root::*;
+pub use fund_merkle_distributor_from_root::*;
+pub use fund_distributor_root::*;
+pub use create_canopy_tree::*;

programs/merkle-distributor/src/instructions/admin/new_distributor.rs

@@ -1,10 +1,12 @@
 use crate::error::ErrorCode::ArithmeticError;
+use crate::state::distributor_root::DistributorRoot;
 use crate::state::merkle_distributor::{ActivationType, ClaimType};
 use crate::{
     error::ErrorCode,
     state::merkle_distributor::{AirdropBonus, MerkleDistributor},
 };
 use anchor_lang::{account, context::Context, prelude::*, Accounts, Key, ToAccountInfo};
+use anchor_spl::associated_token::AssociatedToken;
 use anchor_spl::token::{Mint, Token, TokenAccount};
 
 #[cfg(feature = "localnet")]
@@ -13,10 +15,9 @@ const SECONDS_PER_DAY: i64 = 0;
 #[cfg(not(feature = "localnet"))]
 const SECONDS_PER_DAY: i64 = 24 * 3600; // 24 hours * 3600 seconds
 
-#[derive(AnchorSerialize, AnchorDeserialize, InitSpace)]
+#[derive(AnchorSerialize, AnchorDeserialize, Debug)]
 pub struct NewDistributorParams {
     pub version: u64,
-    pub root: [u8; 32],
     pub total_claim: u64,
     pub max_num_nodes: u64,
     pub start_vesting_ts: i64,
@@ -29,7 +30,7 @@ pub struct NewDistributorParams {
     pub bonus_vesting_duration: u64,
     pub claim_type: u8,
     pub operator: Pubkey,
-    pub locker: Pubkey,
+    pub locker: Pubkey
 }
 
 impl NewDistributorParams {
@@ -123,10 +124,14 @@ pub struct NewDistributor<'info> {
         ],
         bump,
         space = 8 + MerkleDistributor::INIT_SPACE,
-        payer = admin
+        payer = payer
     )]
     pub distributor: AccountLoader<'info, MerkleDistributor>,
 
+    /// The [DistributorRoot].
+    #[account(mut)]
+    pub distributor_root: AccountLoader<'info, DistributorRoot>,
+
     /// Base key of the distributor.
     pub base: Signer<'info>,
 
@@ -140,21 +145,28 @@ pub struct NewDistributor<'info> {
     /// Token vault
     /// Should create previously
     #[account(
+        init_if_needed,
         associated_token::mint = mint,
         associated_token::authority=distributor,
+        payer = payer
     )]
     pub token_vault: Account<'info, TokenAccount>,
 
-    /// Admin wallet, responsible for creating the distributor and paying for the transaction.
-    /// Also has the authority to set the clawback receiver and change itself.
+    /// CHECK: This account is not use to read or write
+    pub admin: UncheckedAccount<'info>,
+
+    /// Payer wallet, responsible for creating the distributor and paying for the transaction.
     #[account(mut)]
-    pub admin: Signer<'info>,
+    pub payer: Signer<'info>,
 
     /// The [System] program.
     pub system_program: Program<'info, System>,
 
     /// The [Token] program.
     pub token_program: Program<'info, Token>,
+
+    // Associated token program.
+    pub associated_token_program: Program<'info, AssociatedToken>,
 }
 
 /// Creates a new [MerkleDistributor].
@@ -172,12 +184,9 @@ pub fn handle_new_distributor(
     params: &NewDistributorParams,
 ) -> Result<()> {
     params.validate()?;
-
     let mut distributor = ctx.accounts.distributor.load_init()?;
-
     distributor.bump = *ctx.bumps.get("distributor").unwrap();
     distributor.version = params.version;
-    distributor.root = params.root;
     distributor.mint = ctx.accounts.mint.key();
     distributor.token_vault = ctx.accounts.token_vault.key();
     distributor.max_total_claim = params.get_max_total_claim()?;
@@ -200,6 +209,7 @@ pub fn handle_new_distributor(
     distributor.activation_type = params.activation_type;
     distributor.operator = params.operator;
     distributor.locker = params.locker;
+    distributor.distributor_root = ctx.accounts.distributor_root.key();
 
     // Note: might get truncated, do not rely on
     msg! {
@@ -220,5 +230,11 @@ pub fn handle_new_distributor(
             distributor.claim_type,
     };
 
+    drop(distributor);
+
+    // increase total distributor created
+    let mut distributor_root = ctx.accounts.distributor_root.load_mut()?;
+    distributor_root.update_new_distributor()?;
+
     Ok(())
 }