Scaled UI Amount

import {
  lamports,
  createClient,
  generateKeyPairSigner,
  unwrapOption
} from "@solana/kit";
import { solanaRpc, rpcAirdrop } from "@solana/kit-plugin-rpc";
import { generatedPayer, airdropPayer } from "@solana/kit-plugin-signer";
import { getCreateAccountInstruction } from "@solana-program/system";
import {
  amountToUiAmountForMintWithoutSimulation,
  extension,
  fetchMint,
  findAssociatedTokenPda,
  getCreateAssociatedTokenInstructionAsync,
  getInitializeMintInstruction,
  getInitializeScaledUiAmountMintInstruction,
  getMintSize,
  getMintToCheckedInstruction,
  getUpdateMultiplierScaledUiMintInstruction,
  isExtension,
  TOKEN_2022_PROGRAM_ADDRESS
} from "@solana-program/token-2022";

const client = await createClient()
  .use(generatedPayer())
  .use(
    solanaRpc({
      rpcUrl: "http://localhost:8899",
      rpcSubscriptionsUrl: "ws://localhost:8900"
    })
  )
  .use(rpcAirdrop())
  .use(airdropPayer(lamports(1_000_000_000n)));

const mint = await generateKeyPairSigner();
const recipient = await generateKeyPairSigner();
const initialMultiplier = 5.0;
const updatedMultiplier = 10.0;
const tokenAmount = 1_000n;

const scaledUiAmountExtension = extension("ScaledUiAmountConfig", {
  authority: client.payer.address,
  multiplier: initialMultiplier,
  newMultiplierEffectiveTimestamp: 0n,
  newMultiplier: initialMultiplier
});
const mintSpace = BigInt(getMintSize([scaledUiAmountExtension]));
const mintRent = await client.rpc
  .getMinimumBalanceForRentExemption(mintSpace)
  .send();

await client.sendTransaction([
  getCreateAccountInstruction({
    payer: client.payer, // Account funding the new mint account.
    newAccount: mint, // New mint account to create.
    lamports: mintRent, // Lamports funding the mint account rent.
    space: mintSpace, // Account size in bytes for the mint plus ScaledUiAmountConfig.
    programAddress: TOKEN_2022_PROGRAM_ADDRESS // Program that owns the mint account.
  }),
  getInitializeScaledUiAmountMintInstruction({
    mint: mint.address, // Mint account that stores the ScaledUiAmountConfig extension.
    authority: client.payer.address, // Authority allowed to update the multiplier later.
    multiplier: initialMultiplier // Initial multiplier used for displayed UI amounts.
  }),
  getInitializeMintInstruction({
    mint: mint.address, // Mint account to initialize.
    decimals: 0, // Number of decimals for the token.
    mintAuthority: client.payer.address, // Authority allowed to mint new tokens.
    freezeAuthority: client.payer.address // Authority allowed to freeze token accounts.
  })
]);

const [tokenAccount] = await findAssociatedTokenPda({
  mint: mint.address,
  owner: recipient.address,
  tokenProgram: TOKEN_2022_PROGRAM_ADDRESS
});

await client.sendTransaction([
  await getCreateAssociatedTokenInstructionAsync({
    payer: client.payer, // Account funding the associated token account creation.
    mint: mint.address, // Mint for the associated token account.
    owner: recipient.address // Owner of the token account.
  }),
  getMintToCheckedInstruction({
    mint: mint.address, // Mint account that issues the tokens.
    token: tokenAccount, // Token account receiving the newly minted tokens.
    mintAuthority: client.payer, // Signer authorized to mint new tokens.
    amount: tokenAmount, // Token amount in base units.
    decimals: 0 // Decimals defined on the mint.
  })
]);

const calculatedUiAmountBeforeUpdate =
  await amountToUiAmountForMintWithoutSimulation(
    client.rpc,
    mint.address,
    tokenAmount
  );

