From 29a9fd9bca0c476f6a30e29e779ae941d3154c4d Mon Sep 17 00:00:00 2001
From: referencedev <goi65io3903d@protonmail.com>
Date: Fri, 10 Sep 2021 07:51:50 -0700
Subject: [PATCH 1/6] Initial setup for stable swap

---
 ref-exchange/src/lib.rs              |   1 +
 ref-exchange/src/pool.rs             |  21 ++
 ref-exchange/src/simple_pool.rs      |   5 +
 ref-exchange/src/stable_swap/math.rs | 348 +++++++++++++++++++++++++++
 ref-exchange/src/stable_swap/mod.rs  | 222 +++++++++++++++++
 ref-exchange/src/views.rs            |   7 +
 6 files changed, 604 insertions(+)
 create mode 100644 ref-exchange/src/stable_swap/math.rs
 create mode 100644 ref-exchange/src/stable_swap/mod.rs

diff --git a/ref-exchange/src/lib.rs b/ref-exchange/src/lib.rs
index 367a9db..414b4fd 100644
--- a/ref-exchange/src/lib.rs
+++ b/ref-exchange/src/lib.rs
@@ -28,6 +28,7 @@ mod multi_fungible_token;
 mod owner;
 mod pool;
 mod simple_pool;
+mod stable_swap;
 mod storage_impl;
 mod token_receiver;
 mod utils;
diff --git a/ref-exchange/src/pool.rs b/ref-exchange/src/pool.rs
index 62f650a..8d9d0b2 100644
--- a/ref-exchange/src/pool.rs
+++ b/ref-exchange/src/pool.rs
@@ -2,6 +2,7 @@ use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::{AccountId, Balance};
 
 use crate::simple_pool::SimplePool;
+use crate::stable_swap::StableSwapPool;
 use crate::utils::SwapVolume;
 
 /// Generic Pool, providing wrapper around different implementations of swap pools.
