token-swap: Improve pool token supply on initialization, deposit, and withdrawal (#508)

* token-swap: Add token supply in invariant calculation

* Refactor state classes into curve components for future use
* Align pool initialization with Uniswap using geometric mean of token
  amounts
* Fix deposit and withdraw instruction to work as a proportion of pool
  tokens
* Add math utilities to calculate the geometric mean with u64

* Improve variable names

* Use a fixed starting pool size

* Run cargo fmt

* Update js tests with new pool numbers

* Run linting

* Remove math

* Fix BN type issues found by flow

Author: joncinque

bpf-sdk-install.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 set -e
 
-channel=${1:-v1.3.9}
+channel=${1:-v1.3.12}
 installDir="$(dirname "$0")"/bin
 cacheDir=~/.cache/solana-bpf-sdk/"$channel"
 

token-swap/js/cli/token-swap-test.js

@@ -39,6 +39,10 @@ let tokenAccountB: PublicKey;
 const BASE_AMOUNT = 1000;
 // Amount passed to instructions
 const USER_AMOUNT = 100;
+// Pool token amount minted on init
+const DEFAULT_POOL_TOKEN_AMOUNT = 1000000000;
+// Pool token amount to withdraw / deposit
+const POOL_TOKEN_AMOUNT = 1000000;
 
 function assert(condition, message) {
   if (!condition) {
@@ -219,16 +223,23 @@ export async function createTokenSwap(): Promise<void> {
 }
 
 export async function deposit(): Promise<void> {
+  const poolMintInfo = await tokenPool.getMintInfo();
+  const supply = poolMintInfo.supply.toNumber();
+  const swapTokenA = await mintA.getAccountInfo(tokenAccountA);
+  const tokenA = (swapTokenA.amount.toNumber() * POOL_TOKEN_AMOUNT) / supply;
+  const swapTokenB = await mintB.getAccountInfo(tokenAccountB);
+  const tokenB = (swapTokenB.amount.toNumber() * POOL_TOKEN_AMOUNT) / supply;
+
   console.log('Creating depositor token a account');
-  let userAccountA = await mintA.createAccount(owner.publicKey);
-  await mintA.mintTo(userAccountA, owner, [], USER_AMOUNT);
-  await mintA.approve(userAccountA, authority, owner, [], USER_AMOUNT);
+  const userAccountA = await mintA.createAccount(owner.publicKey);
+  await mintA.mintTo(userAccountA, owner, [], tokenA);
+  await mintA.approve(userAccountA, authority, owner, [], tokenA);
   console.log('Creating depositor token b account');
-  let userAccountB = await mintB.createAccount(owner.publicKey);
-  await mintB.mintTo(userAccountB, owner, [], USER_AMOUNT);
-  await mintB.approve(userAccountB, authority, owner, [], USER_AMOUNT);
+  const userAccountB = await mintB.createAccount(owner.publicKey);
+  await mintB.mintTo(userAccountB, owner, [], tokenB);
+  await mintB.approve(userAccountB, authority, owner, [], tokenB);
   console.log('Creating depositor pool token account');
-  let newAccountPool = await tokenPool.createAccount(owner.publicKey);
+  const newAccountPool = await tokenPool.createAccount(owner.publicKey);
   const [tokenProgramId] = await GetPrograms(connection);
 
   console.log('Depositing into swap');
@@ -241,7 +252,7 @@ export async function deposit(): Promise<void> {
     tokenPool.publicKey,
     newAccountPool,
     tokenProgramId,
-    USER_AMOUNT,
+    POOL_TOKEN_AMOUNT,
   );
 
   let info;
@@ -250,21 +261,34 @@ export async function deposit(): Promise<void> {
   info = await mintB.getAccountInfo(userAccountB);
   assert(info.amount.toNumber() == 0);
   info = await mintA.getAccountInfo(tokenAccountA);
-  assert(info.amount.toNumber() == BASE_AMOUNT + USER_AMOUNT);
+  assert(info.amount.toNumber() == BASE_AMOUNT + tokenA);
   info = await mintB.getAccountInfo(tokenAccountB);
-  assert(info.amount.toNumber() == BASE_AMOUNT + USER_AMOUNT);
+  assert(info.amount.toNumber() == BASE_AMOUNT + tokenB);
   info = await tokenPool.getAccountInfo(newAccountPool);
-  assert(info.amount.toNumber() == USER_AMOUNT);
+  assert(info.amount.toNumber() == POOL_TOKEN_AMOUNT);
 }
 
 export async function withdraw(): Promise<void> {
+  const poolMintInfo = await tokenPool.getMintInfo();
+  const supply = poolMintInfo.supply.toNumber();
+  let swapTokenA = await mintA.getAccountInfo(tokenAccountA);
+  let swapTokenB = await mintB.getAccountInfo(tokenAccountB);
+  const tokenA = (swapTokenA.amount.toNumber() * POOL_TOKEN_AMOUNT) / supply;
+  const tokenB = (swapTokenB.amount.toNumber() * POOL_TOKEN_AMOUNT) / supply;
+
   console.log('Creating withdraw token A account');
   let userAccountA = await mintA.createAccount(owner.publicKey);
   console.log('Creating withdraw token B account');
   let userAccountB = await mintB.createAccount(owner.publicKey);
 
   console.log('Approving withdrawal from pool account');
-  await tokenPool.approve(tokenAccountPool, authority, owner, [], USER_AMOUNT);
+  await tokenPool.approve(
+    tokenAccountPool,
+    authority,
+    owner,
+    [],
+    POOL_TOKEN_AMOUNT,
+  );
   const [tokenProgramId] = await GetPrograms(connection);
 
   console.log('Withdrawing pool tokens for A and B tokens');
@@ -277,19 +301,23 @@ export async function withdraw(): Promise<void> {
     userAccountA,
     userAccountB,
     tokenProgramId,
-    USER_AMOUNT,
+    POOL_TOKEN_AMOUNT,
   );
 
+  //const poolMintInfo = await tokenPool.getMintInfo();
+  swapTokenA = await mintA.getAccountInfo(tokenAccountA);
+  swapTokenB = await mintB.getAccountInfo(tokenAccountB);
+
   let info = await tokenPool.getAccountInfo(tokenAccountPool);
-  assert(info.amount.toNumber() == BASE_AMOUNT - USER_AMOUNT);
-  info = await mintA.getAccountInfo(tokenAccountA);
-  assert(info.amount.toNumber() == BASE_AMOUNT);
-  info = await mintB.getAccountInfo(tokenAccountB);
-  assert(info.amount.toNumber() == BASE_AMOUNT);
+  assert(
+    info.amount.toNumber() == DEFAULT_POOL_TOKEN_AMOUNT - POOL_TOKEN_AMOUNT,
+  );
+  assert(swapTokenA.amount.toNumber() == BASE_AMOUNT);
+  assert(swapTokenB.amount.toNumber() == BASE_AMOUNT);
   info = await mintA.getAccountInfo(userAccountA);
-  assert(info.amount.toNumber() == USER_AMOUNT);
+  assert(info.amount.toNumber() == tokenA);
   info = await mintB.getAccountInfo(userAccountB);
-  assert(info.amount.toNumber() == USER_AMOUNT);
+  assert(info.amount.toNumber() == tokenB);
 }
 
 export async function swap(): Promise<void> {
@@ -322,5 +350,7 @@ export async function swap(): Promise<void> {
   info = await mintB.getAccountInfo(userAccountB);
   assert(info.amount.toNumber() == 69);
   info = await tokenPool.getAccountInfo(tokenAccountPool);
-  assert(info.amount.toNumber() == BASE_AMOUNT - USER_AMOUNT);
+  assert(
+    info.amount.toNumber() == DEFAULT_POOL_TOKEN_AMOUNT - POOL_TOKEN_AMOUNT,
+  );
 }

token-swap/js/client/token-swap.js

@@ -419,7 +419,7 @@ export class TokenSwap {
    * @param poolToken Pool token
    * @param poolAccount Pool account to deposit the generated tokens
    * @param tokenProgramId Token program id
-   * @param amount Amount of token A to transfer, token B amount is set by the exchange rate
+   * @param amount Amount of pool token to deposit, token A and B amount are set by the exchange rate relative to the total pool token supply
    */
   async deposit(
     authority: PublicKey,
@@ -510,7 +510,7 @@ export class TokenSwap {
    * @param userAccountA Token A user account
    * @param userAccountB token B user account
    * @param tokenProgramId Token program id
-   * @param amount Amount of token A to transfer, token B amount is set by the exchange rate
+   * @param amount Amount of pool token to withdraw, token A and B amount are set by the exchange rate relative to the total pool token supply
    */
   async withdraw(
     authority: PublicKey,

token-swap/program/src/curve.rs

@@ -0,0 +1,165 @@
+//! Swap calculations and curve implementations
+
+/// Initial amount of pool tokens for swap contract, hard-coded to something
+/// "sensible" given a maximum of u64.
+/// Note that on Ethereum, Uniswap uses the geometric mean of all provided
+/// input amounts, and Balancer uses 100 * 10 ^ 18.
+pub const INITIAL_SWAP_POOL_AMOUNT: u64 = 1_000_000_000;
+
+/// Encodes all results of swapping from a source token to a destination token
+pub struct SwapResult {
+    /// New amount of source token
+    pub new_source_amount: u64,
+    /// New amount of destination token
+    pub new_destination_amount: u64,
+    /// Amount of destination token swapped
+    pub amount_swapped: u64,
+}
+
+impl SwapResult {
+    /// SwapResult for swap from one currency into another, given pool information
+    /// and fee
+    pub fn swap_to(
+        source_amount: u64,
+        swap_source_amount: u64,
+        swap_destination_amount: u64,
+        fee_numerator: u64,
+        fee_denominator: u64,
+    ) -> Option<SwapResult> {
+        let invariant = swap_source_amount.checked_mul(swap_destination_amount)?;
+        let new_source_amount = swap_source_amount.checked_add(source_amount)?;
+        let new_destination_amount = invariant.checked_div(new_source_amount)?;
+        let remove = swap_destination_amount.checked_sub(new_destination_amount)?;
+        let fee = remove
+            .checked_mul(fee_numerator)?
+            .checked_div(fee_denominator)?;
+        let new_destination_amount = new_destination_amount.checked_add(fee)?;
+        let amount_swapped = remove.checked_sub(fee)?;
+        Some(SwapResult {
+            new_source_amount,
+            new_destination_amount,
+            amount_swapped,
+        })
+    }
+}
+
+/// The Uniswap invariant calculator.
+pub struct ConstantProduct {
+    /// Token A
+    pub token_a: u64,
+    /// Token B
+    pub token_b: u64,
+    /// Fee numerator
+    pub fee_numerator: u64,
+    /// Fee denominator
+    pub fee_denominator: u64,
+}
+
+impl ConstantProduct {
+    /// Swap token a to b
+    pub fn swap_a_to_b(&mut self, token_a: u64) -> Option<u64> {
+        let result = SwapResult::swap_to(
+            token_a,
+            self.token_a,
+            self.token_b,
+            self.fee_numerator,
+            self.fee_denominator,
+        )?;
+        self.token_a = result.new_source_amount;
+        self.token_b = result.new_destination_amount;
+        Some(result.amount_swapped)
+    }
+
+    /// Swap token b to a
+    pub fn swap_b_to_a(&mut self, token_b: u64) -> Option<u64> {
+        let result = SwapResult::swap_to(
+            token_b,
+            self.token_b,
+            self.token_a,
+            self.fee_numerator,
+            self.fee_denominator,
+        )?;
+        self.token_b = result.new_source_amount;
+        self.token_a = result.new_destination_amount;
+        Some(result.amount_swapped)
+    }
+}
+
+/// Conversions for pool tokens, how much to deposit / withdraw, along with
+/// proper initialization
+pub struct PoolTokenConverter {
+    /// Total supply
+    pub supply: u64,
+    /// Token A amount
+    pub token_a: u64,
+    /// Token B amount
+    pub token_b: u64,
+}
+
+impl PoolTokenConverter {
+    /// Create a converter based on existing market information
+    pub fn new_existing(supply: u64, token_a: u64, token_b: u64) -> Self {
+        Self {
+            supply,
+            token_a,
+            token_b,
+        }
+    }
+
+    /// Create a converter for a new pool token, no supply present yet.
+    /// According to Uniswap, the geometric mean protects the pool creator
+    /// in case the initial ratio is off the market.
+    pub fn new_pool(token_a: u64, token_b: u64) -> Self {
+        let supply = INITIAL_SWAP_POOL_AMOUNT;
+        Self {
+            supply,
+            token_a,
+            token_b,
+        }
+    }
+
+    /// A tokens for pool tokens
+    pub fn token_a_rate(&self, pool_tokens: u64) -> Option<u64> {
+        pool_tokens
+            .checked_mul(self.token_a)?
+            .checked_div(self.supply)
+    }
+
+    /// B tokens for pool tokens
+    pub fn token_b_rate(&self, pool_tokens: u64) -> Option<u64> {
+        pool_tokens
+            .checked_mul(self.token_b)?
+            .checked_div(self.supply)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn initial_pool_amount() {
+        let token_converter = PoolTokenConverter::new_pool(1, 5);
+        assert_eq!(token_converter.supply, INITIAL_SWAP_POOL_AMOUNT);
+    }
+
+    fn check_pool_token_a_rate(
+        token_a: u64,
+        token_b: u64,
+        deposit: u64,
+        supply: u64,
+        expected: Option<u64>,
+    ) {
+        let calculator = PoolTokenConverter::new_existing(supply, token_a, token_b);
+        assert_eq!(calculator.token_a_rate(deposit), expected);
+    }
+
+    #[test]
+    fn issued_tokens() {
+        check_pool_token_a_rate(2, 50, 5, 10, Some(1));
+        check_pool_token_a_rate(10, 10, 5, 10, Some(5));
+        check_pool_token_a_rate(5, 100, 5, 10, Some(2));
+        check_pool_token_a_rate(5, u64::MAX, 5, 10, Some(2));
+        check_pool_token_a_rate(u64::MAX, u64::MAX, 5, 10, None);
+    }
+}

token-swap/program/src/lib.rs

@@ -2,6 +2,7 @@
 
 //! An Uniswap-like program for the Solana blockchain.
 
+pub mod curve;
 pub mod entrypoint;
 pub mod error;
 pub mod instruction;