const updateMultiplierInstruction = getUpdateMultiplierScaledUiMintInstruction({
  mint: mint.address, // Mint account that stores the ScaledUiAmountConfig extension.
  authority: client.payer, // Signer authorized to update the multiplier.
  multiplier: updatedMultiplier, // New multiplier used for displayed UI amounts.
  effectiveTimestamp: 0n // Unix timestamp when the new multiplier takes effect.
});

await client.sendTransaction([updateMultiplierInstruction]);

const calculatedUiAmountAfterUpdate =
  await amountToUiAmountForMintWithoutSimulation(
    client.rpc,
    mint.address,
    tokenAmount
  );

const mintAccount = await fetchMint(client.rpc, mint.address);
const scaledUiAmountConfig = (
  unwrapOption(mintAccount.data.extensions) ?? []
).find((item) => isExtension("ScaledUiAmountConfig", item));

console.log("Mint Address:", mint.address);
console.log("Token Account:", tokenAccount);
console.log(
  "Calculated UI Amount Before Update:",
  calculatedUiAmountBeforeUpdate
);
console.log(
  "Calculated UI Amount After Update:",
  calculatedUiAmountAfterUpdate
);
console.log("ScaledUiAmountConfig:", scaledUiAmountConfig);

import {
  Connection,
  Keypair,
  PublicKey,
  sendAndConfirmTransaction,
  SystemProgram,
  Transaction,
  LAMPORTS_PER_SOL
} from "@solana/web3.js";
import {
  ASSOCIATED_TOKEN_PROGRAM_ID,
  createAssociatedTokenAccountInstruction,
  createInitializeMintInstruction,
  createInitializeScaledUiAmountConfigInstruction,
  createMintToCheckedInstruction,
  createUpdateMultiplierDataInstruction,
  ExtensionType,
  getAssociatedTokenAddressSync,
  getMint,
  getMintLen,
  getScaledUiAmountConfig,
  TOKEN_2022_PROGRAM_ID
} from "@solana/spl-token";

async function calculateScaledUiAmount(
  connection: Connection,
  mintPublicKey: PublicKey,
  tokenAmount: bigint
) {
  const mintAccount = await getMint(
    connection,
    mintPublicKey,
    "confirmed",
    TOKEN_2022_PROGRAM_ID
  );
  const scaledUiAmountConfig = getScaledUiAmountConfig(mintAccount);
  if (!scaledUiAmountConfig) {
    throw new Error("ScaledUiAmountConfig not found");
  }

  const clockAccount = await connection.getParsedAccountInfo(
    new PublicKey("SysvarC1ock11111111111111111111111111111111")
  );
  if (
    !clockAccount.value ||
    typeof clockAccount.value.data !== "object" ||
    !("parsed" in clockAccount.value.data)
  ) {
    throw new Error("Failed to fetch clock sysvar");
  }

  const unixTimestamp = Number(
    clockAccount.value.data.parsed.info.unixTimestamp
  );
  const multiplier =
    unixTimestamp >=
    Number(scaledUiAmountConfig.newMultiplierEffectiveTimestamp)
      ? scaledUiAmountConfig.newMultiplier
      : scaledUiAmountConfig.multiplier;
  const scaledAmount = Math.trunc(Number(tokenAmount) * multiplier);
  const calculatedUiAmount = scaledAmount / 10 ** mintAccount.decimals;

  return {
    calculatedUiAmount: calculatedUiAmount.toString(),
    scaledUiAmountConfig
  };
}

const connection = new Connection("http://localhost:8899", "confirmed");

const feePayer = Keypair.generate();
const recipient = Keypair.generate();
const initialMultiplier = 5.0;
const updatedMultiplier = 10.0;
const tokenAmount = 1_000n;

const airdropSignature = await connection.requestAirdrop(
  feePayer.publicKey,
  2 * LAMPORTS_PER_SOL
);
await connection.confirmTransaction(airdropSignature, "confirmed");