@@ -9,6 +10,7 @@ use crate::utils::SwapVolume;
 #[derive(BorshSerialize, BorshDeserialize)]
 pub enum Pool {
     SimplePool(SimplePool),
+    StableSwapPool(StableSwapPool),
 }
 
 impl Pool {
@@ -16,6 +18,7 @@ impl Pool {
     pub fn kind(&self) -> String {
         match self {
             Pool::SimplePool(_) => "SIMPLE_POOL".to_string(),
+            Pool::StableSwapPool(_) => "STABLE_SWAP".to_string(),
         }
     }
 
@@ -23,6 +26,7 @@ impl Pool {
     pub fn tokens(&self) -> &[AccountId] {
         match self {
             Pool::SimplePool(pool) => pool.tokens(),
+            Pool::StableSwapPool(pool) => pool.tokens(),
         }
     }
 
@@ -31,6 +35,7 @@ impl Pool {
     pub fn add_liquidity(&mut self, sender_id: &AccountId, amounts: &mut Vec<Balance>) -> Balance {
         match self {
             Pool::SimplePool(pool) => pool.add_liquidity(sender_id, amounts),
+            Pool::StableSwapPool(pool) => pool.add_liquidity(sender_id, amounts),
         }
     }
 
@@ -43,6 +48,7 @@ impl Pool {
     ) -> Vec<Balance> {
         match self {
             Pool::SimplePool(pool) => pool.remove_liquidity(sender_id, shares, min_amounts),
+            Pool::StableSwapPool(pool) => pool.remove_liquidity(sender_id, shares, min_amounts),
         }
     }
 
@@ -55,6 +61,7 @@ impl Pool {
     ) -> Balance {
         match self {
             Pool::SimplePool(pool) => pool.get_return(token_in, amount_in, token_out),
+            Pool::StableSwapPool(pool) => pool.get_return(token_in, amount_in, token_out),
         }
     }
 
@@ -62,6 +69,7 @@ impl Pool {
     pub fn get_fee(&self) -> u32 {
         match self {
             Pool::SimplePool(pool) => pool.get_fee(),
+            Pool::StableSwapPool(pool) => pool.get_fee(),
         }
     }
 
@@ -69,6 +77,7 @@ impl Pool {
     pub fn get_volumes(&self) -> Vec<SwapVolume> {
         match self {
             Pool::SimplePool(pool) => pool.get_volumes(),
+            Pool::StableSwapPool(pool) => pool.get_volumes(),
         }
     }
 
@@ -91,30 +100,42 @@ impl Pool {
                 exchange_id,
                 referral_id,
             ),
+            Pool::StableSwapPool(pool) => pool.swap(
+                token_in,
+                amount_in,
+                token_out,
+                min_amount_out,
+                exchange_id,
+                referral_id,
+            ),
         }
     }
 
     pub fn share_total_balance(&self) -> Balance {
         match self {
             Pool::SimplePool(pool) => pool.share_total_balance(),
+            Pool::StableSwapPool(pool) => pool.share_total_balance(),
         }
     }
 
     pub fn share_balances(&self, account_id: &AccountId) -> Balance {
         match self {
             Pool::SimplePool(pool) => pool.share_balance_of(account_id),
+            Pool::StableSwapPool(pool) => pool.share_balance_of(account_id),
         }
     }
 
     pub fn share_transfer(&mut self, sender_id: &AccountId, receiver_id: &AccountId, amount: u128) {
         match self {
             Pool::SimplePool(pool) => pool.share_transfer(sender_id, receiver_id, amount),
+            Pool::StableSwapPool(pool) => pool.share_transfer(sender_id, receiver_id, amount),
         }
     }
 
     pub fn share_register(&mut self, account_id: &AccountId) {
         match self {
             Pool::SimplePool(pool) => pool.share_register(account_id),
+            Pool::StableSwapPool(pool) => pool.share_register(account_id),
         }
     }
 }
diff --git a/ref-exchange/src/simple_pool.rs b/ref-exchange/src/simple_pool.rs
index fc1d874..d75e0ca 100644
--- a/ref-exchange/src/simple_pool.rs
+++ b/ref-exchange/src/simple_pool.rs
@@ -170,6 +170,11 @@ impl SimplePool {
         shares: Balance,
         min_amounts: Vec<Balance>,
     ) -> Vec<Balance> {
+        assert_eq!(
+            min_amounts.len(),
+            self.token_account_ids.len(),
+            "ERR_WRONG_TOKEN_COUNT"
+        );
         let prev_shares_amount = self.shares.get(&sender_id).expect("ERR_NO_SHARES");
         assert!(prev_shares_amount >= shares, "ERR_NOT_ENOUGH_SHARES");
         let mut result = vec![];
diff --git a/ref-exchange/src/stable_swap/math.rs b/ref-exchange/src/stable_swap/math.rs
new file mode 100644
index 0000000..a6c25c6
--- /dev/null
+++ b/ref-exchange/src/stable_swap/math.rs
@@ -0,0 +1,348 @@
+///! Calculator to maintain the invariant on adding/removing liquidity and on swapping.
+///! Large part of the code was taken from https://github.com/saber-hq/stable-swap/blob/master/stable-swap-math/src/curve.rs
+use near_sdk::{Balance, Timestamp};
+
+use crate::utils::U256;
+
+/// Number of coins in the pool.
+pub const N_COINS: u64 = 2;
+/// Timestamp at 0
+pub const ZERO_TS: i64 = 0;
+/// Minimum ramp duration
+pub const MIN_RAMP_DURATION: i64 = 86400;
+/// Min amplification coefficient
+pub const MIN_AMP: u64 = 1;
+/// Max amplification coefficient
+pub const MAX_AMP: u64 = 1_000_000;
+/// Max number of tokens to swap at once.
+pub const MAX_TOKENS_IN: u64 = u64::MAX >> 4;
+
+pub struct Fees {
+    pub trade_fee: u64,
+    pub admin_fee: u64,
+}
+
+impl Fees {
+    pub fn trade_fee(&self, amount: Balance) -> Option<Balance> {
+        Some(0)
+    }
+
+    pub fn admin_trade_fee(&self, amount: Balance) -> Option<Balance> {
+        Some(0)
+    }
+
+    pub fn normalized_trade_fee(&self, num_coins: u64, amount: Balance) -> Option<Balance> {
+        Some(0)
+    }
+}
+
+/// 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: Balance,
+    /// New amount of destination token.
+    pub new_destination_amount: Balance,
+    /// Amount of destination token swapped.
+    pub amount_swapped: Balance,
+    /// Admin fee for the swap.
+    pub admin_fee: Balance,
+    /// Fee for the swap.
+    pub fee: Balance,
+}
+
+/// The StableSwap invariant calculator.
+pub struct StableSwap {
+    /// Initial amplification coefficient (A)
+    initial_amp_factor: u64,
+    /// Target amplification coefficient (A)
+    target_amp_factor: u64,
+    /// Current unix timestamp
+    current_ts: Timestamp,
+    /// Ramp A start timestamp
+    start_ramp_ts: Timestamp,
+    /// Ramp A stop timestamp
+    stop_ramp_ts: Timestamp,
+}
+
+impl StableSwap {
+    pub fn new(
+        initial_amp_factor: u64,
+        target_amp_factor: u64,
+        current_ts: Timestamp,
+        start_ramp_ts: Timestamp,
+        stop_ramp_ts: Timestamp,
+    ) -> Self {
+        Self {
+            initial_amp_factor,
+            target_amp_factor,
+            current_ts,
+            start_ramp_ts,
+            stop_ramp_ts,
+        }
+    }
+
+    fn compute_next_d(
+        &self,
+        amp_factor: u64,
+        d_init: U256,
+        d_prod: U256,
+        sum_x: Balance,
+    ) -> Option<U256> {
+        let ann = amp_factor.checked_mul(N_COINS.into())?;
+        let leverage = (sum_x as u128).checked_mul(ann.into())?;
+        // d = (ann * sum_x + d_prod * n_coins) * d / ((ann - 1) * d + (n_coins + 1) * d_prod)
+        let numerator = d_init.checked_mul(
+            d_prod
+                .checked_mul(N_COINS.into())?
+                .checked_add(leverage.into())?,
+        )?;
+        let denominator = d_init
+            .checked_mul(ann.checked_sub(1)?.into())?
+            .checked_add(d_prod.checked_mul((N_COINS + 1).into())?)?;
+        numerator.checked_div(denominator)
+    }
+
+    /// Compute the amplification coefficient (A)
+    pub fn compute_amp_factor(&self) -> Option<Balance> {
+        if self.current_ts < self.stop_ramp_ts {
+            let time_range = self.stop_ramp_ts.checked_sub(self.start_ramp_ts)?;
+            let time_delta = self.current_ts.checked_sub(self.start_ramp_ts)?;
+
+            // Compute amp factor based on ramp time
+            if self.target_amp_factor >= self.initial_amp_factor {
+                // Ramp up
+                let amp_range = self
+                    .target_amp_factor
+                    .checked_sub(self.initial_amp_factor)?;
+                let amp_delta = (amp_range as u128)
+                    .checked_mul(time_delta as u128)?
+                    .checked_div(time_range as u128)? as u64;
+                self.initial_amp_factor
+                    .checked_add(amp_delta)
+                    .map(|x| x as u128)
+            } else {
+                // Ramp down
+                let amp_range = self
+                    .initial_amp_factor
+                    .checked_sub(self.target_amp_factor)?;
+                let amp_delta = (amp_range as u128)
+                    .checked_mul(time_delta as u128)?
+                    .checked_div(time_range as u128)? as u64;
+                self.initial_amp_factor
+                    .checked_sub(amp_delta)
+                    .map(|x| x as u128)
+            }
+        } else {
+            // when stop_ramp_ts == 0 or current_ts >= stop_ramp_ts
+            Some(self.target_amp_factor as u128)
+        }
+    }
+
+    /// Compute stable swap invariant (D)
+    /// Equation:
+    /// A * sum(x_i) * n**n + D = A * D * n**n + D**(n+1) / (n**n * prod(x_i))
+    pub fn compute_d(&self, amount_a: Balance, amount_b: Balance) -> Option<U256> {
+        let sum_x = amount_a.checked_add(amount_b)?; // sum(x_i), a.k.a S
+        if sum_x == 0 {
+            Some(0.into())
+        } else {
+            let amp_factor = self.compute_amp_factor()?;
+            let amount_a_times_coins = amount_a.checked_mul(N_COINS.into())?;
+            let amount_b_times_coins = amount_b.checked_mul(N_COINS.into())?;
+
+            // Newton's method to approximate D
+            let mut d_prev: U256;
+            let mut d: U256 = sum_x.into();
+            for _ in 0..256 {
+                let mut d_prod = d;
+                d_prod = d_prod
+                    .checked_mul(d)?
+                    .checked_div(amount_a_times_coins.into())?;
+                d_prod = d_prod
+                    .checked_mul(d)?
+                    .checked_div(amount_b_times_coins.into())?;
+                d_prev = d;
+                d = self.compute_next_d(amp_factor as u64, d, d_prod, sum_x)?;
+                // Equality with the precision of 1
+                if d > d_prev {
+                    if d.checked_sub(d_prev)? <= 1.into() {
+                        break;
+                    }
+                } else if d_prev.checked_sub(d)? <= 1.into() {
+                    break;
+                }
+            }
+
+            Some(d)
+        }
+    }
+
+    /// Compute the amount of LP tokens to mint after a deposit
+    pub fn compute_lp_amount_for_deposit(
+        &self,
+        deposit_amount_a: Balance,
+        deposit_amount_b: Balance,
+        swap_amount_a: Balance,
+        swap_amount_b: Balance,
+        pool_token_supply: Balance,
+        fees: &Fees,
+    ) -> Option<Balance> {
+        // Initial invariant
+        let d_0 = self.compute_d(swap_amount_a, swap_amount_b)?;
+        let old_balances = [swap_amount_a, swap_amount_b];
+        let mut new_balances = [
+            swap_amount_a.checked_add(deposit_amount_a)?,
+            swap_amount_b.checked_add(deposit_amount_b)?,
+        ];
+        // Invariant after change
+        let d_1 = self.compute_d(new_balances[0], new_balances[1])?;
+        if d_1 <= d_0 {
+            None
+        } else {
+            // Recalculate the invariant accounting for fees
+            for i in 0..new_balances.len() {
+                let ideal_balance = d_1
+                    .checked_mul(old_balances[i].into())?
+                    .checked_div(d_0)?
+                    .as_u128();
+                let difference = if ideal_balance > new_balances[i] {
+                    ideal_balance.checked_sub(new_balances[i])?
+                } else {
+                    new_balances[i].checked_sub(ideal_balance)?
+                };
+                let fee = fees.normalized_trade_fee(N_COINS, difference)?;
+                new_balances[i] = new_balances[i].checked_sub(fee)?;
+            }
+
+            let d_2 = self.compute_d(new_balances[0], new_balances[1])?;
+            Some(
+                U256::from(pool_token_supply)
+                    .checked_mul(d_2.checked_sub(d_0)?)?
+                    .checked_div(d_0)?
+                    .as_u128(),
+            )
+        }
+    }
+
+    /// Compute swap amount `y` in proportion to `x`
+    /// Solve for y:
+    /// y**2 + y * (sum' - (A*n**n - 1) * D / (A * n**n)) = D ** (n + 1) / (n ** (2 * n) * prod' * A)
+    /// y**2 + b*y = c
+    pub fn compute_y_raw(&self, x: Balance, d: U256) -> Option<U256> {
+        let amp_factor = self.compute_amp_factor()?;
+        let ann = amp_factor.checked_mul(N_COINS.into())?; // A * n ** n
+
+        // sum' = prod' = x
+        // c =  D ** (n + 1) / (n ** (2 * n) * prod' * A)
+        let mut c = d
+            .checked_mul(d)?
+            .checked_div(x.checked_mul(N_COINS.into())?.into())?;
+        c = c
+            .checked_mul(d)?
+            .checked_div(ann.checked_mul(N_COINS.into())?.into())?;
+        // b = sum' - (A*n**n - 1) * D / (A * n**n)
+        let b = d.checked_div(ann.into())?.checked_add(x.into())?; // d is subtracted on line 147
+
+        // Solve for y by approximating: y**2 + b*y = c
+        let mut y_prev: U256;
+        let mut y = d;
+        for _ in 0..256 {
+            y_prev = y;
+            // y = (y * y + c) / (2 * y + b - d);
+            let y_numerator = y.checked_pow(2.into())?.checked_add(c)?;
+            let y_denominator = y.checked_mul(2.into())?.checked_add(b)?.checked_sub(d)?;
+            y = y_numerator.checked_div(y_denominator)?;
+            if y > y_prev {
+                if y.checked_sub(y_prev)? <= 1.into() {
+                    break;
+                }
+            } else if y_prev.checked_sub(y)? <= 1.into() {
+                break;
+            }
+        }
+        Some(y)
+    }
+
+    /// Compute swap amount `y` in proportion to `x`
+    pub fn compute_y(&self, x: Balance, d: U256) -> u128 {
+        self.compute_y_raw(x, d).unwrap().as_u128()
+    }
+
+    /// Calculate withdrawal amount when withdrawing only one type of token
+    /// Calculation:
+    /// 1. Get current D
+    /// 2. Solve Eqn against y_i for D - _token_amount
+    pub fn compute_withdraw_one(
+        &self,
+        pool_token_amount: Balance,
+        pool_token_supply: Balance,
+        swap_base_amount: Balance, // Same denomination of token to be withdrawn
+        swap_quote_amount: Balance, // Counter denomination of token to be withdrawn
+        fees: &Fees,
+    ) -> Option<(Balance, Balance)> {
+        let d_0 = self.compute_d(swap_base_amount, swap_quote_amount)?;
+        let d_1 = d_0.checked_sub(
+            U256::from(pool_token_amount)
+                .checked_mul(d_0)?
+                .checked_div(pool_token_supply.into())?,
+        )?;
+        let new_y = self.compute_y(swap_quote_amount, d_1);
+
+        // expected_base_amount = swap_base_amount * d_1 / d_0 - new_y;
+        let expected_base_amount = U256::from(swap_base_amount)
+            .checked_mul(d_1)?
+            .checked_div(d_0)?
+            .as_u128()
+            .checked_sub(new_y)?;
+        // expected_quote_amount = swap_quote_amount - swap_quote_amount * d_1 / d_0;
+        let expected_quote_amount = swap_quote_amount.checked_sub(
+            U256::from(swap_quote_amount)
+                .checked_mul(d_1)?
+                .checked_div(d_0)?
+                .as_u128(),
+        )?;
+        // new_base_amount = swap_base_amount - expected_base_amount * fee / fee_denominator;
+        let new_base_amount = swap_base_amount
+            .checked_sub(fees.normalized_trade_fee(N_COINS, expected_base_amount)?)?;
+        // new_quote_amount = swap_quote_amount - expected_quote_amount * fee / fee_denominator;
+        let new_quote_amount = swap_quote_amount
+            .checked_sub(fees.normalized_trade_fee(N_COINS, expected_quote_amount)?)?;
+        let dy = new_base_amount
+            .checked_sub(self.compute_y(new_quote_amount, d_1))?
+            .checked_sub(1)?; // Withdraw less to account for rounding errors
+        let dy_0 = swap_base_amount.checked_sub(new_y)?;
+
+        Some((dy, dy_0 - dy))
+    }
+
+    /// Compute SwapResult after an exchange
+    pub fn swap_to(
+        &self,
+        source_amount: Balance,
+        swap_source_amount: Balance,
+        swap_destination_amount: Balance,
+        fees: &Fees,
+    ) -> Option<SwapResult> {
+        let y = self.compute_y(
+            swap_source_amount.checked_add(source_amount)?,
+            self.compute_d(swap_source_amount, swap_destination_amount)?,
+        );
+        let dy = swap_destination_amount.checked_sub(y)?;
+        let dy_fee = fees.trade_fee(dy)?;
+        let admin_fee = fees.admin_trade_fee(dy_fee)?;
+
+        let amount_swapped = dy.checked_sub(dy_fee)?;
+        let new_destination_amount = swap_destination_amount
+            .checked_sub(amount_swapped)?
+            .checked_sub(admin_fee)?;
+        let new_source_amount = swap_source_amount.checked_add(source_amount)?;
+
+        Some(SwapResult {
+            new_source_amount,
+            new_destination_amount,
+            amount_swapped,
+            admin_fee,
+            fee: dy_fee,
+        })
+    }
+}
diff --git a/ref-exchange/src/stable_swap/mod.rs b/ref-exchange/src/stable_swap/mod.rs
new file mode 100644
index 0000000..87c6424
--- /dev/null
+++ b/ref-exchange/src/stable_swap/mod.rs
@@ -0,0 +1,222 @@
+use near_sdk::collections::LookupMap;
+use near_sdk::json_types::ValidAccountId;
+use near_sdk::{env, AccountId, Balance};
+
+use crate::errors::{ERR13_LP_NOT_REGISTERED, ERR14_LP_ALREADY_REGISTERED};
+use crate::stable_swap::math::{Fees, StableSwap};
+use crate::utils::{add_to_collection, SwapVolume};
+use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
+
+mod math;
+
+#[derive(BorshSerialize, BorshDeserialize)]
+pub struct StableSwapPool {
+    /// List of tokens in the pool.
+    pub token_account_ids: Vec<AccountId>,
+    /// How much NEAR this contract has.
+    pub amounts: Vec<Balance>,
+    /// Volumes accumulated by this pool.
+    pub volumes: Vec<SwapVolume>,
+    /// Fee charged for swap (gets divided by FEE_DIVISOR).
+    pub total_fee: u32,
+    /// Portion of the fee going to exchange.
+    pub exchange_fee: u32,
+    /// Portion of the fee going to referral.
+    pub referral_fee: u32,
+    /// Shares of the pool by liquidity providers.
+    pub shares: LookupMap<AccountId, Balance>,
+    /// Total number of shares.
+    pub shares_total_supply: Balance,
+}
+
+impl StableSwapPool {
+    pub fn new(
+        id: u32,
+        token_account_ids: Vec<ValidAccountId>,
+        total_fee: u32,
+        exchange_fee: u32,
+        referral_fee: u32,
+    ) -> Self {
+        assert_eq!(token_account_ids.len(), 2, "ERR_WRONG_TOKEN_COUNT");
+        Self {
+            token_account_ids: token_account_ids.iter().map(|a| a.clone().into()).collect(),
+            amounts: vec![0u128; token_account_ids.len()],
+            volumes: vec![SwapVolume::default(); token_account_ids.len()],
+            total_fee,
+            exchange_fee,
+            referral_fee,
+            shares: LookupMap::new(format!("ss{}", id).into_bytes()),
+            shares_total_supply: 0,
+        }
+    }
+
+    /// Returns token index for given pool.
+    fn token_index(&self, token_id: &AccountId) -> usize {
+        self.token_account_ids
+            .iter()
+            .position(|id| id == token_id)
+            .expect("ERR_MISSING_TOKEN")
+    }
+
+    /// Returns given pool's total fee.
+    pub fn get_fee(&self) -> u32 {
+        self.total_fee
+    }
+
+    /// Returns volumes of the given pool.
+    pub fn get_volumes(&self) -> Vec<SwapVolume> {
+        self.volumes.clone()
+    }
+
+    /// Add liquidity into the pool.
+    /// Allows to add liquidity of a subset of tokens.
+    pub fn add_liquidity(&mut self, sender_id: &AccountId, amounts: &mut Vec<Balance>) -> Balance {
+        assert_eq!(
+            amounts.len(),
+            self.token_account_ids.len(),
+            "ERR_WRONG_TOKEN_COUNT"
+        );
+        let calc = StableSwap::new(0, 0, 0, 0, 0);
+        let new_shares = calc
+            .compute_lp_amount_for_deposit(
+                self.amounts[0],
+                self.amounts[1],
+                0,
+                0,
+                self.shares_total_supply,
+                &Fees {
+                    trade_fee: 0,
+                    admin_fee: 0,
+                },
+            )
+            // TODO: proper error
+            .expect("ERR_FAILED");
+        self.mint_shares(sender_id, new_shares);
+        new_shares
+    }
+
+    /// Mint new shares for given user.
+    fn mint_shares(&mut self, account_id: &AccountId, shares: Balance) {
+        if shares == 0 {
+            return;
+        }
+        self.shares_total_supply += shares;
+        add_to_collection(&mut self.shares, &account_id, shares);
+    }
+
+    /// Remove liquidity from the pool.
+    /// Allows to remove liquidity of a subset of tokens, by providing 0 in `min_amount` for the tokens to not withdraw.
+    pub fn remove_liquidity(
+        &mut self,
+        sender_id: &AccountId,
+        shares: Balance,
+        min_amounts: Vec<Balance>,
+    ) -> Vec<Balance> {
+        assert_eq!(
+            min_amounts.len(),
+            self.token_account_ids.len(),
+            "ERR_WRONG_TOKEN_COUNT"
+        );
+        let prev_shares_amount = self.shares.get(&sender_id).expect("ERR_NO_SHARES");
+        assert!(prev_shares_amount >= shares, "ERR_NOT_ENOUGH_SHARES");
+        let mut result = vec![];
+        result
+    }
+    /// Returns number of tokens in outcome, given amount.
+    /// Tokens are provided as indexes into token list for given pool.
+    fn internal_get_return(
+        &self,
+        token_in: usize,
+        amount_in: Balance,
+        token_out: usize,
+    ) -> Balance {
+        0
+    }
+
+    /// Returns how much token you will receive if swap `token_amount_in` of `token_in` for `token_out`.
+    pub fn get_return(
+        &self,
+        token_in: &AccountId,
+        amount_in: Balance,
+        token_out: &AccountId,
+    ) -> Balance {
+        self.internal_get_return(
+            self.token_index(token_in),
+            amount_in,
+            self.token_index(token_out),
+        )
+    }
+
+    /// Swap `token_amount_in` of `token_in` token into `token_out` and return how much was received.
+    /// Assuming that `token_amount_in` was already received from `sender_id`.
+    pub fn swap(
+        &mut self,
+        token_in: &AccountId,
+        amount_in: Balance,
+        token_out: &AccountId,
+        min_amount_out: Balance,
+        exchange_id: &AccountId,
+        referral_id: &Option<AccountId>,
+    ) -> Balance {
+        0
+    }
+
+    /// Register given account with 0 balance in shares.
+    /// Storage payment should be checked by caller.
+    pub fn share_register(&mut self, account_id: &AccountId) {
+        if self.shares.contains_key(account_id) {
+            env::panic(ERR14_LP_ALREADY_REGISTERED.as_bytes());
+        }
+        self.shares.insert(account_id, &0);
+    }
+
+    /// Transfers shares from predecessor to receiver.
+    pub fn share_transfer(&mut self, sender_id: &AccountId, receiver_id: &AccountId, amount: u128) {
+        let balance = self.shares.get(&sender_id).expect("ERR_NO_SHARES");
+        if let Some(new_balance) = balance.checked_sub(amount) {
+            self.shares.insert(&sender_id, &new_balance);
+        } else {
+            env::panic(b"ERR_NOT_ENOUGH_SHARES");
+        }
+        let balance_out = self
+            .shares
+            .get(&receiver_id)
+            .expect(ERR13_LP_NOT_REGISTERED);
+        self.shares.insert(&receiver_id, &(balance_out + amount));
+    }
+
+    /// Returns balance of shares for given user.
+    pub fn share_balance_of(&self, account_id: &AccountId) -> Balance {
+        self.shares.get(account_id).unwrap_or_default()
+    }
+
+    /// Returns total number of shares in this pool.
+    pub fn share_total_balance(&self) -> Balance {
+        self.shares_total_supply
+    }
+
+    /// Returns list of tokens in this pool.
+    pub fn tokens(&self) -> &[AccountId] {
+        &self.token_account_ids
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use near_sdk::test_utils::{accounts, VMContextBuilder};
+    use near_sdk::{testing_env, MockedBlockchain};
+    use near_sdk_sim::to_yocto;
+
+    #[test]
+    fn test_basics() {
+        let mut context = VMContextBuilder::new();
+        context.predecessor_account_id(accounts(0));
+        testing_env!(context.build());
+        let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 0, 0, 0);
+        let mut amounts = vec![to_yocto("5"), to_yocto("10")];
+        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
+        assert_eq!(amounts, vec![to_yocto("5"), to_yocto("10")]);
+        assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 1);
+    }
+}
diff --git a/ref-exchange/src/views.rs b/ref-exchange/src/views.rs
index 3d62cc1..d7d901d 100644
--- a/ref-exchange/src/views.rs
+++ b/ref-exchange/src/views.rs
@@ -35,6 +35,13 @@ impl From<Pool> for PoolInfo {
                 total_fee: pool.total_fee,
                 shares_total_supply: U128(pool.shares_total_supply),
             },
+            Pool::StableSwapPool(pool) => Self {
+                pool_kind,
+                token_account_ids: pool.token_account_ids,
+                amounts: pool.amounts.into_iter().map(|a| U128(a)).collect(),
+                total_fee: pool.total_fee,
+                shares_total_supply: U128(pool.shares_total_supply),
+            },
         }
     }
 }

