瀏覽代碼

Use u64 behind the scenes for solana-tokens (#13815)

* Use u64 behind the scenes for Allocations

* Fixup readme

* Clippy and remove errant comments
Tyera Eulberg 5 年之前
父節點
當前提交
5e2d38227f
共有 7 個文件被更改,包括 258 次插入178 次删除
  1. 11 10
      tokens/README.md
  2. 3 2
      tokens/src/arg_parser.rs
  3. 2 2
      tokens/src/args.rs
  4. 180 125
      tokens/src/commands.rs
  5. 4 4
      tokens/src/db.rs
  6. 31 25
      tokens/src/spl_token.rs
  7. 27 10
      tokens/src/token_display.rs

+ 11 - 10
tokens/README.md

@@ -135,15 +135,16 @@ The Associated Token Account will be created, and funded by the fee_payer, if it
 does not already exist.
 
 Send SPL tokens to the recipients in `<RECIPIENTS_CSV>`.
+*NOTE:* the CSV expects SPL-token amounts in raw format (no decimals)
 
 Example recipients.csv:
 
 ```text
 recipient,amount
-CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,75.4
-C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,10
-7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,42.1
-7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,20
+CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,75400
+C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,10000
+7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,42100
+7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,20000
 ```
 
 You can check the status of the recipients before beginning a distribution. You
@@ -158,9 +159,9 @@ Example output:
 ```text
 Token: JDte736XZ1jGUtfAS32DLpBUWBR7WGSHy1hSZ36VRQ5V
 Recipient                                             Expected Balance            Actual Balance                Difference
-CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT                     75.40                      0.00                    -75.40
+CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT                    75.400                      0.000                   -75.400
 C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s                    10.000  Associated token account not yet created
-7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr                     42.10                      0.00                    -42.10
+7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr                    42.100                      0.000                   -42.100
 7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1                    20.000  Associated token account not yet created
 ```
 
@@ -195,10 +196,10 @@ Example updated recipients.csv:
 
 ```text
 recipient,amount
-CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,100
-C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,100
-7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,100
-7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,100
+CYRJWqiSjLitBAcRxPvWpgX3s5TvmN2SuRY3eEYypFvT,100000
+C56nwrDVFpPrqwGYsTgQxv1ZraTh81H14PV4RHvZe36s,100000
+7aHDubg5FBYj1SgmyBgU3ZJdtfuqYCQsJQK2pTR5JUqr,100000
+7qQPmVAQxEQ5djPDCtiEUrxaPf8wKtLG1m6SB1brejJ1,100000
 ```
 
 Using dry-run:

+ 3 - 2
tokens/src/arg_parser.rs

@@ -11,6 +11,7 @@ use solana_clap_utils::{
 };
 use solana_cli_config::CONFIG_FILE;
 use solana_remote_wallet::remote_wallet::maybe_wallet_manager;
+use solana_sdk::native_token::sol_to_lamports;
 use std::error::Error;
 use std::ffi::OsString;
 use std::process::exit;
@@ -360,7 +361,7 @@ fn parse_distribute_tokens_args(
         fee_payer,
         stake_args: None,
         spl_token_args: None,
-        transfer_amount: value_of(matches, "transfer_amount"),
+        transfer_amount: value_of(matches, "transfer_amount").map(sol_to_lamports),
     })
 }
 