const mint = Keypair.generate();
const mintLength = getMintLen([ExtensionType.ScaledUiAmountConfig]);
const mintRent = await connection.getMinimumBalanceForRentExemption(mintLength);

const tokenAccount = getAssociatedTokenAddressSync(
  mint.publicKey,
  recipient.publicKey,
  false,
  TOKEN_2022_PROGRAM_ID,
  ASSOCIATED_TOKEN_PROGRAM_ID
);

const createMintAccountInstruction = SystemProgram.createAccount({
  fromPubkey: feePayer.publicKey, // Account funding the new mint account.
  newAccountPubkey: mint.publicKey, // New mint account to create.
  space: mintLength, // Account size in bytes for the mint plus ScaledUiAmountConfig.
  lamports: mintRent, // Lamports funding the mint account rent.
  programId: TOKEN_2022_PROGRAM_ID // Program that owns the mint account.
});

const initializeScaledUiAmountInstruction =
  createInitializeScaledUiAmountConfigInstruction(
    mint.publicKey, // Mint account that stores the ScaledUiAmountConfig extension.
    feePayer.publicKey, // Authority allowed to update the multiplier later.
    initialMultiplier, // Initial multiplier used for displayed UI amounts.
    TOKEN_2022_PROGRAM_ID // Token program that owns the mint.
  );

const initializeMintInstruction = createInitializeMintInstruction(
  mint.publicKey, // Mint account to initialize.
  0, // Number of decimals for the token.
  feePayer.publicKey, // Authority allowed to mint new tokens.
  feePayer.publicKey, // Authority allowed to freeze token accounts.
  TOKEN_2022_PROGRAM_ID // Program that owns the mint account.
);

await sendAndConfirmTransaction(
  connection,
  new Transaction().add(
    createMintAccountInstruction,
    initializeScaledUiAmountInstruction,
    initializeMintInstruction
  ),
  [feePayer, mint]
);

const createTokenAccountInstruction = createAssociatedTokenAccountInstruction(
  feePayer.publicKey, // Account funding the associated token account creation.
  tokenAccount, // Associated token account address to create.
  recipient.publicKey, // Owner of the token account.
  mint.publicKey, // Mint for the associated token account.
  TOKEN_2022_PROGRAM_ID, // Token program that owns the token account.
  ASSOCIATED_TOKEN_PROGRAM_ID // Associated Token Program that creates the account.
);

const mintToTokenAccountInstruction = createMintToCheckedInstruction(
  mint.publicKey, // Mint account that issues the tokens.
  tokenAccount, // Token account receiving the newly minted tokens.
  feePayer.publicKey, // Signer authorized to mint new tokens.
  tokenAmount, // Token amount in base units.
  0, // Decimals defined on the mint.
  [], // Additional multisig signers.
  TOKEN_2022_PROGRAM_ID // Token program that owns the mint and token account.
);

await sendAndConfirmTransaction(
  connection,
  new Transaction().add(
    createTokenAccountInstruction,
    mintToTokenAccountInstruction
  ),
  [feePayer]
);

const { calculatedUiAmount: calculatedUiAmountBeforeUpdate } =
  await calculateScaledUiAmount(connection, mint.publicKey, tokenAmount);

const updateMultiplierInstruction = createUpdateMultiplierDataInstruction(
  mint.publicKey, // Mint account that stores the ScaledUiAmountConfig extension.
  feePayer.publicKey, // Signer authorized to update the multiplier.
  updatedMultiplier, // New multiplier used for displayed UI amounts.
  0n, // Unix timestamp when the new multiplier takes effect.
  [], // Additional multisig signers.
  TOKEN_2022_PROGRAM_ID // Token program that owns the mint.
);

await sendAndConfirmTransaction(
  connection,
  new Transaction().add(updateMultiplierInstruction),
  [feePayer]
);