From 962c844990703e112ac58b27710d0e4d84c10004 Mon Sep 17 00:00:00 2001
From: referencedev <goi65io3903d@protonmail.com>
Date: Mon, 13 Sep 2021 08:37:53 -0700
Subject: [PATCH 2/6] Continue working on stable swap

---
 ref-exchange/src/fees.rs            |  11 ++
 ref-exchange/src/lib.rs             |  38 +++--
 ref-exchange/src/pool.rs            |  29 ++--
 ref-exchange/src/simple_pool.rs     |  42 +++---
 ref-exchange/src/stable_swap/mod.rs | 214 ++++++++++++++++++++++++----
 5 files changed, 250 insertions(+), 84 deletions(-)
 create mode 100644 ref-exchange/src/fees.rs

diff --git a/ref-exchange/src/fees.rs b/ref-exchange/src/fees.rs
new file mode 100644
index 0000000..c6c0ae8
--- /dev/null
+++ b/ref-exchange/src/fees.rs
@@ -0,0 +1,11 @@
+use near_sdk::AccountId;
+
+/// Maintain information about fees.
+pub struct SwapFees {
+    /// Basis points of the fee for exchange.
+    pub exchange_fee: u32,
+    /// Basis points of the fee for referrer.
+    pub referral_fee: u32,
+    pub exchange_id: AccountId,
+    pub referral_id: Option<AccountId>,
+}
diff --git a/ref-exchange/src/lib.rs b/ref-exchange/src/lib.rs
index a8bb51b..1ae08ca 100644
--- a/ref-exchange/src/lib.rs
+++ b/ref-exchange/src/lib.rs
@@ -7,14 +7,15 @@ use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::{LookupMap, UnorderedSet, Vector};
 use near_sdk::json_types::{ValidAccountId, U128};
 use near_sdk::{
-    assert_one_yocto, env, log, near_bindgen, AccountId, Balance, PanicOnDefault, Promise,
-    PromiseResult, StorageUsage, BorshStorageKey
+    assert_one_yocto, env, log, near_bindgen, AccountId, Balance, BorshStorageKey, PanicOnDefault,
+    Promise, PromiseResult, StorageUsage,
 };
 
-use crate::account_deposit::{VAccount, Account};
+use crate::account_deposit::{Account, VAccount};
 pub use crate::action::SwapAction;
 use crate::action::{Action, ActionResult};
 use crate::errors::*;
+use crate::fees::SwapFees;
 use crate::pool::Pool;
 use crate::simple_pool::SimplePool;
 use crate::utils::check_token_duplicates;
@@ -23,6 +24,7 @@ pub use crate::views::PoolInfo;
 mod account_deposit;
 mod action;
 mod errors;
+mod fees;
 mod legacy;
 mod multi_fungible_token;
 mod owner;
@@ -89,7 +91,7 @@ impl Contract {
         )))
     }
 
-    /// [AUDIT_03_reject(NOPE action is allowed by design)] 
+    /// [AUDIT_03_reject(NOPE action is allowed by design)]
     /// [AUDIT_04]
     /// Executes generic set of actions.
     /// If referrer provided, pays referral_fee to it.