@@ -423,7 +424,7 @@ fn parse_distribute_stake_args(
 
     let stake_args = StakeArgs {
         stake_account_address,
-        unlocked_sol: value_t_or_exit!(matches, "unlocked_sol", f64),
+        unlocked_sol: sol_to_lamports(value_t_or_exit!(matches, "unlocked_sol", f64)),
         stake_authority,
         withdraw_authority,
         lockup_authority,

+ 2 - 2
tokens/src/args.rs

@@ -9,11 +9,11 @@ pub struct DistributeTokensArgs {
     pub fee_payer: Box<dyn Signer>,
     pub stake_args: Option<StakeArgs>,
     pub spl_token_args: Option<SplTokenArgs>,
-    pub transfer_amount: Option<f64>,
+    pub transfer_amount: Option<u64>,
 }
 
 pub struct StakeArgs {
-    pub unlocked_sol: f64,
+    pub unlocked_sol: u64,
     pub stake_account_address: Pubkey,
     pub stake_authority: Box<dyn Signer>,
     pub withdraw_authority: Box<dyn Signer>,

+ 180 - 125
tokens/src/commands.rs

@@ -11,7 +11,9 @@ use indexmap::IndexMap;
 use indicatif::{ProgressBar, ProgressStyle};
 use pickledb::PickleDb;
 use serde::{Deserialize, Serialize};
-use solana_account_decoder::parse_token::{pubkey_from_spl_token_v2_0, spl_token_v2_0_pubkey};
+use solana_account_decoder::parse_token::{
+    pubkey_from_spl_token_v2_0, spl_token_v2_0_pubkey, token_amount_to_ui_amount,
+};
 use solana_client::{
     client_error::{ClientError, Result as ClientResult},
     rpc_client::RpcClient,
@@ -42,7 +44,7 @@ use std::{
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct Allocation {
     pub recipient: String,
-    pub amount: f64,
+    pub amount: u64,
     pub lockup_date: String,
 }
 
@@ -105,7 +107,7 @@ fn merge_allocations(allocations: &[Allocation]) -> Vec<Allocation> {
             .entry(&allocation.recipient)
             .or_insert(Allocation {
                 recipient: allocation.recipient.clone(),
-                amount: 0.0,
+                amount: 0,
                 lockup_date: "".to_string(),
             })
             .amount += allocation.amount;
@@ -134,11 +136,11 @@ fn apply_previous_transactions(
                 break;
             } else {
                 amount -= allocation.amount;
-                allocation.amount = 0.0;
+                allocation.amount = 0;
             }
         }
     }
-    allocations.retain(|x| x.amount > f64::EPSILON);
+    allocations.retain(|x| x.amount > 0);
 }
 
 fn transfer<S: Signer>(
@@ -168,7 +170,7 @@ fn distribution_instructions(
     if args.stake_args.is_none() && args.spl_token_args.is_none() {
         let from = args.sender_keypair.pubkey();
         let to = allocation.recipient.parse().unwrap();
-        let lamports = sol_to_lamports(allocation.amount);
+        let lamports = allocation.amount;
         let instruction = system_instruction::transfer(&from, &to, lamports);
         return vec![instruction];
     }
@@ -186,7 +188,7 @@ fn distribution_instructions(
     let mut instructions = stake_instruction::split(
         &stake_args.stake_account_address,
         &stake_authority,
-        sol_to_lamports(allocation.amount - unlocked_sol),
+        allocation.amount - unlocked_sol,
         &new_stake_account_address,
     );
 
@@ -230,7 +232,7 @@ fn distribution_instructions(
     instructions.push(system_instruction::transfer(
         &sender_pubkey,
         &recipient,
-        sol_to_lamports(unlocked_sol),
+        unlocked_sol,
     ));
 
     instructions
@@ -254,31 +256,32 @@ fn distribute_allocations(
             Some(allocation.lockup_date.parse::<DateTime<Utc>>().unwrap())
         };
 
-        let (decimals, do_create_associated_token_account) = if let Some(spl_token_args) =
-            &args.spl_token_args
-        {
-            let wallet_address = allocation.recipient.parse().unwrap();
-            let associated_token_address = get_associated_token_address(
-                &wallet_address,
-                &spl_token_v2_0_pubkey(&spl_token_args.mint),
-            );
-            let do_create_associated_token_account = client
-                .get_multiple_accounts(&[pubkey_from_spl_token_v2_0(&associated_token_address)])?
-                [0]
-            .is_none();
-            if do_create_associated_token_account {
-                created_accounts += 1;
-            }
-            (
-                spl_token_args.decimals as usize,
-                do_create_associated_token_account,
-            )
-        } else {
-            (9, false)
-        };
+        let (display_amount, decimals, do_create_associated_token_account) =
+            if let Some(spl_token_args) = &args.spl_token_args {
+                let wallet_address = allocation.recipient.parse().unwrap();
+                let associated_token_address = get_associated_token_address(
+                    &wallet_address,
+                    &spl_token_v2_0_pubkey(&spl_token_args.mint),
+                );
+                let do_create_associated_token_account =
+                    client.get_multiple_accounts(&[pubkey_from_spl_token_v2_0(
+                        &associated_token_address,
+                    )])?[0]
+                        .is_none();
+                if do_create_associated_token_account {
+                    created_accounts += 1;
+                }
+                (
+                    token_amount_to_ui_amount(allocation.amount, spl_token_args.decimals).ui_amount,
+                    spl_token_args.decimals as usize,
+                    do_create_associated_token_account,
+                )
+            } else {
+                (lamports_to_sol(allocation.amount), 9, false)
+            };
         println!(
             "{:<44}  {:>24.2$}",
-            allocation.recipient, allocation.amount, decimals
+            allocation.recipient, display_amount, decimals
         );
         let instructions = distribution_instructions(
             allocation,
@@ -361,8 +364,9 @@ fn distribute_allocations(
 
 fn read_allocations(
     input_csv: &str,
-    transfer_amount: Option<f64>,
+    transfer_amount: Option<u64>,
     require_lockup_heading: bool,
+    raw_amount: bool,
 ) -> io::Result<Vec<Allocation>> {
     let mut rdr = ReaderBuilder::new().trim(Trim::All).from_path(input_csv)?;
     let allocations = if let Some(amount) = transfer_amount {
@@ -379,7 +383,31 @@ fn read_allocations(
             })
             .collect()
     } else if require_lockup_heading {
-        rdr.deserialize().map(|entry| entry.unwrap()).collect()
+        let recipients: Vec<(String, f64, String)> = rdr
+            .deserialize()
+            .map(|recipient| recipient.unwrap())
+            .collect();
+        recipients
+            .into_iter()
+            .map(|(recipient, amount, lockup_date)| Allocation {
+                recipient,
+                amount: sol_to_lamports(amount),
+                lockup_date,
+            })
+            .collect()
+    } else if raw_amount {
+        let recipients: Vec<(String, u64)> = rdr
+            .deserialize()
+            .map(|recipient| recipient.unwrap())
+            .collect();
+        recipients
+            .into_iter()
+            .map(|(recipient, amount)| Allocation {
+                recipient,
+                amount,
+                lockup_date: "".to_string(),
+            })
+            .collect()
     } else {
         let recipients: Vec<(String, f64)> = rdr
             .deserialize()
@@ -389,7 +417,7 @@ fn read_allocations(
             .into_iter()
             .map(|(recipient, amount)| Allocation {
                 recipient,
-                amount,
+                amount: sol_to_lamports(amount),
                 lockup_date: "".to_string(),
             })
             .collect()
@@ -414,10 +442,15 @@ pub fn process_allocations(
         &args.input_csv,
         args.transfer_amount,
         require_lockup_heading,
+        args.spl_token_args.is_some(),
     )?;
-    let is_sol = args.spl_token_args.is_none();
 
-    let starting_total_tokens = Token::from(allocations.iter().map(|x| x.amount).sum(), is_sol);
+    let starting_total_tokens = allocations.iter().map(|x| x.amount).sum();
+    let starting_total_tokens = if let Some(spl_token_args) = &args.spl_token_args {
+        Token::spl_token(starting_total_tokens, spl_token_args.decimals)
+    } else {
+        Token::sol(starting_total_tokens)
+    };
     println!(
         "{} {}",
         style("Total in input_csv:").bold(),
@@ -437,8 +470,20 @@ pub fn process_allocations(
         return Ok(confirmations);
     }
 
-    let distributed_tokens = Token::from(transaction_infos.iter().map(|x| x.amount).sum(), is_sol);
-    let undistributed_tokens = Token::from(allocations.iter().map(|x| x.amount).sum(), is_sol);
+    let distributed_tokens = transaction_infos.iter().map(|x| x.amount).sum();
+    let undistributed_tokens = allocations.iter().map(|x| x.amount).sum();
+    let (distributed_tokens, undistributed_tokens) =
+        if let Some(spl_token_args) = &args.spl_token_args {
+            (
+                Token::spl_token(distributed_tokens, spl_token_args.decimals),
+                Token::spl_token(undistributed_tokens, spl_token_args.decimals),
+            )
+        } else {
+            (
+                Token::sol(distributed_tokens),
+                Token::sol(undistributed_tokens),
+            )
+        };
     println!("{} {}", style("Distributed:").bold(), distributed_tokens,);
     println!(
         "{} {}",
@@ -555,7 +600,7 @@ fn check_payer_balances(
     client: &RpcClient,
     args: &DistributeTokensArgs,
 ) -> Result<(), Error> {
-    let mut undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
+    let mut undistributed_tokens: u64 = allocations.iter().map(|x| x.amount).sum();
 
     let (_blockhash, fee_calculator) = client.get_recent_blockhash()?;
     let fees = fee_calculator
@@ -564,26 +609,22 @@ fn check_payer_balances(
         .unwrap();
 
     let (distribution_source, unlocked_sol_source) = if let Some(stake_args) = &args.stake_args {
-        let total_unlocked_sol = allocations.len() as f64 * stake_args.unlocked_sol;
+        let total_unlocked_sol = allocations.len() as u64 * stake_args.unlocked_sol;
         undistributed_tokens -= total_unlocked_sol;
         (
             stake_args.stake_account_address,
-            Some((
-                args.sender_keypair.pubkey(),
-                sol_to_lamports(total_unlocked_sol),
-            )),
+            Some((args.sender_keypair.pubkey(), total_unlocked_sol)),
         )
     } else {
         (args.sender_keypair.pubkey(), None)
     };
-    let allocation_lamports = sol_to_lamports(undistributed_tokens);
 
     if let Some((unlocked_sol_source, total_unlocked_sol)) = unlocked_sol_source {
         let staker_balance = client.get_balance(&distribution_source)?;
-        if staker_balance < allocation_lamports {
+        if staker_balance < undistributed_tokens {
             return Err(Error::InsufficientFunds(
                 vec![FundingSource::StakeAccount].into(),
-                lamports_to_sol(allocation_lamports),
+                lamports_to_sol(undistributed_tokens),
             ));
         }
         if args.fee_payer.pubkey() == unlocked_sol_source {
@@ -612,10 +653,10 @@ fn check_payer_balances(
         }
     } else if args.fee_payer.pubkey() == distribution_source {
         let balance = client.get_balance(&args.fee_payer.pubkey())?;
-        if balance < fees + allocation_lamports {
+        if balance < fees + undistributed_tokens {
             return Err(Error::InsufficientFunds(
                 vec![FundingSource::SystemAccount, FundingSource::FeePayer].into(),
-                lamports_to_sol(fees + allocation_lamports),
+                lamports_to_sol(fees + undistributed_tokens),
             ));
         }
     } else {
@@ -627,10 +668,10 @@ fn check_payer_balances(
             ));
         }
         let sender_balance = client.get_balance(&distribution_source)?;
-        if sender_balance < allocation_lamports {
+        if sender_balance < undistributed_tokens {
             return Err(Error::InsufficientFunds(
                 vec![FundingSource::SystemAccount].into(),
-                lamports_to_sol(allocation_lamports),
+                lamports_to_sol(undistributed_tokens),
             ));
         }
     }
@@ -638,7 +679,8 @@ fn check_payer_balances(
 }
 
 pub fn process_balances(client: &RpcClient, args: &BalancesArgs) -> Result<(), Error> {
-    let allocations: Vec<Allocation> = read_allocations(&args.input_csv, None, false)?;
+    let allocations: Vec<Allocation> =
+        read_allocations(&args.input_csv, None, false, args.spl_token_args.is_some())?;
     let allocations = merge_allocations(&allocations);
 
     let token = if let Some(spl_token_args) = &args.spl_token_args {
@@ -662,7 +704,7 @@ pub fn process_balances(client: &RpcClient, args: &BalancesArgs) -> Result<(), E
             print_token_balances(client, allocation, spl_token_args)?;
         } else {
             let address: Pubkey = allocation.recipient.parse().unwrap();
-            let expected = lamports_to_sol(sol_to_lamports(allocation.amount));
+            let expected = lamports_to_sol(allocation.amount);
             let actual = lamports_to_sol(client.get_balance(&address).unwrap());
             println!(
                 "{:<44}  {:>24.9}  {:>24.9}  {:>24.9}",
@@ -689,7 +731,7 @@ use tempfile::{tempdir, NamedTempFile};
 pub fn test_process_distribute_tokens_with_client(
     client: &RpcClient,
     sender_keypair: Keypair,
-    transfer_amount: Option<f64>,
+    transfer_amount: Option<u64>,
 ) {
     let fee_payer = Keypair::new();
     let transaction = transfer(
@@ -707,20 +749,21 @@ pub fn test_process_distribute_tokens_with_client(
         sol_to_lamports(1.0),
     );
 
-    let alice_pubkey = solana_sdk::pubkey::new_rand();
-    let allocation = Allocation {
-        recipient: alice_pubkey.to_string(),
-        amount: if let Some(amount) = transfer_amount {
-            amount
-        } else {
-            1000.0
-        },
-        lockup_date: "".to_string(),
+    let expected_amount = if let Some(amount) = transfer_amount {
+        amount
+    } else {
+        sol_to_lamports(1000.0)
     };
+    let alice_pubkey = solana_sdk::pubkey::new_rand();
     let allocations_file = NamedTempFile::new().unwrap();
     let input_csv = allocations_file.path().to_str().unwrap().to_string();
     let mut wtr = csv::WriterBuilder::new().from_writer(allocations_file);
-    wtr.serialize(&allocation).unwrap();
+    wtr.write_record(&["recipient", "amount"]).unwrap();
+    wtr.write_record(&[
+        alice_pubkey.to_string(),
+        lamports_to_sol(expected_amount).to_string(),
+    ])
+    .unwrap();
     wtr.flush().unwrap();
 
     let dir = tempdir().unwrap();
@@ -752,13 +795,9 @@ pub fn test_process_distribute_tokens_with_client(
         db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
     assert_eq!(transaction_infos.len(), 1);
     assert_eq!(transaction_infos[0].recipient, alice_pubkey);
-    let expected_amount = sol_to_lamports(allocation.amount);
-    assert_eq!(
-        sol_to_lamports(transaction_infos[0].amount),
-        expected_amount
-    );
+    assert_eq!(transaction_infos[0].amount, expected_amount);
 
-    assert_eq!(client.get_balance(&alice_pubkey).unwrap(), expected_amount,);
+    assert_eq!(client.get_balance(&alice_pubkey).unwrap(), expected_amount);
 
     check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
 
@@ -768,13 +807,9 @@ pub fn test_process_distribute_tokens_with_client(
         db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
     assert_eq!(transaction_infos.len(), 1);
     assert_eq!(transaction_infos[0].recipient, alice_pubkey);
-    let expected_amount = sol_to_lamports(allocation.amount);
-    assert_eq!(
-        sol_to_lamports(transaction_infos[0].amount),
-        expected_amount
-    );
+    assert_eq!(transaction_infos[0].amount, expected_amount);
 
-    assert_eq!(client.get_balance(&alice_pubkey).unwrap(), expected_amount,);
+    assert_eq!(client.get_balance(&alice_pubkey).unwrap(), expected_amount);
 
     check_output_file(&output_path, &db::open_db(&transaction_db, true).unwrap());
 }
@@ -817,16 +852,19 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
         .send_and_confirm_transaction_with_spinner(&transaction)
         .unwrap();
 
+    let expected_amount = sol_to_lamports(1000.0);
     let alice_pubkey = solana_sdk::pubkey::new_rand();
-    let allocation = Allocation {
-        recipient: alice_pubkey.to_string(),
-        amount: 1000.0,
-        lockup_date: "".to_string(),
-    };
     let file = NamedTempFile::new().unwrap();
     let input_csv = file.path().to_str().unwrap().to_string();
     let mut wtr = csv::WriterBuilder::new().from_writer(file);
-    wtr.serialize(&allocation).unwrap();
+    wtr.write_record(&["recipient", "amount", "lockup_date"])
+        .unwrap();
+    wtr.write_record(&[
+        alice_pubkey.to_string(),
+        lamports_to_sol(expected_amount).to_string(),
+        "".to_string(),
+    ])
+    .unwrap();
     wtr.flush().unwrap();
 
     let dir = tempdir().unwrap();
@@ -845,7 +883,7 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
         stake_authority: Box::new(stake_authority),
         withdraw_authority: Box::new(withdraw_authority),
         lockup_authority: None,
-        unlocked_sol: 1.0,
+        unlocked_sol: sol_to_lamports(1.0),
     };
     let args = DistributeTokensArgs {
         fee_payer: Box::new(fee_payer),
@@ -865,11 +903,7 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
         db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
     assert_eq!(transaction_infos.len(), 1);
     assert_eq!(transaction_infos[0].recipient, alice_pubkey);
-    let expected_amount = sol_to_lamports(allocation.amount);
-    assert_eq!(
-        sol_to_lamports(transaction_infos[0].amount),
-        expected_amount
-    );
+    assert_eq!(transaction_infos[0].amount, expected_amount);
 
     assert_eq!(
         client.get_balance(&alice_pubkey).unwrap(),
@@ -889,11 +923,7 @@ pub fn test_process_distribute_stake_with_client(client: &RpcClient, sender_keyp
         db::read_transaction_infos(&db::open_db(&transaction_db, true).unwrap());
     assert_eq!(transaction_infos.len(), 1);
     assert_eq!(transaction_infos[0].recipient, alice_pubkey);
-    let expected_amount = sol_to_lamports(allocation.amount);
-    assert_eq!(
-        sol_to_lamports(transaction_infos[0].amount),
-        expected_amount
-    );
+    assert_eq!(transaction_infos[0].amount, expected_amount);
 
     assert_eq!(
         client.get_balance(&alice_pubkey).unwrap(),
@@ -948,7 +978,7 @@ mod tests {
         } = TestValidator::run();
         let url = get_rpc_request_str(leader_data.rpc, false);
         let client = RpcClient::new_with_commitment(url, CommitmentConfig::recent());
-        test_process_distribute_tokens_with_client(&client, alice, Some(1.5));
+        test_process_distribute_tokens_with_client(&client, alice, Some(sol_to_lamports(1.5)));
 
         // Explicit cleanup, otherwise "pure virtual method called" crash in Docker
         server.close().unwrap();
@@ -978,7 +1008,7 @@ mod tests {
         let alice_pubkey = solana_sdk::pubkey::new_rand();
         let allocation = Allocation {
             recipient: alice_pubkey.to_string(),
-            amount: 42.0,
+            amount: 42,
             lockup_date: "".to_string(),
         };
         let file = NamedTempFile::new().unwrap();
@@ -988,12 +1018,27 @@ mod tests {
         wtr.flush().unwrap();
 
         assert_eq!(
-            read_allocations(&input_csv, None, false).unwrap(),
-            vec![allocation.clone()]
+            read_allocations(&input_csv, None, false, true).unwrap(),
+            vec![allocation]
+        );
+
+        let allocation_sol = Allocation {
+            recipient: alice_pubkey.to_string(),
+            amount: sol_to_lamports(42.0),
+            lockup_date: "".to_string(),
+        };
+
+        assert_eq!(
+            read_allocations(&input_csv, None, true, true).unwrap(),
+            vec![allocation_sol.clone()]
         );
         assert_eq!(
-            read_allocations(&input_csv, None, true).unwrap(),
-            vec![allocation]
+            read_allocations(&input_csv, None, false, false).unwrap(),
+            vec![allocation_sol.clone()]
+        );
+        assert_eq!(
+            read_allocations(&input_csv, None, true, false).unwrap(),
+            vec![allocation_sol]
         );
     }
 
@@ -1013,17 +1058,17 @@ mod tests {
         let expected_allocations = vec![
             Allocation {
                 recipient: pubkey0.to_string(),
-                amount: 42.0,
+                amount: sol_to_lamports(42.0),
                 lockup_date: "".to_string(),
             },
             Allocation {
                 recipient: pubkey1.to_string(),
-                amount: 43.0,
+                amount: sol_to_lamports(43.0),
                 lockup_date: "".to_string(),
             },
         ];
         assert_eq!(
-            read_allocations(&input_csv, None, false).unwrap(),
+            read_allocations(&input_csv, None, false, false).unwrap(),
             expected_allocations
         );
     }
@@ -1045,17 +1090,17 @@ mod tests {
         let expected_allocations = vec![
             Allocation {
                 recipient: pubkey0.to_string(),
-                amount: 42.0,
+                amount: sol_to_lamports(42.0),
                 lockup_date: "".to_string(),
             },
             Allocation {
                 recipient: pubkey1.to_string(),
-                amount: 43.0,
+                amount: sol_to_lamports(43.0),
                 lockup_date: "".to_string(),
             },
         ];
         assert_eq!(
-            read_allocations(&input_csv, None, true).unwrap(),
+            read_allocations(&input_csv, None, true, false).unwrap(),
             expected_allocations
         );
     }
@@ -1074,7 +1119,7 @@ mod tests {
         wtr.serialize(&pubkey2.to_string()).unwrap();
         wtr.flush().unwrap();
 
-        let amount = 1.5;
+        let amount = sol_to_lamports(1.5);
 
         let expected_allocations = vec![
             Allocation {
@@ -1094,7 +1139,7 @@ mod tests {
             },
         ];
         assert_eq!(
-            read_allocations(&input_csv, Some(amount), false).unwrap(),
+            read_allocations(&input_csv, Some(amount), false, false).unwrap(),
             expected_allocations
         );
     }
@@ -1106,18 +1151,18 @@ mod tests {
         let mut allocations = vec![
             Allocation {
                 recipient: alice.to_string(),
-                amount: 1.0,
+                amount: sol_to_lamports(1.0),
                 lockup_date: "".to_string(),
             },
             Allocation {
                 recipient: bob.to_string(),
-                amount: 1.0,
+                amount: sol_to_lamports(1.0),
                 lockup_date: "".to_string(),
             },
         ];
         let transaction_infos = vec![TransactionInfo {
             recipient: bob,
-            amount: 1.0,
+            amount: sol_to_lamports(1.0),
             ..TransactionInfo::default()
         }];
         apply_previous_transactions(&mut allocations, &transaction_infos);
@@ -1136,12 +1181,12 @@ mod tests {
         let lockup1 = "9999-12-31T23:59:59Z".to_string();
         let alice_alloc = Allocation {
             recipient: alice_pubkey.to_string(),
-            amount: 1.0,
+            amount: sol_to_lamports(1.0),
             lockup_date: "".to_string(),
         };
         let alice_alloc_lockup0 = Allocation {
             recipient: alice_pubkey.to_string(),
-            amount: 1.0,
+            amount: sol_to_lamports(1.0),
             lockup_date: lockup0.clone(),
         };
         let alice_info = TransactionInfo {
@@ -1184,7 +1229,7 @@ mod tests {
         let lockup_date_str = "2021-01-07T00:00:00Z";
         let allocation = Allocation {
             recipient: Pubkey::default().to_string(),
-            amount: 1.0,
+            amount: sol_to_lamports(1.0),
             lockup_date: lockup_date_str.to_string(),
         };
         let stake_account_address = solana_sdk::pubkey::new_rand();
@@ -1195,7 +1240,7 @@ mod tests {
             stake_authority: Box::new(Keypair::new()),
             withdraw_authority: Box::new(Keypair::new()),
             lockup_authority: Some(Box::new(lockup_authority)),
-            unlocked_sol: 1.0,
+            unlocked_sol: sol_to_lamports(1.0),
         };
         let args = DistributeTokensArgs {
             fee_payer: Box::new(Keypair::new()),
@@ -1235,7 +1280,7 @@ mod tests {
     }
 
     fn initialize_check_payer_balances_inputs(
-        allocation_amount: f64,
+        allocation_amount: u64,
         sender_keypair_file: &str,
         fee_payer: &str,
         stake_args: Option<StakeArgs>,
@@ -1291,7 +1336,7 @@ mod tests {
 
         // Fully funded payer
         let (allocations, mut args) = initialize_check_payer_balances_inputs(
-            allocation_amount,
+            sol_to_lamports(allocation_amount),
             &sender_keypair_file,
             &sender_keypair_file,
             None,
@@ -1408,7 +1453,7 @@ mod tests {
 
         // Fully funded payers
         let (allocations, mut args) = initialize_check_payer_balances_inputs(
-            allocation_amount,
+            sol_to_lamports(allocation_amount),
             &funded_payer_keypair_file,
             &sender_keypair_file,
             None,
@@ -1452,8 +1497,8 @@ mod tests {
     }
 
     fn initialize_stake_account(
-        stake_account_amount: f64,
-        unlocked_sol: f64,
+        stake_account_amount: u64,
+        unlocked_sol: u64,
         sender_keypair: &Keypair,
         client: &RpcClient,
     ) -> StakeArgs {
@@ -1472,7 +1517,7 @@ mod tests {
             &stake_account_address,
             &authorized,
             &lockup,
-            sol_to_lamports(stake_account_amount),
+            stake_account_amount,
         );
         let message = Message::new(&instructions, Some(&sender_keypair.pubkey()));
         let signers = [sender_keypair, &stake_account_keypair];
@@ -1521,11 +1566,16 @@ mod tests {
 
         let allocation_amount = 1000.0;
         let unlocked_sol = 1.0;
-        let stake_args = initialize_stake_account(allocation_amount, unlocked_sol, &alice, &client);
+        let stake_args = initialize_stake_account(
+            sol_to_lamports(allocation_amount),
+            sol_to_lamports(unlocked_sol),
+            &alice,
+            &client,
+        );
 
         // Fully funded payer & stake account
         let (allocations, mut args) = initialize_check_payer_balances_inputs(
-            allocation_amount,
+            sol_to_lamports(allocation_amount),
             &sender_keypair_file,
             &sender_keypair_file,
             Some(stake_args),
@@ -1536,7 +1586,7 @@ mod tests {
         let expensive_allocation_amount = 5000.0;
         let expensive_allocations = vec![Allocation {
             recipient: solana_sdk::pubkey::new_rand().to_string(),
-            amount: expensive_allocation_amount,
+            amount: sol_to_lamports(expensive_allocation_amount),
             lockup_date: "".to_string(),
         }];
         let err_result =
@@ -1642,7 +1692,12 @@ mod tests {
 
         let allocation_amount = 1000.0;
         let unlocked_sol = 1.0;
-        let stake_args = initialize_stake_account(allocation_amount, unlocked_sol, &alice, &client);
+        let stake_args = initialize_stake_account(
+            sol_to_lamports(allocation_amount),
+            sol_to_lamports(unlocked_sol),
+            &alice,
+            &client,
+        );
 
         let funded_payer = Keypair::new();
         let funded_payer_keypair_file = tmp_file_path("keypair_file", &funded_payer.pubkey());
@@ -1660,7 +1715,7 @@ mod tests {
 
         // Fully funded payers
         let (allocations, mut args) = initialize_check_payer_balances_inputs(
-            allocation_amount,
+            sol_to_lamports(allocation_amount),
             &funded_payer_keypair_file,
             &sender_keypair_file,
             Some(stake_args),

+ 4 - 4
tokens/src/db.rs

@@ -8,7 +8,7 @@ use std::{cmp::Ordering, fs, io, path::Path};
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct TransactionInfo {
     pub recipient: Pubkey,
-    pub amount: f64,
+    pub amount: u64,
     pub new_stake_account_address: Option<Pubkey>,
     pub finalized_date: Option<DateTime<Utc>>,
     pub transaction: Transaction,
@@ -19,7 +19,7 @@ pub struct TransactionInfo {
 #[derive(Serialize, Deserialize, Debug, Default, PartialEq)]
 struct SignedTransactionInfo {
     recipient: String,
-    amount: f64,
+    amount: u64,
     #[serde(skip_serializing_if = "String::is_empty", default)]
     new_stake_account_address: String,
     finalized_date: Option<DateTime<Utc>>,
@@ -34,7 +34,7 @@ impl Default for TransactionInfo {
         };
         Self {
             recipient: Pubkey::default(),
-            amount: 0.0,
+            amount: 0,
             new_stake_account_address: None,
             finalized_date: None,
             transaction,
@@ -104,7 +104,7 @@ pub fn read_transaction_infos(db: &PickleDb) -> Vec<TransactionInfo> {
 pub fn set_transaction_info(
     db: &mut PickleDb,
     recipient: &Pubkey,
-    amount: f64,
+    amount: u64,
     transaction: &Transaction,
     new_stake_account_address: Option<&Pubkey>,
     finalized: bool,

+ 31 - 25
tokens/src/spl_token.rs

@@ -75,7 +75,7 @@ pub fn build_spl_token_instructions(
         &associated_token_address,
         &spl_token_v2_0_pubkey(&args.sender_keypair.pubkey()),
         &[],
-        spl_token_amount(allocation.amount, spl_token_args.decimals),
+        allocation.amount,
         spl_token_args.decimals,
     )
     .unwrap();
@@ -94,8 +94,7 @@ pub fn check_spl_token_balances(
         .spl_token_args
         .as_ref()
         .expect("spl_token_args must be some");
-    let undistributed_tokens: f64 = allocations.iter().map(|x| x.amount).sum();
-    let allocation_amount = spl_token_amount(undistributed_tokens, spl_token_args.decimals);
+    let allocation_amount: u64 = allocations.iter().map(|x| x.amount).sum();
 
     let fee_calculator = client.get_recent_blockhash()?.1;
     let fees = fee_calculator
@@ -140,30 +139,37 @@ pub fn print_token_balances(
     let recipient_account = client
         .get_account(&pubkey_from_spl_token_v2_0(&associated_token_address))
         .unwrap_or_default();
-    let (actual, difference) =
-        if let Ok(recipient_token) = SplTokenAccount::unpack(&recipient_account.data) {
-            let actual = token_amount_to_ui_amount(recipient_token.amount, spl_token_args.decimals)
-                .ui_amount;
-            (
-                style(format!(
-                    "{:>24.1$}",
-                    actual, spl_token_args.decimals as usize
-                )),
-                format!(
-                    "{:>24.1$}",
-                    actual - expected,
-                    spl_token_args.decimals as usize
-                ),
-            )
-        } else {
-            (
-                style("Associated token account not yet created".to_string()).yellow(),
-                "".to_string(),
-            )
-        };
+    let (actual, difference) = if let Ok(recipient_token) =
+        SplTokenAccount::unpack(&recipient_account.data)
+    {
+        let actual_ui_amount =
+            token_amount_to_ui_amount(recipient_token.amount, spl_token_args.decimals).ui_amount;
+        let expected_ui_amount =
+            token_amount_to_ui_amount(expected, spl_token_args.decimals).ui_amount;
+        (
+            style(format!(
+                "{:>24.1$}",
+                actual_ui_amount, spl_token_args.decimals as usize
+            )),
+            format!(
+                "{:>24.1$}",
+                actual_ui_amount - expected_ui_amount,
+                spl_token_args.decimals as usize
+            ),
+        )
+    } else {
+        (
+            style("Associated token account not yet created".to_string()).yellow(),
+            "".to_string(),
+        )
+    };
     println!(
         "{:<44}  {:>24.4$}  {:>24}  {:>24}",
-        allocation.recipient, expected, actual, difference, spl_token_args.decimals as usize
+        allocation.recipient,
+        token_amount_to_ui_amount(expected, spl_token_args.decimals).ui_amount,
+        actual,
+        difference,
+        spl_token_args.decimals as usize
     );
     Ok(())
 }

+ 27 - 10
tokens/src/token_display.rs

@@ -1,3 +1,5 @@
+use solana_account_decoder::parse_token::token_amount_to_ui_amount;
+use solana_sdk::native_token::lamports_to_sol;
 use std::{
     fmt::{Debug, Display, Formatter, Result},
     ops::Add,
@@ -12,25 +14,39 @@ pub enum TokenType {
 }
 
 pub struct Token {
-    amount: f64,
+    amount: u64,
+    decimals: u8,
     token_type: TokenType,
 }
 
 impl Token {
     fn write_with_symbol(&self, f: &mut Formatter) -> Result {
         match &self.token_type {
-            TokenType::Sol => write!(f, "{}{}", SOL_SYMBOL, self.amount,),
-            TokenType::SplToken => write!(f, "{} tokens", self.amount,),
+            TokenType::Sol => {
+                let amount = lamports_to_sol(self.amount);
+                write!(f, "{}{}", SOL_SYMBOL, amount)
+            }
+            TokenType::SplToken => {
+                let amount = token_amount_to_ui_amount(self.amount, self.decimals).ui_amount;
+                write!(f, "{} tokens", amount)
+            }
         }
     }
 
-    pub fn from(amount: f64, is_sol: bool) -> Self {
-        let token_type = if is_sol {
-            TokenType::Sol
-        } else {
-            TokenType::SplToken
-        };
-        Self { amount, token_type }
+    pub fn sol(amount: u64) -> Self {
+        Self {
+            amount,
+            decimals: 9,
+            token_type: TokenType::Sol,
+        }
+    }
+
+    pub fn spl_token(amount: u64, decimals: u8) -> Self {
+        Self {
+            amount,
+            decimals,
+            token_type: TokenType::SplToken,
+        }
     }
 }
 
@@ -53,6 +69,7 @@ impl Add for Token {
         if self.token_type == other.token_type {
             Self {
                 amount: self.amount + other.amount,
+                decimals: self.decimals,
                 token_type: self.token_type,
             }
         } else {