const {
  calculatedUiAmount: calculatedUiAmountAfterUpdate,
  scaledUiAmountConfig
} = await calculateScaledUiAmount(connection, mint.publicKey, tokenAmount);

console.log("Mint Address:", mint.publicKey.toBase58());
console.log("Token Account:", tokenAccount.toBase58());
console.log(
  "Calculated UI Amount Before Update:",
  calculatedUiAmountBeforeUpdate
);
console.log(
  "Calculated UI Amount After Update:",
  calculatedUiAmountAfterUpdate
);
console.log("ScaledUiAmountConfig:", scaledUiAmountConfig);

use anyhow::Result;
use solana_client::nonblocking::rpc_client::RpcClient;
use solana_commitment_config::CommitmentConfig;
use solana_sdk::{
    sysvar::clock::{self, Clock},
    signature::{Keypair, Signer},
    transaction::Transaction,
};
use solana_system_interface::instruction::create_account;
use spl_associated_token_account_interface::{
    address::get_associated_token_address_with_program_id,
    instruction::create_associated_token_account,
};
use spl_token_2022_interface::{
    extension::{
        scaled_ui_amount::{instruction, ScaledUiAmountConfig},
        BaseStateWithExtensions, ExtensionType, StateWithExtensions,
    },
    instruction::{initialize_mint, mint_to_checked},
    state::Mint,
    ID as TOKEN_2022_PROGRAM_ID,
};