@@ -208,7 +210,6 @@ impl Contract {
 
 /// Internal methods implementation.
 impl Contract {
-
     /// Check how much storage taken costs and refund the left over back.
     fn internal_check_storage(&self, prev_storage: StorageUsage) {
         let storage_cost = env::storage_usage()
@@ -298,8 +299,12 @@ impl Contract {
             amount_in,
             token_out,
             min_amount_out,
-            &self.owner_id,
-            referral_id,
+            SwapFees {
+                exchange_fee: self.exchange_fee,
+                exchange_id: self.owner_id.clone(),
+                referral_fee: self.referral_fee,
+                referral_id: referral_id.clone(),
+            },
         );
         self.pools.replace(pool_id, &pool);
         amount_out
@@ -746,7 +751,7 @@ mod tests {
         // account(0) -- swap contract
         // account(1) -- token0 contract
         // account(2) -- token1 contract
-        // account(3) -- user account 
+        // account(3) -- user account
         // account(4) -- another user account
         let (mut context, mut contract) = setup_contract();
         deposit_tokens(
@@ -769,20 +774,14 @@ mod tests {
             contract.mft_balance_of(":0".to_string(), accounts(3)).0,
             to_yocto("1")
         );
-        assert_eq!(
-            contract.mft_total_supply(":0".to_string()).0,
-            to_yocto("1")
-        );
+        assert_eq!(contract.mft_total_supply(":0".to_string()).0, to_yocto("1"));
         testing_env!(context.attached_deposit(1).build());
         contract.add_liquidity(id, vec![U128(to_yocto("50")), U128(to_yocto("50"))], None);
         assert_eq!(
             contract.mft_balance_of(":0".to_string(), accounts(3)).0,
             to_yocto("2")
         );
-        assert_eq!(
-            contract.mft_total_supply(":0".to_string()).0,
-            to_yocto("2")
-        );
+        assert_eq!(contract.mft_total_supply(":0".to_string()).0, to_yocto("2"));
 
         // register another user
         testing_env!(context
@@ -804,10 +803,7 @@ mod tests {
             contract.mft_balance_of(":0".to_string(), accounts(4)).0,
             to_yocto("1")
         );
-        assert_eq!(
-            contract.mft_total_supply(":0".to_string()).0,
-            to_yocto("2")
-        );
+        assert_eq!(contract.mft_total_supply(":0".to_string()).0, to_yocto("2"));
         // remove lpt for account 3
         testing_env!(context
             .predecessor_account_id(accounts(3))
@@ -860,7 +856,7 @@ mod tests {
         // account(0) -- swap contract
         // account(1) -- token0 contract
         // account(2) -- token1 contract
-        // account(3) -- user account 
+        // account(3) -- user account
         let (mut context, mut contract) = setup_contract();
         deposit_tokens(
             &mut context,
diff --git a/ref-exchange/src/pool.rs b/ref-exchange/src/pool.rs
index 8d9d0b2..9a78984 100644
--- a/ref-exchange/src/pool.rs
+++ b/ref-exchange/src/pool.rs
@@ -1,6 +1,7 @@
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::{AccountId, Balance};
 
+use crate::fees::SwapFees;
 use crate::simple_pool::SimplePool;
 use crate::stable_swap::StableSwapPool;
 use crate::utils::SwapVolume;
@@ -61,7 +62,8 @@ impl Pool {
     ) -> Balance {
         match self {
             Pool::SimplePool(pool) => pool.get_return(token_in, amount_in, token_out),
-            Pool::StableSwapPool(pool) => pool.get_return(token_in, amount_in, token_out),
+            _ => 0
+            // Pool::StableSwapPool(pool) => pool.get_return(token_in, amount_in, token_out),
         }
     }
 
@@ -88,26 +90,15 @@ impl Pool {
         amount_in: Balance,
         token_out: &AccountId,
         min_amount_out: Balance,
-        exchange_id: &AccountId,
-        referral_id: &Option<AccountId>,
+        fees: SwapFees,
     ) -> Balance {
         match self {
-            Pool::SimplePool(pool) => pool.swap(
-                token_in,
-                amount_in,
-                token_out,
-                min_amount_out,
-                exchange_id,
-                referral_id,
-            ),
-            Pool::StableSwapPool(pool) => pool.swap(
-                token_in,
-                amount_in,
-                token_out,
-                min_amount_out,
-                exchange_id,
-                referral_id,
-            ),
+            Pool::SimplePool(pool) => {
+                pool.swap(token_in, amount_in, token_out, min_amount_out, fees)
+            }
+            Pool::StableSwapPool(pool) => {
+                pool.swap(token_in, amount_in, token_out, min_amount_out, fees)
+            }
         }
     }
 
diff --git a/ref-exchange/src/simple_pool.rs b/ref-exchange/src/simple_pool.rs
index 18cd7e2..af08715 100644
--- a/ref-exchange/src/simple_pool.rs
+++ b/ref-exchange/src/simple_pool.rs
@@ -1,14 +1,15 @@
 use std::cmp::min;
 
+use crate::StorageKey;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::LookupMap;
 use near_sdk::json_types::ValidAccountId;
 use near_sdk::{env, AccountId, Balance};
-use crate::StorageKey;
 
 use crate::errors::{
     ERR13_LP_NOT_REGISTERED, ERR14_LP_ALREADY_REGISTERED, ERR31_ZERO_AMOUNT, ERR32_ZERO_SHARES,
 };
+use crate::fees::SwapFees;
 use crate::utils::{
     add_to_collection, integer_sqrt, SwapVolume, FEE_DIVISOR, INIT_SHARES_SUPPLY, U256,
 };
@@ -51,7 +52,11 @@ impl SimplePool {
             "ERR_FEE_TOO_LARGE"
         );
         // [AUDIT_10]
-        assert_eq!(token_account_ids.len(), NUM_TOKENS, "ERR_SHOULD_HAVE_2_TOKENS");
+        assert_eq!(
+            token_account_ids.len(),
+            NUM_TOKENS,
+            "ERR_SHOULD_HAVE_2_TOKENS"
+        );
         Self {
             token_account_ids: token_account_ids.iter().map(|a| a.clone().into()).collect(),
             amounts: vec![0u128; token_account_ids.len()],
@@ -60,9 +65,7 @@ impl SimplePool {
             exchange_fee,
             referral_fee,
             // [AUDIT_11]
-            shares: LookupMap::new(StorageKey::Shares {
-                pool_id: id,
-            }),
+            shares: LookupMap::new(StorageKey::Shares { pool_id: id }),
             shares_total_supply: 0,
         }
     }
@@ -272,8 +275,7 @@ impl SimplePool {
         amount_in: Balance,
         token_out: &AccountId,
         min_amount_out: Balance,
-        exchange_id: &AccountId,
-        referral_id: &Option<AccountId>,
+        fees: SwapFees,
     ) -> Balance {
         let in_idx = self.token_index(token_in);
         let out_idx = self.token_index(token_out);
@@ -304,14 +306,14 @@ impl SimplePool {
         // Allocate exchange fee as fraction of total fee by issuing LP shares proportionally.
         if self.exchange_fee > 0 && numerator > U256::zero() {
             let denominator = new_invariant * self.total_fee / self.exchange_fee;
-            self.mint_shares(&exchange_id, (numerator / denominator).as_u128());
+            self.mint_shares(&fees.exchange_id, (numerator / denominator).as_u128());
         }
 
         // If there is referral provided and the account already registered LP, allocate it % of LP rewards.
-        if let Some(referral_id) = referral_id {
+        if let Some(referral_id) = fees.referral_id {
             if self.referral_fee > 0
                 && numerator > U256::zero()
-                && self.shares.contains_key(referral_id)
+                && self.shares.contains_key(&referral_id)
             {
                 let denominator = new_invariant * self.total_fee / self.referral_fee;
                 self.mint_shares(&referral_id, (numerator / denominator).as_u128());
@@ -354,8 +356,12 @@ mod tests {
             one_near,
             accounts(2).as_ref(),
             1,
-            accounts(3).as_ref(),
-            &None,
+            SwapFees {
+                exchange_fee: 0,
+                exchange_id: accounts(3).as_ref().clone(),
+                referral_fee: 0,
+                referral_id: None,
+            },
         );
         assert_eq!(
             pool.share_balance_of(accounts(0).as_ref()),
@@ -400,8 +406,12 @@ mod tests {
             one_near,
             accounts(2).as_ref(),
             1,
-            accounts(3).as_ref(),
-            &None,
+            SwapFees {
+                exchange_fee: 100,
+                exchange_id: accounts(3).as_ref().clone(),
+                referral_fee: 0,
+                referral_id: None,
+            },
         );
         assert_eq!(
             pool.share_balance_of(accounts(0).as_ref()),
@@ -426,9 +436,7 @@ mod tests {
             exchange_fee: 5,
             referral_fee: 1,
             shares_total_supply: 35967818779820559673547466,
-            shares: LookupMap::new(StorageKey::Shares {
-                pool_id: 0,
-            }),
+            shares: LookupMap::new(StorageKey::Shares { pool_id: 0 }),
         };
         let mut amounts = vec![145782, 1];
         let _ = pool.add_liquidity(&accounts(2).to_string(), &mut amounts);
diff --git a/ref-exchange/src/stable_swap/mod.rs b/ref-exchange/src/stable_swap/mod.rs
index 87c6424..92f8744 100644
--- a/ref-exchange/src/stable_swap/mod.rs
+++ b/ref-exchange/src/stable_swap/mod.rs
@@ -1,10 +1,12 @@
 use near_sdk::collections::LookupMap;
 use near_sdk::json_types::ValidAccountId;
-use near_sdk::{env, AccountId, Balance};
+use near_sdk::{env, AccountId, Balance, Timestamp};
 
 use crate::errors::{ERR13_LP_NOT_REGISTERED, ERR14_LP_ALREADY_REGISTERED};
-use crate::stable_swap::math::{Fees, StableSwap};
-use crate::utils::{add_to_collection, SwapVolume};
+use crate::fees::SwapFees;
+use crate::stable_swap::math::{Fees, StableSwap, SwapResult, N_COINS};
+use crate::utils::{add_to_collection, SwapVolume, FEE_DIVISOR};
+use crate::StorageKey;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 
 mod math;
@@ -19,34 +21,44 @@ pub struct StableSwapPool {
     pub volumes: Vec<SwapVolume>,
     /// Fee charged for swap (gets divided by FEE_DIVISOR).
     pub total_fee: u32,
-    /// Portion of the fee going to exchange.
-    pub exchange_fee: u32,
-    /// Portion of the fee going to referral.
-    pub referral_fee: u32,
     /// Shares of the pool by liquidity providers.
     pub shares: LookupMap<AccountId, Balance>,
     /// Total number of shares.
     pub shares_total_supply: Balance,
+    /// Initial amplification coefficient.
+    pub init_amp_factor: u64,
+    /// Future amplification coefficient.
+    pub future_amp_factor: u64,
+    /// Initial amplification time.
+    pub init_amp_time: Timestamp,
+    /// Future amplification time.
+    pub future_amp_time: Timestamp,
 }
 
 impl StableSwapPool {
     pub fn new(
         id: u32,
         token_account_ids: Vec<ValidAccountId>,
+        amp_factor: u64,
         total_fee: u32,
-        exchange_fee: u32,
-        referral_fee: u32,
     ) -> Self {
-        assert_eq!(token_account_ids.len(), 2, "ERR_WRONG_TOKEN_COUNT");
+        assert_eq!(
+            token_account_ids.len() as u64,
+            math::N_COINS,
+            "ERR_WRONG_TOKEN_COUNT"
+        );
+        assert!(total_fee < FEE_DIVISOR, "ERR_FEE_TOO_LARGE");
         Self {
             token_account_ids: token_account_ids.iter().map(|a| a.clone().into()).collect(),
             amounts: vec![0u128; token_account_ids.len()],
             volumes: vec![SwapVolume::default(); token_account_ids.len()],
             total_fee,
-            exchange_fee,
-            referral_fee,
-            shares: LookupMap::new(format!("ss{}", id).into_bytes()),
+            shares: LookupMap::new(StorageKey::Shares { pool_id: id }),
             shares_total_supply: 0,
+            init_amp_factor: amp_factor,
+            future_amp_factor: amp_factor,
+            init_amp_time: 0,
+            future_amp_time: 0,
         }
     }
 
@@ -76,21 +88,32 @@ impl StableSwapPool {
             self.token_account_ids.len(),
             "ERR_WRONG_TOKEN_COUNT"
         );
-        let calc = StableSwap::new(0, 0, 0, 0, 0);
-        let new_shares = calc
+        let invariant = StableSwap::new(
+            self.init_amp_factor,
+            self.future_amp_factor,
+            env::block_timestamp(),
+            self.init_amp_time,
+            self.future_amp_time,
+        );
+        let new_shares = invariant
             .compute_lp_amount_for_deposit(
                 self.amounts[0],
                 self.amounts[1],
-                0,
-                0,
+                amounts[0],
+                amounts[1],
                 self.shares_total_supply,
                 &Fees {
-                    trade_fee: 0,
+                    trade_fee: self.total_fee as u64,
                     admin_fee: 0,
                 },
             )
             // TODO: proper error
-            .expect("ERR_FAILED");
+            .expect("ERR_CALC_FAILED");
+
+        // TODO: add slippage check on the LP tokens.
+        self.amounts[0] += amounts[0];
+        self.amounts[1] += amounts[1];
+
         self.mint_shares(sender_id, new_shares);
         new_shares
     }
@@ -120,6 +143,37 @@ impl StableSwapPool {
         let prev_shares_amount = self.shares.get(&sender_id).expect("ERR_NO_SHARES");
         assert!(prev_shares_amount >= shares, "ERR_NOT_ENOUGH_SHARES");
         let mut result = vec![];
+        let invariant = StableSwap::new(
+            self.init_amp_factor,
+            self.future_amp_factor,
+            env::block_timestamp(),
+            self.init_amp_time,
+            self.future_amp_time,
+        );
+        for (idx, min_amount) in min_amounts.iter().enumerate() {
+            if *min_amount != 0 {
+                let (amount_out, fee) = invariant
+                    .compute_withdraw_one(
+                        shares,
+                        self.shares_total_supply,
+                        self.amounts[idx],
+                        self.amounts[1 - idx],
+                        &Fees {
+                            trade_fee: self.total_fee as u64,
+                            admin_fee: 0,
+                        },
+                    )
+                    .expect("ERR_CALC");
+                assert!(amount_out >= *min_amount, "ERR_SLIPPAGE");
+                // todo: fees
+                result[idx] = amount_out;
+            }
+        }
+        for i in 0..N_COINS {
+            self.amounts[i as usize] = self.amounts[i as usize]
+                .checked_sub(result[i as usize])
+                .expect("ERR_CALC");
+        }
         result
     }
     /// Returns number of tokens in outcome, given amount.
@@ -129,8 +183,26 @@ impl StableSwapPool {
         token_in: usize,
         amount_in: Balance,
         token_out: usize,
-    ) -> Balance {
-        0
+        fees: SwapFees,
+    ) -> SwapResult {
+        let invariant = StableSwap::new(
+            self.init_amp_factor,
+            self.future_amp_factor,
+            env::block_timestamp(),
+            self.init_amp_time,
+            self.future_amp_time,
+        );
+        invariant
+            .swap_to(
+                self.amounts[token_in],
+                amount_in,
+                self.amounts[token_out],
+                &Fees {
+                    trade_fee: self.total_fee as u64,
+                    admin_fee: 0,
+                },
+            )
+            .expect("ERR_CALC")
     }
 
     /// Returns how much token you will receive if swap `token_amount_in` of `token_in` for `token_out`.
@@ -139,12 +211,15 @@ impl StableSwapPool {
         token_in: &AccountId,
         amount_in: Balance,
         token_out: &AccountId,
+        fees: SwapFees,
     ) -> Balance {
         self.internal_get_return(
             self.token_index(token_in),
             amount_in,
             self.token_index(token_out),
+            fees,
         )
+        .amount_swapped
     }
 
     /// Swap `token_amount_in` of `token_in` token into `token_out` and return how much was received.
@@ -155,10 +230,26 @@ impl StableSwapPool {
         amount_in: Balance,
         token_out: &AccountId,
         min_amount_out: Balance,
-        exchange_id: &AccountId,
-        referral_id: &Option<AccountId>,
+        fees: SwapFees,
     ) -> Balance {
-        0
+        let in_idx = self.token_index(token_in);
+        let out_idx = self.token_index(token_out);
+        let result = self.internal_get_return(in_idx, amount_in, out_idx, fees);
+        assert!(result.amount_swapped >= min_amount_out, "ERR_MIN_AMOUNT");
+        env::log(
+            format!(
+                "Swapped {} {} for {} {}",
+                amount_in, token_in, result.amount_swapped, token_out
+            )
+            .as_bytes(),
+        );
+
+        self.amounts[in_idx] = result.new_source_amount;
+        self.amounts[out_idx] = result.new_destination_amount;
+
+        // TODO: add admin / referral fee here.
+
+        result.amount_swapped
     }
 
     /// Register given account with 0 balance in shares.
@@ -199,6 +290,18 @@ impl StableSwapPool {
     pub fn tokens(&self) -> &[AccountId] {
         &self.token_account_ids
     }
+
+    /// [Admin function] increase the amplification factor.
+    pub fn ramp_amplification(&mut self, future_amp_factor: u64, future_amp_time: Timestamp) {
+        // TODO: proper implementation
+        self.future_amp_factor = future_amp_time;
+        self.future_amp_factor = future_amp_factor;
+    }
+
+    /// [Admin function] Stop increase of amplification factor.
+    pub fn stop_ramp_amplification(&mut self) {
+        // TODO: implement
+    }
 }
 
 #[cfg(test)]
@@ -211,12 +314,69 @@ mod tests {
     #[test]
     fn test_basics() {
         let mut context = VMContextBuilder::new();
-        context.predecessor_account_id(accounts(0));
-        testing_env!(context.build());
-        let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 0, 0, 0);
+        testing_env!(context.predecessor_account_id(accounts(0)).build());
+        let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 1, 0);
+        assert_eq!(
+            pool.tokens(),
+            vec![accounts(1).to_string(), accounts(2).to_string()]
+        );
+
+        let mut amounts = vec![to_yocto("5"), to_yocto("10")];
+        let _ = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
+
+        // Add only one side of the capital.
+        let mut amounts2 = vec![to_yocto("5"), to_yocto("0")];
+        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts2);
+
+        // Withdraw on another side of the capital.
+        let amounts_out = pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![0, 1]);
+
+        // assert?
+
+        let out = pool.swap(
+            accounts(1).as_ref(),
+            to_yocto("1"),
+            accounts(2).as_ref(),
+            1,
+            SwapFees {
+                exchange_fee: 0,
+                exchange_id: accounts(1).as_ref().clone(),
+                referral_fee: 0,
+                referral_id: None,
+            },
+        );
+
+        // assert out to_yocto("2")
+        assert_eq!(pool.amounts, vec![to_yocto("6"), to_yocto("8")]);
+    }
+
+    /// Test that adding and then removing all of the liquidity leaves the pool empty and with no shares.
+    #[test]
+    fn test_add_transfer_remove_liquidity() {
+        let mut context = VMContextBuilder::new();
+        testing_env!(context.predecessor_account_id(accounts(0)).build());
+        let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 1, 0);
         let mut amounts = vec![to_yocto("5"), to_yocto("10")];
         let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
         assert_eq!(amounts, vec![to_yocto("5"), to_yocto("10")]);
         assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 1);
+        assert_eq!(pool.share_total_balance(), 1);
+
+        // Move shares to another account.
+        pool.share_transfer(accounts(0).as_ref(), accounts(3).as_ref(), num_shares);
+        assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 0);
+        assert_eq!(pool.share_balance_of(accounts(3).as_ref()), num_shares);
+        assert_eq!(pool.share_total_balance(), 1);
+
+        // Remove all liquidity.
+        testing_env!(context.predecessor_account_id(accounts(1)).build());
+        let out_amounts = pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1]);
+
+        // Check it's all taken out.
+        assert_eq!(amounts, out_amounts);
+        assert_eq!(pool.share_total_balance(), 0);
+        assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 0);
+        assert_eq!(pool.share_balance_of(accounts(3).as_ref()), 0);
+        assert_eq!(pool.amounts, vec![0, 0]);
     }
 }

From a9332fc9bf2ff41219beb5371fbf4c346aaf158b Mon Sep 17 00:00:00 2001
From: referencedev <goi65io3903d@protonmail.com>
Date: Mon, 13 Sep 2021 15:03:02 -0700
Subject: [PATCH 3/6] Basic tests pass. Fees are WIP

---
 ref-exchange/src/simple_pool.rs      |   3 +-
 ref-exchange/src/stable_swap/math.rs |   1 +
 ref-exchange/src/stable_swap/mod.rs  | 125 ++++++++++++++++++---------
 test-token/src/lib.rs                |   2 +-
 4 files changed, 87 insertions(+), 44 deletions(-)

diff --git a/ref-exchange/src/simple_pool.rs b/ref-exchange/src/simple_pool.rs
index af08715..7eccea7 100644
--- a/ref-exchange/src/simple_pool.rs
+++ b/ref-exchange/src/simple_pool.rs
@@ -191,7 +191,7 @@ impl SimplePool {
             result.push(amount);
         }
         if prev_shares_amount == shares {
-            // [AUDIT_13] never unregister a LP when he remove liqudity.
+            // [AUDIT_13] Never unregister an LP when liquidity is removed.
             self.shares.insert(&sender_id, &0);
         } else {
             self.shares
@@ -277,6 +277,7 @@ impl SimplePool {
         min_amount_out: Balance,
         fees: SwapFees,
     ) -> Balance {
+        assert_ne!(token_in, token_out, "ERR_SAME_TOKEN_SWAP");
         let in_idx = self.token_index(token_in);
         let out_idx = self.token_index(token_out);
         let amount_out = self.internal_get_return(in_idx, amount_in, out_idx);
diff --git a/ref-exchange/src/stable_swap/math.rs b/ref-exchange/src/stable_swap/math.rs
index a6c25c6..49b3216 100644
--- a/ref-exchange/src/stable_swap/math.rs
+++ b/ref-exchange/src/stable_swap/math.rs
@@ -37,6 +37,7 @@ impl Fees {
 }
 
 /// Encodes all results of swapping from a source token to a destination token.
+#[derive(Debug)]
 pub struct SwapResult {
     /// New amount of source token.
     pub new_source_amount: Balance,
diff --git a/ref-exchange/src/stable_swap/mod.rs b/ref-exchange/src/stable_swap/mod.rs
index 92f8744..b981d25 100644
--- a/ref-exchange/src/stable_swap/mod.rs
+++ b/ref-exchange/src/stable_swap/mod.rs
@@ -95,20 +95,28 @@ impl StableSwapPool {
             self.init_amp_time,
             self.future_amp_time,
         );
-        let new_shares = invariant
-            .compute_lp_amount_for_deposit(
-                self.amounts[0],
-                self.amounts[1],
-                amounts[0],
-                amounts[1],
-                self.shares_total_supply,
-                &Fees {
-                    trade_fee: self.total_fee as u64,
-                    admin_fee: 0,
-                },
-            )
-            // TODO: proper error
-            .expect("ERR_CALC_FAILED");
+        let new_shares = if self.shares_total_supply == 0 {
+            // Bootstrapping the pool.
+            invariant
+                .compute_d(amounts[0], amounts[1])
+                .expect("ERR_CALC_FAILED")
+                .as_u128()
+        } else {
+            invariant
+                .compute_lp_amount_for_deposit(
+                    amounts[0],
+                    amounts[1],
+                    self.amounts[0],
+                    self.amounts[1],
+                    self.shares_total_supply,
+                    &Fees {
+                        trade_fee: self.total_fee as u64,
+                        admin_fee: 0,
+                    },
+                )
+                // TODO: proper error
+                .expect("ERR_CALC_FAILED")
+        };
 
         // TODO: add slippage check on the LP tokens.
         self.amounts[0] += amounts[0];
@@ -142,7 +150,7 @@ impl StableSwapPool {
         );
         let prev_shares_amount = self.shares.get(&sender_id).expect("ERR_NO_SHARES");
         assert!(prev_shares_amount >= shares, "ERR_NOT_ENOUGH_SHARES");
-        let mut result = vec![];
+        let mut result = vec![0u128; N_COINS as usize];
         let invariant = StableSwap::new(
             self.init_amp_factor,
             self.future_amp_factor,
@@ -174,6 +182,22 @@ impl StableSwapPool {
                 .checked_sub(result[i as usize])
                 .expect("ERR_CALC");
         }
+        // Never unregister an LP when liquidity is removed.
+        self.shares
+            .insert(&sender_id, &(prev_shares_amount - shares));
+        env::log(
+            format!(
+                "{} shares of liquidity removed: receive back {:?}",
+                shares,
+                result
+                    .iter()
+                    .zip(self.token_account_ids.iter())
+                    .map(|(amount, token_id)| format!("{} {}", amount, token_id))
+                    .collect::<Vec<String>>(),
+            )
+            .as_bytes(),
+        );
+        self.shares_total_supply -= shares;
         result
     }
     /// Returns number of tokens in outcome, given amount.
@@ -194,8 +218,8 @@ impl StableSwapPool {
         );
         invariant
             .swap_to(
-                self.amounts[token_in],
                 amount_in,
+                self.amounts[token_in],
                 self.amounts[token_out],
                 &Fees {
                     trade_fee: self.total_fee as u64,
@@ -232,9 +256,11 @@ impl StableSwapPool {
         min_amount_out: Balance,
         fees: SwapFees,
     ) -> Balance {
+        assert_ne!(token_in, token_out, "ERR_SAME_TOKEN_SWAP");
         let in_idx = self.token_index(token_in);
         let out_idx = self.token_index(token_out);
         let result = self.internal_get_return(in_idx, amount_in, out_idx, fees);
+        println!("{:?}", result);
         assert!(result.amount_swapped >= min_amount_out, "ERR_MIN_AMOUNT");
         env::log(
             format!(
@@ -311,6 +337,26 @@ mod tests {
     use near_sdk::{testing_env, MockedBlockchain};
     use near_sdk_sim::to_yocto;
 
+    fn swap(
+        pool: &mut StableSwapPool,
+        token_in: usize,
+        amount_in: Balance,
+        token_out: usize,
+    ) -> Balance {
+        pool.swap(
+            accounts(token_in).as_ref(),
+            amount_in,
+            accounts(token_out).as_ref(),
+            1,
+            SwapFees {
+                exchange_fee: 0,
+                exchange_id: accounts(1).as_ref().clone(),
+                referral_fee: 0,
+                referral_id: None,
+            },
+        )
+    }
+
     #[test]
     fn test_basics() {
         let mut context = VMContextBuilder::new();
@@ -324,30 +370,20 @@ mod tests {
         let mut amounts = vec![to_yocto("5"), to_yocto("10")];
         let _ = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
 
+        let out = swap(&mut pool, 1, to_yocto("1"), 2);
+        assert_eq!(out, 1313682630255414606428571);
+        assert_eq!(pool.amounts, vec![to_yocto("6"), 8686317369744585393571429]);
+        let out2 = swap(&mut pool, 2, out, 1);
+        assert_eq!(out2, to_yocto("1") + 2); // due to precision difference.
+        assert_eq!(pool.amounts, vec![to_yocto("5") - 2, to_yocto("10")]);
+
         // Add only one side of the capital.
         let mut amounts2 = vec![to_yocto("5"), to_yocto("0")];
         let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts2);
 
         // Withdraw on another side of the capital.
         let amounts_out = pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![0, 1]);
-
-        // assert?
-
-        let out = pool.swap(
-            accounts(1).as_ref(),
-            to_yocto("1"),
-            accounts(2).as_ref(),
-            1,
-            SwapFees {
-                exchange_fee: 0,
-                exchange_id: accounts(1).as_ref().clone(),
-                referral_fee: 0,
-                referral_id: None,
-            },
-        );
-
-        // assert out to_yocto("2")
-        assert_eq!(pool.amounts, vec![to_yocto("6"), to_yocto("8")]);
+        assert_eq!(amounts_out, vec![0, to_yocto("5")]);
     }
 
     /// Test that adding and then removing all of the liquidity leaves the pool empty and with no shares.
@@ -359,24 +395,29 @@ mod tests {
         let mut amounts = vec![to_yocto("5"), to_yocto("10")];
         let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
         assert_eq!(amounts, vec![to_yocto("5"), to_yocto("10")]);
-        assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 1);
-        assert_eq!(pool.share_total_balance(), 1);
+        assert!(num_shares > 1);
+        assert_eq!(num_shares, pool.share_balance_of(accounts(0).as_ref()));
+        assert_eq!(pool.share_total_balance(), num_shares);
 
         // Move shares to another account.
+        pool.share_register(accounts(3).as_ref());
         pool.share_transfer(accounts(0).as_ref(), accounts(3).as_ref(), num_shares);
         assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 0);
         assert_eq!(pool.share_balance_of(accounts(3).as_ref()), num_shares);
-        assert_eq!(pool.share_total_balance(), 1);
+        assert_eq!(pool.share_total_balance(), num_shares);
 
         // Remove all liquidity.
-        testing_env!(context.predecessor_account_id(accounts(1)).build());
-        let out_amounts = pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1]);
+        testing_env!(context.predecessor_account_id(accounts(3)).build());
+        let out_amounts = pool.remove_liquidity(accounts(3).as_ref(), num_shares, vec![1, 1]);
 
-        // Check it's all taken out.
-        assert_eq!(amounts, out_amounts);
+        // Check it's all taken out. Due to precision there is ~1 yN.
+        assert_eq!(
+            vec![amounts[0], amounts[1]],
+            vec![out_amounts[0] + 1, out_amounts[1] + 1]
+        );
         assert_eq!(pool.share_total_balance(), 0);
         assert_eq!(pool.share_balance_of(accounts(0).as_ref()), 0);
         assert_eq!(pool.share_balance_of(accounts(3).as_ref()), 0);
-        assert_eq!(pool.amounts, vec![0, 0]);
+        assert_eq!(pool.amounts, vec![1, 1]);
     }
 }
diff --git a/test-token/src/lib.rs b/test-token/src/lib.rs
index bb438fc..5ea935e 100644
--- a/test-token/src/lib.rs
+++ b/test-token/src/lib.rs
@@ -4,7 +4,7 @@ use near_contract_standards::fungible_token::metadata::{
 use near_contract_standards::fungible_token::FungibleToken;
 use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::json_types::{ValidAccountId, U128};
-use near_sdk::{env, near_bindgen, AccountId, PanicOnDefault, PromiseOrValue};
+use near_sdk::{near_bindgen, AccountId, PanicOnDefault, PromiseOrValue};
 
 near_sdk::setup_alloc!();
 

From 0a789f29b7a27ab9f63b7f740989b7201b8e2507 Mon Sep 17 00:00:00 2001
From: referencedev <goi65io3903d@protonmail.com>
Date: Tue, 14 Sep 2021 15:29:11 -0700
Subject: [PATCH 4/6] Handling fees across add/withdraw liquidity and swap

---
 ref-exchange/src/fees.rs             |  17 ++++-
 ref-exchange/src/lib.rs              |  17 ++++-
 ref-exchange/src/pool.rs             |  16 +++--
 ref-exchange/src/stable_swap/math.rs |  47 ++++++------
 ref-exchange/src/stable_swap/mod.rs  | 104 +++++++++++++++------------
 5 files changed, 131 insertions(+), 70 deletions(-)

diff --git a/ref-exchange/src/fees.rs b/ref-exchange/src/fees.rs
index c6c0ae8..3670215 100644
--- a/ref-exchange/src/fees.rs
+++ b/ref-exchange/src/fees.rs
@@ -1,4 +1,4 @@
-use near_sdk::AccountId;
+use near_sdk::{env, AccountId};
 
 /// Maintain information about fees.
 pub struct SwapFees {
@@ -9,3 +9,18 @@ pub struct SwapFees {
     pub exchange_id: AccountId,
     pub referral_id: Option<AccountId>,
 }
+
+impl SwapFees {
+    pub fn new(exchange_fee: u32) -> Self {
+        SwapFees {
+            exchange_fee,
+            exchange_id: env::current_account_id(),
+            referral_fee: 0,
+            referral_id: None,
+        }
+    }
+
+    pub fn zero() -> Self {
+        Self::new(0)
+    }
+}
diff --git a/ref-exchange/src/lib.rs b/ref-exchange/src/lib.rs
index 1ae08ca..428a4bd 100644
--- a/ref-exchange/src/lib.rs
+++ b/ref-exchange/src/lib.rs
@@ -160,7 +160,16 @@ impl Contract {
         let mut amounts: Vec<u128> = amounts.into_iter().map(|amount| amount.into()).collect();
         let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
         // Add amounts given to liquidity first. It will return the balanced amounts.
-        pool.add_liquidity(&sender_id, &mut amounts);
+        pool.add_liquidity(
+            &sender_id,
+            &mut amounts,
+            SwapFees {
+                exchange_fee: self.exchange_fee,
+                exchange_id: self.owner_id.clone(),
+                referral_fee: 0,
+                referral_id: None,
+            },
+        );
         if let Some(min_amounts) = min_amounts {
             // Check that all amounts are above request min amounts in case of front running that changes the exchange rate.
             for (amount, min_amount) in amounts.iter().zip(min_amounts.iter()) {
@@ -192,6 +201,12 @@ impl Contract {
                 .into_iter()
                 .map(|amount| amount.into())
                 .collect(),
+            SwapFees {
+                exchange_fee: self.exchange_fee,
+                exchange_id: self.owner_id.clone(),
+                referral_fee: 0,
+                referral_id: None,
+            },
         );
         self.pools.replace(pool_id, &pool);
         let tokens = pool.tokens();
diff --git a/ref-exchange/src/pool.rs b/ref-exchange/src/pool.rs
index 9a78984..4016326 100644
--- a/ref-exchange/src/pool.rs
+++ b/ref-exchange/src/pool.rs
@@ -33,10 +33,15 @@ impl Pool {
 
     /// Adds liquidity into underlying pool.
     /// Updates amounts to amount kept in the pool.
-    pub fn add_liquidity(&mut self, sender_id: &AccountId, amounts: &mut Vec<Balance>) -> Balance {
+    pub fn add_liquidity(
+        &mut self,
+        sender_id: &AccountId,
+        amounts: &mut Vec<Balance>,
+        fees: SwapFees,
+    ) -> Balance {
         match self {
             Pool::SimplePool(pool) => pool.add_liquidity(sender_id, amounts),
-            Pool::StableSwapPool(pool) => pool.add_liquidity(sender_id, amounts),
+            Pool::StableSwapPool(pool) => pool.add_liquidity(sender_id, amounts, &fees),
         }
     }
 
@@ -46,10 +51,13 @@ impl Pool {
         sender_id: &AccountId,
         shares: Balance,
         min_amounts: Vec<Balance>,
+        fees: SwapFees,
     ) -> Vec<Balance> {
         match self {
             Pool::SimplePool(pool) => pool.remove_liquidity(sender_id, shares, min_amounts),
-            Pool::StableSwapPool(pool) => pool.remove_liquidity(sender_id, shares, min_amounts),
+            Pool::StableSwapPool(pool) => {
+                pool.remove_liquidity(sender_id, shares, min_amounts, &fees)
+            }
         }
     }
 
@@ -97,7 +105,7 @@ impl Pool {
                 pool.swap(token_in, amount_in, token_out, min_amount_out, fees)
             }
             Pool::StableSwapPool(pool) => {
-                pool.swap(token_in, amount_in, token_out, min_amount_out, fees)
+                pool.swap(token_in, amount_in, token_out, min_amount_out, &fees)
             }
         }
     }
diff --git a/ref-exchange/src/stable_swap/math.rs b/ref-exchange/src/stable_swap/math.rs
index 49b3216..1a81a81 100644
--- a/ref-exchange/src/stable_swap/math.rs
+++ b/ref-exchange/src/stable_swap/math.rs
@@ -2,37 +2,44 @@
 ///! Large part of the code was taken from https://github.com/saber-hq/stable-swap/blob/master/stable-swap-math/src/curve.rs
 use near_sdk::{Balance, Timestamp};
 
-use crate::utils::U256;
+use crate::fees::SwapFees;
+use crate::utils::{FEE_DIVISOR, U256};
 
 /// Number of coins in the pool.
-pub const N_COINS: u64 = 2;
-/// Timestamp at 0
-pub const ZERO_TS: i64 = 0;
-/// Minimum ramp duration
+pub const N_COINS: u32 = 2;
+/// Minimum ramp duration.
 pub const MIN_RAMP_DURATION: i64 = 86400;
-/// Min amplification coefficient
+/// Min amplification coefficient.
 pub const MIN_AMP: u64 = 1;
-/// Max amplification coefficient
+/// Max amplification coefficient.
 pub const MAX_AMP: u64 = 1_000_000;
 /// Max number of tokens to swap at once.
 pub const MAX_TOKENS_IN: u64 = u64::MAX >> 4;
 
+/// Stable Swap Fee calculator.
 pub struct Fees {
-    pub trade_fee: u64,
-    pub admin_fee: u64,
+    pub trade_fee: u32,
+    pub admin_fee: u32,
 }
 
 impl Fees {
-    pub fn trade_fee(&self, amount: Balance) -> Option<Balance> {
-        Some(0)
+    pub fn new(total_fee: u32, fees: &SwapFees) -> Self {
+        Self {
+            trade_fee: total_fee - fees.exchange_fee,
+            admin_fee: fees.exchange_fee,
+        }
+    }
+    pub fn trade_fee(&self, amount: Balance) -> Balance {
+        amount * (self.trade_fee as u128) / (FEE_DIVISOR as u128)
     }
 
-    pub fn admin_trade_fee(&self, amount: Balance) -> Option<Balance> {
-        Some(0)
+    pub fn admin_trade_fee(&self, amount: Balance) -> Balance {
+        amount * (self.admin_fee as u128) / (FEE_DIVISOR as u128)
     }
 
-    pub fn normalized_trade_fee(&self, num_coins: u64, amount: Balance) -> Option<Balance> {
-        Some(0)
+    pub fn normalized_trade_fee(&self, num_coins: u32, amount: Balance) -> Balance {
+        let adjusted_trade_fee = (self.trade_fee * num_coins) / (4 * (num_coins - 1));
+        amount * (adjusted_trade_fee as u128) / (FEE_DIVISOR as u128)
     }
 }
 
@@ -211,7 +218,7 @@ impl StableSwap {
                 } else {
                     new_balances[i].checked_sub(ideal_balance)?
                 };
-                let fee = fees.normalized_trade_fee(N_COINS, difference)?;
+                let fee = fees.normalized_trade_fee(N_COINS, difference);
                 new_balances[i] = new_balances[i].checked_sub(fee)?;
             }
 
@@ -304,10 +311,10 @@ impl StableSwap {
         )?;
         // new_base_amount = swap_base_amount - expected_base_amount * fee / fee_denominator;
         let new_base_amount = swap_base_amount
-            .checked_sub(fees.normalized_trade_fee(N_COINS, expected_base_amount)?)?;
+            .checked_sub(fees.normalized_trade_fee(N_COINS, expected_base_amount))?;
         // new_quote_amount = swap_quote_amount - expected_quote_amount * fee / fee_denominator;
         let new_quote_amount = swap_quote_amount
-            .checked_sub(fees.normalized_trade_fee(N_COINS, expected_quote_amount)?)?;
+            .checked_sub(fees.normalized_trade_fee(N_COINS, expected_quote_amount))?;
         let dy = new_base_amount
             .checked_sub(self.compute_y(new_quote_amount, d_1))?
             .checked_sub(1)?; // Withdraw less to account for rounding errors
@@ -329,8 +336,8 @@ impl StableSwap {
             self.compute_d(swap_source_amount, swap_destination_amount)?,
         );
         let dy = swap_destination_amount.checked_sub(y)?;
-        let dy_fee = fees.trade_fee(dy)?;
-        let admin_fee = fees.admin_trade_fee(dy_fee)?;
+        let dy_fee = fees.trade_fee(dy);
+        let admin_fee = fees.admin_trade_fee(dy_fee);
 
         let amount_swapped = dy.checked_sub(dy_fee)?;
         let new_destination_amount = swap_destination_amount
diff --git a/ref-exchange/src/stable_swap/mod.rs b/ref-exchange/src/stable_swap/mod.rs
index b981d25..c75b5b9 100644
--- a/ref-exchange/src/stable_swap/mod.rs
+++ b/ref-exchange/src/stable_swap/mod.rs
@@ -1,13 +1,13 @@
+use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 use near_sdk::collections::LookupMap;
 use near_sdk::json_types::ValidAccountId;
 use near_sdk::{env, AccountId, Balance, Timestamp};
 
 use crate::errors::{ERR13_LP_NOT_REGISTERED, ERR14_LP_ALREADY_REGISTERED};
 use crate::fees::SwapFees;
-use crate::stable_swap::math::{Fees, StableSwap, SwapResult, N_COINS};
+use crate::stable_swap::math::{Fees, StableSwap, SwapResult, MAX_AMP, MIN_AMP, N_COINS};
 use crate::utils::{add_to_collection, SwapVolume, FEE_DIVISOR};
 use crate::StorageKey;
-use near_sdk::borsh::{self, BorshDeserialize, BorshSerialize};
 
 mod math;
 
@@ -42,8 +42,12 @@ impl StableSwapPool {
         amp_factor: u64,
         total_fee: u32,
     ) -> Self {
+        assert!(
+            amp_factor >= MIN_AMP && amp_factor <= MAX_AMP,
+            "ERR_WRONG_AMP"
+        );
         assert_eq!(
-            token_account_ids.len() as u64,
+            token_account_ids.len() as u32,
             math::N_COINS,
             "ERR_WRONG_TOKEN_COUNT"
         );
@@ -82,7 +86,12 @@ impl StableSwapPool {
 
     /// Add liquidity into the pool.
     /// Allows to add liquidity of a subset of tokens.
-    pub fn add_liquidity(&mut self, sender_id: &AccountId, amounts: &mut Vec<Balance>) -> Balance {
+    pub fn add_liquidity(
+        &mut self,
+        sender_id: &AccountId,
+        amounts: &mut Vec<Balance>,
+        fees: &SwapFees,
+    ) -> Balance {
         assert_eq!(
             amounts.len(),
             self.token_account_ids.len(),
@@ -109,10 +118,7 @@ impl StableSwapPool {
                     self.amounts[0],
                     self.amounts[1],
                     self.shares_total_supply,
-                    &Fees {
-                        trade_fee: self.total_fee as u64,
-                        admin_fee: 0,
-                    },
+                    &Fees::new(self.total_fee, &fees),
                 )
                 // TODO: proper error
                 .expect("ERR_CALC_FAILED")
@@ -126,15 +132,6 @@ impl StableSwapPool {
         new_shares
     }
 
-    /// Mint new shares for given user.
-    fn mint_shares(&mut self, account_id: &AccountId, shares: Balance) {
-        if shares == 0 {
-            return;
-        }
-        self.shares_total_supply += shares;
-        add_to_collection(&mut self.shares, &account_id, shares);
-    }
-
     /// Remove liquidity from the pool.
     /// Allows to remove liquidity of a subset of tokens, by providing 0 in `min_amount` for the tokens to not withdraw.
     pub fn remove_liquidity(
@@ -142,6 +139,7 @@ impl StableSwapPool {
         sender_id: &AccountId,
         shares: Balance,
         min_amounts: Vec<Balance>,
+        fees: &SwapFees,
     ) -> Vec<Balance> {
         assert_eq!(
             min_amounts.len(),
@@ -158,6 +156,8 @@ impl StableSwapPool {
             self.init_amp_time,
             self.future_amp_time,
         );
+        let mut fee_amounts = vec![0u128; N_COINS as usize];
+        let stable_swap_fees = Fees::new(self.total_fee, &fees);
         for (idx, min_amount) in min_amounts.iter().enumerate() {
             if *min_amount != 0 {
                 let (amount_out, fee) = invariant
@@ -166,17 +166,15 @@ impl StableSwapPool {
                         self.shares_total_supply,
                         self.amounts[idx],
                         self.amounts[1 - idx],
-                        &Fees {
-                            trade_fee: self.total_fee as u64,
-                            admin_fee: 0,
-                        },
+                        &stable_swap_fees,
                     )
                     .expect("ERR_CALC");
                 assert!(amount_out >= *min_amount, "ERR_SLIPPAGE");
-                // todo: fees
+                fee_amounts[idx] += fee;
                 result[idx] = amount_out;
             }
         }
+        println!("fees: {:?}", fee_amounts);
         for i in 0..N_COINS {
             self.amounts[i as usize] = self.amounts[i as usize]
                 .checked_sub(result[i as usize])
@@ -207,7 +205,7 @@ impl StableSwapPool {
         token_in: usize,
         amount_in: Balance,
         token_out: usize,
-        fees: SwapFees,
+        fees: &SwapFees,
     ) -> SwapResult {
         let invariant = StableSwap::new(
             self.init_amp_factor,
@@ -221,10 +219,7 @@ impl StableSwapPool {
                 amount_in,
                 self.amounts[token_in],
                 self.amounts[token_out],
-                &Fees {
-                    trade_fee: self.total_fee as u64,
-                    admin_fee: 0,
-                },
+                &Fees::new(self.total_fee, &fees),
             )
             .expect("ERR_CALC")
     }
@@ -235,13 +230,13 @@ impl StableSwapPool {
         token_in: &AccountId,
         amount_in: Balance,
         token_out: &AccountId,
-        fees: SwapFees,
+        fees: &SwapFees,
     ) -> Balance {
         self.internal_get_return(
             self.token_index(token_in),
             amount_in,
             self.token_index(token_out),
-            fees,
+            &fees,
         )
         .amount_swapped
     }
@@ -254,13 +249,12 @@ impl StableSwapPool {
         amount_in: Balance,
         token_out: &AccountId,
         min_amount_out: Balance,
-        fees: SwapFees,
+        fees: &SwapFees,
     ) -> Balance {
         assert_ne!(token_in, token_out, "ERR_SAME_TOKEN_SWAP");
         let in_idx = self.token_index(token_in);
         let out_idx = self.token_index(token_out);
-        let result = self.internal_get_return(in_idx, amount_in, out_idx, fees);
-        println!("{:?}", result);
+        let result = self.internal_get_return(in_idx, amount_in, out_idx, &fees);
         assert!(result.amount_swapped >= min_amount_out, "ERR_MIN_AMOUNT");
         env::log(
             format!(
@@ -278,6 +272,15 @@ impl StableSwapPool {
         result.amount_swapped
     }
 
+    /// Mint new shares for given user.
+    fn mint_shares(&mut self, account_id: &AccountId, shares: Balance) {
+        if shares == 0 {
+            return;
+        }
+        self.shares_total_supply += shares;
+        add_to_collection(&mut self.shares, &account_id, shares);
+    }
+
     /// Register given account with 0 balance in shares.
     /// Storage payment should be checked by caller.
     pub fn share_register(&mut self, account_id: &AccountId) {
@@ -332,11 +335,12 @@ impl StableSwapPool {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use near_sdk::test_utils::{accounts, VMContextBuilder};
     use near_sdk::{testing_env, MockedBlockchain};
     use near_sdk_sim::to_yocto;
 
+    use super::*;
+
     fn swap(
         pool: &mut StableSwapPool,
         token_in: usize,
@@ -348,12 +352,7 @@ mod tests {
             amount_in,
             accounts(token_out).as_ref(),
             1,
-            SwapFees {
-                exchange_fee: 0,
-                exchange_id: accounts(1).as_ref().clone(),
-                referral_fee: 0,
-                referral_id: None,
-            },
+            &SwapFees::zero(),
         )
     }
 
@@ -361,6 +360,7 @@ mod tests {
     fn test_basics() {
         let mut context = VMContextBuilder::new();
         testing_env!(context.predecessor_account_id(accounts(0)).build());
+        let fees = SwapFees::zero();
         let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 1, 0);
         assert_eq!(
             pool.tokens(),
@@ -368,7 +368,7 @@ mod tests {
         );
 
         let mut amounts = vec![to_yocto("5"), to_yocto("10")];
-        let _ = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
+        let _ = pool.add_liquidity(accounts(0).as_ref(), &mut amounts, &fees);
 
         let out = swap(&mut pool, 1, to_yocto("1"), 2);
         assert_eq!(out, 1313682630255414606428571);
@@ -379,13 +379,27 @@ mod tests {
 
         // Add only one side of the capital.
         let mut amounts2 = vec![to_yocto("5"), to_yocto("0")];
-        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts2);
+        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts2, &fees);
 
         // Withdraw on another side of the capital.
-        let amounts_out = pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![0, 1]);
+        let amounts_out =
+            pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![0, 1], &fees);
         assert_eq!(amounts_out, vec![0, to_yocto("5")]);
     }
 
+    #[test]
+    fn test_with_fees() {
+        let mut context = VMContextBuilder::new();
+        testing_env!(context.predecessor_account_id(accounts(0)).build());
+        let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 1, 2000);
+        let mut amounts = vec![to_yocto("5"), to_yocto("10")];
+        let fees = SwapFees::new(1000);
+        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts, &fees);
+        let amounts_out =
+            pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1], &fees);
+        println!("amount out: {:?}", amounts_out);
+    }
+
     /// Test that adding and then removing all of the liquidity leaves the pool empty and with no shares.
     #[test]
     fn test_add_transfer_remove_liquidity() {
@@ -393,7 +407,8 @@ mod tests {
         testing_env!(context.predecessor_account_id(accounts(0)).build());
         let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 1, 0);
         let mut amounts = vec![to_yocto("5"), to_yocto("10")];
-        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts);
+        let fees = SwapFees::zero();
+        let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts, &fees);
         assert_eq!(amounts, vec![to_yocto("5"), to_yocto("10")]);
         assert!(num_shares > 1);
         assert_eq!(num_shares, pool.share_balance_of(accounts(0).as_ref()));
@@ -408,7 +423,8 @@ mod tests {
 
         // Remove all liquidity.
         testing_env!(context.predecessor_account_id(accounts(3)).build());
-        let out_amounts = pool.remove_liquidity(accounts(3).as_ref(), num_shares, vec![1, 1]);
+        let out_amounts =
+            pool.remove_liquidity(accounts(3).as_ref(), num_shares, vec![1, 1], &fees);
 
         // Check it's all taken out. Due to precision there is ~1 yN.
         assert_eq!(

From f6343f74ea6d267b5af8f6e3c6b92827f07c9725 Mon Sep 17 00:00:00 2001
From: referencedev <goi65io3903d@protonmail.com>
Date: Fri, 17 Sep 2021 07:39:05 -0700
Subject: [PATCH 5/6] Adding start/stop ramp for amplification

---
 ref-exchange/src/stable_swap/math.rs |  31 +++---
 ref-exchange/src/stable_swap/mod.rs  | 141 ++++++++++++++++++++++-----
 2 files changed, 132 insertions(+), 40 deletions(-)

diff --git a/ref-exchange/src/stable_swap/math.rs b/ref-exchange/src/stable_swap/math.rs
index 1a81a81..e5ac8e9 100644
--- a/ref-exchange/src/stable_swap/math.rs
+++ b/ref-exchange/src/stable_swap/math.rs
@@ -8,13 +8,13 @@ use crate::utils::{FEE_DIVISOR, U256};
 /// Number of coins in the pool.
 pub const N_COINS: u32 = 2;
 /// Minimum ramp duration.
-pub const MIN_RAMP_DURATION: i64 = 86400;
+pub const MIN_RAMP_DURATION: Timestamp = 86400;
 /// Min amplification coefficient.
-pub const MIN_AMP: u64 = 1;
+pub const MIN_AMP: u128 = 1;
 /// Max amplification coefficient.
-pub const MAX_AMP: u64 = 1_000_000;
-/// Max number of tokens to swap at once.
-pub const MAX_TOKENS_IN: u64 = u64::MAX >> 4;
+pub const MAX_AMP: u128 = 1_000_000;
+/// Max amplification change.
+pub const MAX_AMP_CHANGE: u128 = 10;
 
 /// Stable Swap Fee calculator.
 pub struct Fees {
@@ -30,6 +30,11 @@ impl Fees {
         }
     }
     pub fn trade_fee(&self, amount: Balance) -> Balance {
+        println!(
+            "trade fee: {} {}",
+            amount * (self.trade_fee as u128) / (FEE_DIVISOR as u128),
+            amount
+        );
         amount * (self.trade_fee as u128) / (FEE_DIVISOR as u128)
     }
 
@@ -61,9 +66,9 @@ pub struct SwapResult {
 /// The StableSwap invariant calculator.
 pub struct StableSwap {
     /// Initial amplification coefficient (A)
-    initial_amp_factor: u64,
+    initial_amp_factor: u128,
     /// Target amplification coefficient (A)
-    target_amp_factor: u64,
+    target_amp_factor: u128,
     /// Current unix timestamp
     current_ts: Timestamp,
     /// Ramp A start timestamp
@@ -74,8 +79,8 @@ pub struct StableSwap {
 
 impl StableSwap {
     pub fn new(
-        initial_amp_factor: u64,
-        target_amp_factor: u64,
+        initial_amp_factor: u128,
+        target_amp_factor: u128,
         current_ts: Timestamp,
         start_ramp_ts: Timestamp,
         stop_ramp_ts: Timestamp,
@@ -91,7 +96,7 @@ impl StableSwap {
 
     fn compute_next_d(
         &self,
-        amp_factor: u64,
+        amp_factor: u128,
         d_init: U256,
         d_prod: U256,
         sum_x: Balance,
@@ -124,7 +129,7 @@ impl StableSwap {
                     .checked_sub(self.initial_amp_factor)?;
                 let amp_delta = (amp_range as u128)
                     .checked_mul(time_delta as u128)?
-                    .checked_div(time_range as u128)? as u64;
+                    .checked_div(time_range as u128)?;
                 self.initial_amp_factor
                     .checked_add(amp_delta)
                     .map(|x| x as u128)
@@ -135,7 +140,7 @@ impl StableSwap {
                     .checked_sub(self.target_amp_factor)?;
                 let amp_delta = (amp_range as u128)
                     .checked_mul(time_delta as u128)?
-                    .checked_div(time_range as u128)? as u64;
+                    .checked_div(time_range as u128)?;
                 self.initial_amp_factor
                     .checked_sub(amp_delta)
                     .map(|x| x as u128)
@@ -170,7 +175,7 @@ impl StableSwap {
                     .checked_mul(d)?
                     .checked_div(amount_b_times_coins.into())?;
                 d_prev = d;
-                d = self.compute_next_d(amp_factor as u64, d, d_prod, sum_x)?;
+                d = self.compute_next_d(amp_factor, d, d_prod, sum_x)?;
                 // Equality with the precision of 1
                 if d > d_prev {
                     if d.checked_sub(d_prev)? <= 1.into() {
diff --git a/ref-exchange/src/stable_swap/mod.rs b/ref-exchange/src/stable_swap/mod.rs
index c75b5b9..5414f69 100644
--- a/ref-exchange/src/stable_swap/mod.rs
+++ b/ref-exchange/src/stable_swap/mod.rs
@@ -5,7 +5,9 @@ use near_sdk::{env, AccountId, Balance, Timestamp};
 
 use crate::errors::{ERR13_LP_NOT_REGISTERED, ERR14_LP_ALREADY_REGISTERED};
 use crate::fees::SwapFees;
-use crate::stable_swap::math::{Fees, StableSwap, SwapResult, MAX_AMP, MIN_AMP, N_COINS};
+use crate::stable_swap::math::{
+    Fees, StableSwap, SwapResult, MAX_AMP, MAX_AMP_CHANGE, MIN_AMP, MIN_RAMP_DURATION, N_COINS,
+};
 use crate::utils::{add_to_collection, SwapVolume, FEE_DIVISOR};
 use crate::StorageKey;
 
@@ -26,20 +28,20 @@ pub struct StableSwapPool {
     /// Total number of shares.
     pub shares_total_supply: Balance,
     /// Initial amplification coefficient.
-    pub init_amp_factor: u64,
-    /// Future amplification coefficient.
-    pub future_amp_factor: u64,
+    pub init_amp_factor: u128,
+    /// Target for ramping up amplification coefficient.
+    pub target_amp_factor: u128,
     /// Initial amplification time.
     pub init_amp_time: Timestamp,
-    /// Future amplification time.
-    pub future_amp_time: Timestamp,
+    /// Stop ramp up amplification time.
+    pub stop_amp_time: Timestamp,
 }
 
 impl StableSwapPool {
     pub fn new(
         id: u32,
         token_account_ids: Vec<ValidAccountId>,
-        amp_factor: u64,
+        amp_factor: u128,
         total_fee: u32,
     ) -> Self {
         assert!(
@@ -60,9 +62,9 @@ impl StableSwapPool {
             shares: LookupMap::new(StorageKey::Shares { pool_id: id }),
             shares_total_supply: 0,
             init_amp_factor: amp_factor,
-            future_amp_factor: amp_factor,
+            target_amp_factor: amp_factor,
             init_amp_time: 0,
-            future_amp_time: 0,
+            stop_amp_time: 0,
         }
     }
 
@@ -99,10 +101,10 @@ impl StableSwapPool {
         );
         let invariant = StableSwap::new(
             self.init_amp_factor,
-            self.future_amp_factor,
+            self.target_amp_factor,
             env::block_timestamp(),
             self.init_amp_time,
-            self.future_amp_time,
+            self.stop_amp_time,
         );
         let new_shares = if self.shares_total_supply == 0 {
             // Bootstrapping the pool.
@@ -151,10 +153,10 @@ impl StableSwapPool {
         let mut result = vec![0u128; N_COINS as usize];
         let invariant = StableSwap::new(
             self.init_amp_factor,
-            self.future_amp_factor,
+            self.target_amp_factor,
             env::block_timestamp(),
             self.init_amp_time,
-            self.future_amp_time,
+            self.stop_amp_time,
         );
         let mut fee_amounts = vec![0u128; N_COINS as usize];
         let stable_swap_fees = Fees::new(self.total_fee, &fees);
@@ -180,9 +182,7 @@ impl StableSwapPool {
                 .checked_sub(result[i as usize])
                 .expect("ERR_CALC");
         }
-        // Never unregister an LP when liquidity is removed.
-        self.shares
-            .insert(&sender_id, &(prev_shares_amount - shares));
+        self.burn_shares(&sender_id, prev_shares_amount, shares);
         env::log(
             format!(
                 "{} shares of liquidity removed: receive back {:?}",
@@ -195,7 +195,6 @@ impl StableSwapPool {
             )
             .as_bytes(),
         );
-        self.shares_total_supply -= shares;
         result
     }
     /// Returns number of tokens in outcome, given amount.
@@ -209,10 +208,10 @@ impl StableSwapPool {
     ) -> SwapResult {
         let invariant = StableSwap::new(
             self.init_amp_factor,
-            self.future_amp_factor,
+            self.target_amp_factor,
             env::block_timestamp(),
             self.init_amp_time,
-            self.future_amp_time,
+            self.stop_amp_time,
         );
         invariant
             .swap_to(
@@ -269,6 +268,9 @@ impl StableSwapPool {
 
         // TODO: add admin / referral fee here.
 
+        // mint
+        println!("{:?}", self.amounts);
+
         result.amount_swapped
     }
 
@@ -281,6 +283,22 @@ impl StableSwapPool {
         add_to_collection(&mut self.shares, &account_id, shares);
     }
 
+    /// Burn shares from given user's balance.
+    fn burn_shares(
+        &mut self,
+        account_id: &AccountId,
+        prev_shares_amount: Balance,
+        shares: Balance,
+    ) {
+        if shares == 0 {
+            return;
+        }
+        // Never remove shares from storage to allow to bring it back without extra storage deposit.
+        self.shares_total_supply -= shares;
+        self.shares
+            .insert(&account_id, &(prev_shares_amount - shares));
+    }
+
     /// Register given account with 0 balance in shares.
     /// Storage payment should be checked by caller.
     pub fn share_register(&mut self, account_id: &AccountId) {
@@ -321,15 +339,55 @@ impl StableSwapPool {
     }
 
     /// [Admin function] increase the amplification factor.
-    pub fn ramp_amplification(&mut self, future_amp_factor: u64, future_amp_time: Timestamp) {
-        // TODO: proper implementation
-        self.future_amp_factor = future_amp_time;
-        self.future_amp_factor = future_amp_factor;
+    pub fn ramp_amplification(&mut self, future_amp_factor: u128, future_amp_time: Timestamp) {
+        let current_time = env::block_timestamp();
+        assert!(
+            current_time >= self.init_amp_time + MIN_RAMP_DURATION,
+            "ERR_RAMP_LOCKED"
+        );
+        assert!(
+            future_amp_time >= current_time + MIN_RAMP_DURATION,
+            "ERR_INSUFFICIENT_RAMP_TIME"
+        );
+        let invariant = StableSwap::new(
+            self.init_amp_factor,
+            self.target_amp_factor,
+            current_time,
+            self.init_amp_time,
+            self.stop_amp_time,
+        );
+        let amp_factor = invariant.compute_amp_factor().expect("ERR_CALC");
+        assert!(
+            future_amp_factor > 0 && future_amp_factor < MAX_AMP,
+            "ERR_INVALID_AMP_FACTOR"
+        );
+        assert!(
+            (future_amp_factor >= amp_factor && future_amp_factor <= amp_factor * MAX_AMP_CHANGE)
+                || (future_amp_factor < amp_factor
+                    && future_amp_factor * MAX_AMP_CHANGE >= amp_factor),
+            "ERR_AMP_LARGE_CHANGE"
+        );
+        self.init_amp_factor = amp_factor;
+        self.init_amp_time = current_time;
+        self.target_amp_factor = future_amp_factor;
+        self.stop_amp_time = future_amp_time;
     }
 
     /// [Admin function] Stop increase of amplification factor.
     pub fn stop_ramp_amplification(&mut self) {
-        // TODO: implement
+        let current_time = env::block_timestamp();
+        let invariant = StableSwap::new(
+            self.init_amp_factor,
+            self.target_amp_factor,
+            current_time,
+            self.init_amp_time,
+            self.stop_amp_time,
+        );
+        let amp_factor = invariant.compute_amp_factor().expect("ERR_CALC");
+        self.init_amp_factor = amp_factor;
+        self.target_amp_factor = amp_factor;
+        self.init_amp_time = current_time;
+        self.stop_amp_time = current_time;
     }
 }
 
@@ -387,6 +445,7 @@ mod tests {
         assert_eq!(amounts_out, vec![0, to_yocto("5")]);
     }
 
+    /// Test everything with fees.
     #[test]
     fn test_with_fees() {
         let mut context = VMContextBuilder::new();
@@ -395,9 +454,17 @@ mod tests {
         let mut amounts = vec![to_yocto("5"), to_yocto("10")];
         let fees = SwapFees::new(1000);
         let num_shares = pool.add_liquidity(accounts(0).as_ref(), &mut amounts, &fees);
-        let amounts_out =
-            pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1], &fees);
-        println!("amount out: {:?}", amounts_out);
+        let amount_out = pool.swap(
+            accounts(1).as_ref(),
+            to_yocto("1"),
+            accounts(2).as_ref(),
+            1,
+            &fees,
+        );
+        println!("swap out: {}", amount_out);
+        // let amounts_out =
+        //     pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1], &fees);
+        // println!("amount out: {:?}", amounts_out);
     }
 
     /// Test that adding and then removing all of the liquidity leaves the pool empty and with no shares.
@@ -436,4 +503,24 @@ mod tests {
         assert_eq!(pool.share_balance_of(accounts(3).as_ref()), 0);
         assert_eq!(pool.amounts, vec![1, 1]);
     }
+
+    /// Test ramping up amplification factor, ramping it even more and then stopping.
+    #[test]
+    fn test_ramp_amp() {
+        let mut context = VMContextBuilder::new();
+        testing_env!(context.predecessor_account_id(accounts(0)).build());
+        let mut pool = StableSwapPool::new(0, vec![accounts(1), accounts(2)], 1, 0);
+
+        let start_ts = 1_000_000_000;
+        testing_env!(context.block_timestamp(start_ts).build());
+        pool.ramp_amplification(5, start_ts + MIN_RAMP_DURATION * 10);
+        testing_env!(context
+            .block_timestamp(start_ts + MIN_RAMP_DURATION * 3)
+            .build());
+        pool.ramp_amplification(15, start_ts + MIN_RAMP_DURATION * 20);
+        testing_env!(context
+            .block_timestamp(start_ts + MIN_RAMP_DURATION * 5)
+            .build());
+        pool.stop_ramp_amplification();
+    }
 }

From ed57bf470e5c8d9527c1381320fe8b6dfe7bd153 Mon Sep 17 00:00:00 2001
From: referencedev <goi65io3903d@protonmail.com>
Date: Fri, 17 Sep 2021 07:57:45 -0700
Subject: [PATCH 6/6] Integrating stable swap into main exchange

---
 ref-exchange/src/lib.rs             | 17 +++++++++++++++++
 ref-exchange/src/owner.rs           | 29 +++++++++++++++++++++++++++++
 ref-exchange/src/stable_swap/mod.rs |  6 +++---
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/ref-exchange/src/lib.rs b/ref-exchange/src/lib.rs
index 428a4bd..a039d4f 100644
--- a/ref-exchange/src/lib.rs
+++ b/ref-exchange/src/lib.rs
@@ -18,6 +18,7 @@ use crate::errors::*;
 use crate::fees::SwapFees;
 use crate::pool::Pool;
 use crate::simple_pool::SimplePool;
+use crate::stable_swap::StableSwapPool;
 use crate::utils::check_token_duplicates;
 pub use crate::views::PoolInfo;
 
@@ -91,6 +92,22 @@ impl Contract {
         )))
     }
 
+    #[payable]
+    pub fn add_stable_swap_pool(
+        &mut self,
+        tokens: Vec<ValidAccountId>,
+        fee: u32,
+        amp_factor: u64,
+    ) -> u64 {
+        check_token_duplicates(&tokens);
+        self.internal_add_pool(Pool::StableSwapPool(StableSwapPool::new(
+            self.pools.len() as u32,
+            tokens,
+            amp_factor as u128,
+            fee + self.exchange_fee + self.referral_fee,
+        )))
+    }
+
     /// [AUDIT_03_reject(NOPE action is allowed by design)]
     /// [AUDIT_04]
     /// Executes generic set of actions.
diff --git a/ref-exchange/src/owner.rs b/ref-exchange/src/owner.rs
index 8b3c4ff..e5c5a55 100644
--- a/ref-exchange/src/owner.rs
+++ b/ref-exchange/src/owner.rs
@@ -1,5 +1,7 @@
 //! Implement all the relevant logic for owner of this contract.
 
+use near_sdk::json_types::WrappedTimestamp;
+
 use crate::*;
 
 #[near_bindgen]
@@ -49,6 +51,33 @@ impl Contract {
             "ERR_NOT_ALLOWED"
         );
     }
+
+    pub fn stable_swap_ramp_amp(
+        &mut self,
+        pool_id: u64,
+        future_amp_factor: u64,
+        future_amp_time: WrappedTimestamp,
+    ) {
+        self.assert_owner();
+        let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
+        match &mut pool {
+            Pool::StableSwapPool(pool) => {
+                pool.ramp_amplification(future_amp_factor as u128, future_amp_time.0)
+            }
+            _ => env::panic(b"ERR_NOT_STABLE_POOL"),
+        }
+        self.pools.replace(pool_id, &pool);
+    }
+
+    pub fn stable_swap_stop_ramp_amp(&mut self, pool_id: u64) {
+        let mut pool = self.pools.get(pool_id).expect("ERR_NO_POOL");
+        match &mut pool {
+            Pool::StableSwapPool(pool) => pool.stop_ramp_amplification(),
+            _ => env::panic(b"ERR_NOT_STABLE_POOL"),
+        }
+        self.assert_owner();
+        self.pools.replace(pool_id, &pool);
+    }
 }
 
 #[cfg(target_arch = "wasm32")]
diff --git a/ref-exchange/src/stable_swap/mod.rs b/ref-exchange/src/stable_swap/mod.rs
index 5414f69..f301bba 100644
--- a/ref-exchange/src/stable_swap/mod.rs
+++ b/ref-exchange/src/stable_swap/mod.rs
@@ -462,9 +462,9 @@ mod tests {
             &fees,
         );
         println!("swap out: {}", amount_out);
-        // let amounts_out =
-        //     pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1], &fees);
-        // println!("amount out: {:?}", amounts_out);
+        let amounts_out =
+            pool.remove_liquidity(accounts(0).as_ref(), num_shares, vec![1, 1], &fees);
+        println!("amount out: {:?}", amounts_out);
     }
 
     /// Test that adding and then removing all of the liquidity leaves the pool empty and with no shares.