#[tokio::main]
async fn main() -> Result<()> {
    let client = RpcClient::new_with_commitment(
        String::from("http://localhost:8899"),
        CommitmentConfig::confirmed(),
    );
    let fee_payer = Keypair::new();
    let recipient = Keypair::new();
    let initial_multiplier = 5.0;
    let updated_multiplier = 10.0;
    let token_amount = 1_000u64;

    let airdrop_signature = client
        .request_airdrop(&fee_payer.pubkey(), 2_000_000_000)
        .await?;

    loop {
        let confirmed = client.confirm_transaction(&airdrop_signature).await?;
        if confirmed {
            break;
        }
    }

    let mint = Keypair::new();
    let mint_space =
        ExtensionType::try_calculate_account_len::<Mint>(&[ExtensionType::ScaledUiAmount])?;
    let mint_rent = client
        .get_minimum_balance_for_rent_exemption(mint_space)
        .await?;

    let create_mint_account_instruction = create_account(
        &fee_payer.pubkey(),    // Account funding the new mint account.
        &mint.pubkey(),         // New mint account to create.
        mint_rent,              // Lamports funding the mint account rent.
        mint_space as u64,      // Account size in bytes for the mint plus ScaledUiAmountConfig.
        &TOKEN_2022_PROGRAM_ID, // Program that owns the mint account.
    );

    let initialize_scaled_ui_amount_instruction = instruction::initialize(
        &TOKEN_2022_PROGRAM_ID, // Token program that owns the mint.
        &mint.pubkey(), // Mint account that stores the ScaledUiAmountConfig extension.
        Some(fee_payer.pubkey()), // Authority allowed to update the multiplier later.
        initial_multiplier, // Initial multiplier used for displayed UI amounts.
    )?;

    let initialize_mint_instruction = initialize_mint(
        &TOKEN_2022_PROGRAM_ID,    // Program that owns the mint account.
        &mint.pubkey(),            // Mint account to initialize.
        &fee_payer.pubkey(),       // Authority allowed to mint new tokens.
        Some(&fee_payer.pubkey()), // Authority allowed to freeze token accounts.
        0,                         // Number of decimals for the token.
    )?;

    let create_mint_transaction = Transaction::new_signed_with_payer(
        &[
            create_mint_account_instruction,
            initialize_scaled_ui_amount_instruction,
            initialize_mint_instruction,
        ],
        Some(&fee_payer.pubkey()),
        &[&fee_payer, &mint],
        client.get_latest_blockhash().await?,
    );

    client
        .send_and_confirm_transaction(&create_mint_transaction)
        .await?;

    let token_account = get_associated_token_address_with_program_id(
        &recipient.pubkey(),
        &mint.pubkey(),
        &TOKEN_2022_PROGRAM_ID,
    );

    let create_token_account_instruction = create_associated_token_account(
        &fee_payer.pubkey(), // Account funding the associated token account creation.
        &recipient.pubkey(), // Owner of the token account.
        &mint.pubkey(), // Mint for the associated token account.
        &TOKEN_2022_PROGRAM_ID, // Token program that owns the token account.
    );

    let mint_to_token_account_instruction = mint_to_checked(
        &TOKEN_2022_PROGRAM_ID, // Token program that owns the mint and token account.
        &mint.pubkey(), // Mint account that issues the tokens.
        &token_account, // Token account receiving the newly minted tokens.
        &fee_payer.pubkey(), // Signer authorized to mint new tokens.
        &[], // Additional multisig signers.
        token_amount, // Token amount in base units.
        0, // Decimals defined on the mint.
    )?;

    let create_token_account_transaction = Transaction::new_signed_with_payer(
        &[
            create_token_account_instruction,
            mint_to_token_account_instruction,
        ],
        Some(&fee_payer.pubkey()),
        &[&fee_payer],
        client.get_latest_blockhash().await?,
    );

    client
        .send_and_confirm_transaction(&create_token_account_transaction)
        .await?;

    let mint_account_before_update = client.get_account(&mint.pubkey()).await?;
    let mint_state_before_update = StateWithExtensions::<Mint>::unpack(&mint_account_before_update.data)?;
    let scaled_ui_amount_config_before_update =
        mint_state_before_update.get_extension::<ScaledUiAmountConfig>()?;

    let clock_account = client.get_account(&clock::ID).await?;
    let clock: Clock = clock_account.deserialize_data()?;
    let calculated_ui_amount_before_update = scaled_ui_amount_config_before_update
        .amount_to_ui_amount(
            token_amount,
            mint_state_before_update.base.decimals,
            clock.unix_timestamp,
        )
        .unwrap();

    let update_multiplier_instruction = instruction::update_multiplier(
        &TOKEN_2022_PROGRAM_ID, // Token program that owns the mint.
        &mint.pubkey(), // Mint account that stores the ScaledUiAmountConfig extension.
        &fee_payer.pubkey(), // Signer authorized to update the multiplier.
        &[], // Additional multisig signers.
        updated_multiplier, // New multiplier used for displayed UI amounts.
        0, // Unix timestamp when the new multiplier takes effect.
    )?;

    let update_multiplier_transaction = Transaction::new_signed_with_payer(
        &[update_multiplier_instruction],
        Some(&fee_payer.pubkey()),
        &[&fee_payer],
        client.get_latest_blockhash().await?,
    );

    client
        .send_and_confirm_transaction(&update_multiplier_transaction)
        .await?;

    let mint_account_after_update = client.get_account(&mint.pubkey()).await?;
    let mint_state_after_update = StateWithExtensions::<Mint>::unpack(&mint_account_after_update.data)?;
    let scaled_ui_amount_config_after_update =
        mint_state_after_update.get_extension::<ScaledUiAmountConfig>()?;
    let calculated_ui_amount_after_update = scaled_ui_amount_config_after_update
        .amount_to_ui_amount(
            token_amount,
            mint_state_after_update.base.decimals,
            clock.unix_timestamp,
        )
        .unwrap();

    println!("Mint Address: {}", mint.pubkey());
    println!("Token Account: {}", token_account);
    println!(
        "Calculated UI Amount Before Update: {}",
        calculated_ui_amount_before_update
    );
    println!(
        "Calculated UI Amount After Update: {}",
        calculated_ui_amount_after_update
    );
    println!("ScaledUiAmountConfig: {:#?}", scaled_ui_amount_config_after_update);

    Ok(())
}