token-swap: Add offset curve and math (#935)

* Add offset curve

* Fix for math

* Add PreciseNumber

* Use Balancer formulation for trading token -> pool token conversion

* Add round-trip conversion testing

* Add offset curve to JS

* Run cargo fmt

* Update JS test numbers for new calcs

* Integrate review feedback

* Allow for withdrawals when one side is 0

* Run cargo fmt

* Disallow deposits for offset curve

* Run cargo fmt

* Allow for withdrawals through 0

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

@@ -63,9 +63,9 @@ let currentFeeAmount = 0;
 // Swap instruction constants
 // Because there is no withdraw fee in the production version, these numbers
 // need to get slightly tweaked in the two cases.
-const SWAP_AMOUNT_IN = 100000;
+const SWAP_AMOUNT_IN = 99999;
 const SWAP_AMOUNT_OUT = SWAP_PROGRAM_OWNER_FEE_ADDRESS ? 90661 : 90674;
-const SWAP_FEE = SWAP_PROGRAM_OWNER_FEE_ADDRESS ? 22272 : 22276;
+const SWAP_FEE = SWAP_PROGRAM_OWNER_FEE_ADDRESS ? 21820 : 21823;
 const HOST_SWAP_FEE = SWAP_PROGRAM_OWNER_FEE_ADDRESS
   ? Math.floor((SWAP_FEE * HOST_FEE_NUMERATOR) / HOST_FEE_DENOMINATOR)
   : 0;
@@ -291,9 +291,9 @@ export async function deposit(): Promise<void> {
   const poolMintInfo = await tokenPool.getMintInfo();
   const supply = poolMintInfo.supply.toNumber();
   const swapTokenA = await mintA.getAccountInfo(tokenAccountA);
-  const tokenA = (swapTokenA.amount.toNumber() * POOL_TOKEN_AMOUNT) / supply;
+  const tokenA = Math.floor((swapTokenA.amount.toNumber() * POOL_TOKEN_AMOUNT) / (supply + POOL_TOKEN_AMOUNT));
   const swapTokenB = await mintB.getAccountInfo(tokenAccountB);
-  const tokenB = (swapTokenB.amount.toNumber() * POOL_TOKEN_AMOUNT) / supply;
+  const tokenB = Math.floor((swapTokenB.amount.toNumber() * POOL_TOKEN_AMOUNT) / (supply + POOL_TOKEN_AMOUNT));
 
   console.log('Creating depositor token a account');
   const userAccountA = await mintA.createAccount(owner.publicKey);
@@ -496,7 +496,7 @@ export async function swap(): Promise<void> {
 
   info = await mintA.getAccountInfo(tokenAccountA);
   assert(info.amount.toNumber() == currentSwapTokenA + SWAP_AMOUNT_IN);
-  currentSwapTokenA -= SWAP_AMOUNT_IN;
+  currentSwapTokenA += SWAP_AMOUNT_IN;
 
   info = await mintB.getAccountInfo(tokenAccountB);
   assert(info.amount.toNumber() == currentSwapTokenB - SWAP_AMOUNT_OUT);

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

@@ -83,6 +83,7 @@ export const TokenSwapLayout: typeof BufferLayout.Structure = BufferLayout.struc
 export const CurveType = Object.freeze({
   ConstantProduct: 0, // Constant product curve, Uniswap-style
   ConstantPrice: 1, // Constant price curve, always X amount of A token for 1 B token, where X is defined at init
+  Offset: 3, // Offset curve, like Uniswap, but with an additional offset on the token B side
 });
 
 /**

token-swap/program/src/curve/base.rs

@@ -10,6 +10,7 @@ use crate::curve::{
     constant_price::ConstantPriceCurve,
     constant_product::ConstantProductCurve,
     fees::Fees,
+    offset::OffsetCurve,
     stable::StableCurve,
 };
 use arrayref::{array_mut_ref, array_ref, array_refs, mut_array_refs};
@@ -24,8 +25,10 @@ pub enum CurveType {
     ConstantProduct,
     /// Flat line, always providing 1:1 from one token to another
     ConstantPrice,
-    /// Stable, Like uniswap, but with wide zone of 1:1 instead of one point
+    /// Stable, like uniswap, but with wide zone of 1:1 instead of one point
     Stable,
+    /// Offset curve, like Uniswap, but the token B side has a faked offset
+    Offset,
 }
 
 /// Encodes all results of swapping from a source token to a destination token
@@ -181,6 +184,7 @@ impl Pack for SwapCurve {
                     Box::new(ConstantPriceCurve::unpack_from_slice(calculator)?)
                 }
                 CurveType::Stable => Box::new(StableCurve::unpack_from_slice(calculator)?),
+                CurveType::Offset => Box::new(OffsetCurve::unpack_from_slice(calculator)?),
             },
         })
     }
@@ -210,6 +214,7 @@ impl TryFrom<u8> for CurveType {
             0 => Ok(CurveType::ConstantProduct),
             1 => Ok(CurveType::ConstantPrice),
             2 => Ok(CurveType::Stable),
+            3 => Ok(CurveType::Offset),
             _ => Err(ProgramError::InvalidAccountData),
         }
     }

token-swap/program/src/curve/calculator.rs

@@ -1,6 +1,6 @@
 //! Swap calculations
 
-use crate::error::SwapError;
+use crate::{curve::math::PreciseNumber, error::SwapError};
 use std::fmt::Debug;
 
 /// Initial amount of pool tokens for swap contract, hard-coded to something
@@ -11,7 +11,7 @@ pub const INITIAL_SWAP_POOL_AMOUNT: u128 = 1_000_000_000;
 
 /// Hardcode the number of token types in a pool, used to calculate the
 /// equivalent pool tokens for the owner trading fee.
-const TOKENS_IN_POOL: u128 = 2;
+pub const TOKENS_IN_POOL: u128 = 2;
 
 /// Helper function for mapping to SwapError::CalculationFailure
 pub fn map_zero_to_none(x: u128) -> Option<u128> {
@@ -115,10 +115,18 @@ pub trait CurveCalculator: Debug + DynPack {
             TradeDirection::AtoB => swap_token_a_amount,
             TradeDirection::BtoA => swap_token_b_amount,
         };
-        pool_supply
+        let swap_source_amount = PreciseNumber::new(swap_source_amount)?;
-            .checked_mul(source_amount)?
+        let source_amount = PreciseNumber::new(source_amount)?;
-            .checked_div(swap_source_amount)?
+        let ratio = source_amount.checked_div(&swap_source_amount)?;
-            .checked_div(TOKENS_IN_POOL)
+        let one = PreciseNumber::new(1)?;
+        let two = PreciseNumber::new(2)?;
+        let base = one.checked_add(&ratio)?;
+        let guess = base.checked_div(&two)?;
+        let root = base
+            .newtonian_root_approximation(&two, guess)?
+            .checked_sub(&one)?;
+        let pool_supply = PreciseNumber::new(pool_supply)?;
+        pool_supply.checked_mul(&root)?.to_imprecise()
     }
 
     /// Validate that the given curve has no bad parameters
@@ -135,4 +143,90 @@ pub trait CurveCalculator: Debug + DynPack {
         }
         Ok(())
     }
+
+    /// Some curves will function best and prevent attacks if we prevent
+    /// deposits after initialization
+    fn allows_deposits(&self) -> bool {
+        true
+    }
+}
+
+#[cfg(test)]
+pub mod test {
+    use super::*;
+
+    /// Check that two numbers are within 1 of each other
+    fn almost_equal(a: u128, b: u128) {
+        if a >= b {
+            assert!(a - b <= 1);
+        } else {
+            assert!(b - a <= 1);
+        }
+    }
+
+    pub fn check_pool_token_conversion(
+        curve: &dyn CurveCalculator,
+        swap_token_a_amount: u128,
+        swap_token_b_amount: u128,
+        token_a_amount: u128,
+    ) {
+        // check that depositing token A is the same as swapping for token B
+        // and depositing the result
+        let swap_results = curve
+            .swap_without_fees(
+                token_a_amount,
+                swap_token_a_amount,
+                swap_token_b_amount,
+                TradeDirection::AtoB,
+            )
+            .unwrap();
+        let token_a_amount = swap_results.source_amount_swapped;
+        let token_b_amount = swap_results.destination_amount_swapped;
+        let pool_supply = curve.new_pool_supply();
+        let pool_tokens_from_a = curve
+            .trading_tokens_to_pool_tokens(
+                token_a_amount,
+                swap_token_a_amount + token_a_amount,
+                swap_token_b_amount,
+                pool_supply,
+                TradeDirection::AtoB,
+            )
+            .unwrap();
+        let pool_tokens_from_b = curve
+            .trading_tokens_to_pool_tokens(
+                token_b_amount,
+                swap_token_a_amount + token_a_amount,
+                swap_token_b_amount,
+                pool_supply,
+                TradeDirection::BtoA,
+            )
+            .unwrap();
+        let deposit_token_a = curve
+            .pool_tokens_to_trading_tokens(
+                pool_tokens_from_a,
+                pool_supply + pool_tokens_from_a,
+                swap_token_a_amount,
+                swap_token_b_amount,
+            )
+            .unwrap();
+
+        let deposit_token_b = curve
+            .pool_tokens_to_trading_tokens(
+                pool_tokens_from_b,
+                pool_supply + pool_tokens_from_b,
+                swap_token_a_amount,
+                swap_token_b_amount,
+            )
+            .unwrap();
+
+        // They should be within 1 token because truncation
+        almost_equal(
+            deposit_token_b.token_a_amount,
+            deposit_token_a.token_a_amount,
+        );
+        almost_equal(
+            deposit_token_b.token_b_amount,
+            deposit_token_b.token_b_amount,
+        );
+    }
 }

token-swap/program/src/curve/constant_product.rs

@@ -93,7 +93,7 @@ impl DynPack for ConstantProductCurve {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::curve::calculator::INITIAL_SWAP_POOL_AMOUNT;
+    use crate::curve::calculator::{test::check_pool_token_conversion, INITIAL_SWAP_POOL_AMOUNT};
 
     #[test]
     fn initial_pool_amount() {
@@ -213,4 +213,24 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn pool_token_conversion() {
+        let tests: &[(u128, u128, u128)] = &[
+            (1_000_000, 2400112, 100_000),
+            (1_000, 100, 100),
+            (30, 1_288, 100_000),
+            (1_000, 1_288, 100_000),
+            (212, 10_000, 100_000),
+        ];
+        for (swap_token_a_amount, swap_token_b_amount, token_a_amount) in tests.iter() {
+            let curve = ConstantProductCurve {};
+            check_pool_token_conversion(
+                &curve,
+                *swap_token_a_amount,
+                *swap_token_b_amount,
+                *token_a_amount,
+            );
+        }
+    }
 }

token-swap/program/src/curve/math.rs

@@ -39,3 +39,407 @@ impl U256 {
         }
     }
 }
+
+/// The representation of the number one as a precise number
+pub const ONE: u128 = 10_000_000_000;
+
+/// Maximum weight for token in swap.  This number is meant to stay small to
+/// so that it is possible to accurately calculate x^(MAX_WEIGHT / MIN_WEIGHT).
+pub const MAX_WEIGHT: u8 = 100;
+
+/// Minimum weight for token in swap
+pub const MIN_WEIGHT: u8 = 1;
+
+/// Struct encapsulating a fixed-point number that allows for decimal calculations
+#[derive(Clone)]
+pub struct PreciseNumber {
+    /// Wrapper over the inner value, which is multiplied by ONE
+    pub value: U256,
+}
+
+/// The precise-number 1 as a U256
+fn one() -> U256 {
+    U256::from(ONE)
+}
+
+/// The number 0 as a PreciseNumber, used for easier calculations.
+fn zero() -> U256 {
+    U256::from(0)
+}
+
+impl PreciseNumber {
+    /// Correction to apply to avoid truncation errors on division.  Since
+    /// integer operations will always floor the result, we artifically bump it
+    /// up by one half to get the expect result.
+    fn rounding_correction() -> U256 {
+        U256::from(ONE / 2)
+    }
+
+    /// Desired precision for the correction factor applied during each
+    /// iteration of checked_pow_approximation.  Once the correction factor is
+    /// smaller than this number, or we reach the maxmium number of iterations,
+    /// the calculation ends.
+    fn precision() -> U256 {
+        U256::from(100)
+    }
+
+    /// Maximum number iterations to apply on checked_pow_approximation.
+    const MAX_APPROXIMATION_ITERATIONS: u128 = 100;
+
+    /// Minimum base allowed when calculating exponents in checked_pow_fraction
+    /// and checked_pow_approximation.  This simply avoids 0 as a base.
+    fn min_pow_base() -> U256 {
+        U256::from(1)
+    }
+
+    /// Maximum base allowed when calculating exponents in checked_pow_fraction
+    /// and checked_pow_approximation.  The calculation use a Taylor Series
+    /// approxmation around 1, which converges for bases between 0 and 2.  See
+    /// https://en.wikipedia.org/wiki/Binomial_series#Conditions_for_convergence
+    /// for more information.
+    fn max_pow_base() -> U256 {
+        U256::from(2 * ONE)
+    }
+
+    /// Create a precise number from an imprecise u128, should always succeed
+    pub fn new(value: u128) -> Option<Self> {
+        let value = U256::from(value).checked_mul(one())?;
+        Some(Self { value })
+    }
+
+    /// Convert a precise number back to u128
+    pub fn to_imprecise(&self) -> Option<u128> {
+        match self
+            .value
+            .checked_add(Self::rounding_correction())?
+            .checked_div(one())
+        {
+            Some(v) => Some(v.as_u128()),
+            None => None,
+        }
+    }
+
+    /// Checks that two PreciseNumbers are equal within some tolerance
+    pub fn almost_eq(&self, rhs: &Self, precision: U256) -> bool {
+        let (difference, _) = self.unsigned_sub(rhs);
+        difference.value < precision
+    }
+
+    /// Floors a precise value to a precision of ONE
+    pub fn floor(&self) -> Option<Self> {
+        let value = self.value.checked_div(one())?.checked_mul(one())?;
+        Some(Self { value })
+    }
+
+    /// Performs a checked division on two precise numbers
+    pub fn checked_div(&self, rhs: &Self) -> Option<Self> {
+        if rhs.value == zero() {
+            return None;
+        }
+        match self.value.checked_mul(one()) {
+            Some(v) => {
+                let value = v
+                    .checked_add(Self::rounding_correction())?
+                    .checked_div(rhs.value)?;
+                Some(Self { value })
+            }
+            None => {
+                let value = self
+                    .value
+                    .checked_add(Self::rounding_correction())?
+                    .checked_div(rhs.value)?
+                    .checked_mul(one())?;
+                Some(Self { value })
+            }
+        }
+    }
+
+    /// Performs a multiplication on two precise numbers
+    pub fn checked_mul(&self, rhs: &Self) -> Option<Self> {
+        match self.value.checked_mul(rhs.value) {
+            Some(v) => {
+                let value = v
+                    .checked_add(Self::rounding_correction())?
+                    .checked_div(one())?;
+                Some(Self { value })
+            }
+            None => {
+                let value = if self.value >= rhs.value {
+                    self.value.checked_div(one())?.checked_mul(rhs.value)?
+                } else {
+                    rhs.value.checked_div(one())?.checked_mul(self.value)?
+                };
+                Some(Self { value })
+            }
+        }
+    }
+
+    /// Performs addition of two precise numbers
+    pub fn checked_add(&self, rhs: &Self) -> Option<Self> {
+        let value = self.value.checked_add(rhs.value)?;
+        Some(Self { value })
+    }
+
+    /// Subtracts the argument from self
+    pub fn checked_sub(&self, rhs: &Self) -> Option<Self> {
+        let value = self.value.checked_sub(rhs.value)?;
+        Some(Self { value })
+    }
+
+    /// Performs a subtraction, returning the result and whether the result is negative
+    pub fn unsigned_sub(&self, rhs: &Self) -> (Self, bool) {
+        match self.value.checked_sub(rhs.value) {
+            None => {
+                let value = rhs.value.checked_sub(self.value).unwrap();
+                (Self { value }, true)
+            }
+            Some(value) => (Self { value }, false),
+        }
+    }
+
+    /// Performs pow on a precise number
+    pub fn checked_pow(&self, exponent: u128) -> Option<Self> {
+        // For odd powers, start with a multiplication by base since we halve the
+        // exponent at the start
+        let value = if exponent.checked_rem(2)? == 0 {
+            one()
+        } else {
+            self.value
+        };
+        let mut result = Self { value };
+
+        // To minimize the number of operations, we keep squaring the base, and
+        // only push to the result on odd exponents, like a binary decomposition
+        // of the exponent.
+        let mut squared_base = self.clone();
+        let mut current_exponent = exponent.checked_div(2)?;
+        while current_exponent != 0 {
+            squared_base = squared_base.checked_mul(&squared_base)?;
+
+            // For odd exponents, "push" the base onto the value
+            if current_exponent.checked_rem(2)? != 0 {
+                result = result.checked_mul(&squared_base)?;
+            }
+
+            current_exponent = current_exponent.checked_div(2)?;
+        }
+        Some(result)
+    }
+
+    /// Approximate the nth root of a number using a Taylor Series around 1 on
+    /// x ^ n, where 0 < n < 1, result is a precise number.
+    /// Refine the guess for each term, using:
+    ///                                  1                    2
+    /// f(x) = f(a) + f'(a) * (x - a) + --- * f''(a) * (x - a)  + ...
+    ///                                  2!
+    /// For x ^ n, this gives:
+    ///  n    n         n-1           1                  n-2        2
+    /// x  = a  + n * a    (x - a) + --- * n * (n - 1) a     (x - a)  + ...
+    ///                               2!
+    ///
+    /// More simply, this means refining the term at each iteration with:
+    ///
+    /// t_k+1 = t_k * (x - a) * (n + 1 - k) / k
+    ///
+    /// where a = 1, n = power, x = precise_num
+    pub fn checked_pow_approximation(&self, exponent: &Self, max_iterations: u128) -> Option<Self> {
+        assert!(self.value >= Self::min_pow_base());
+        assert!(self.value <= Self::max_pow_base());
+        let one = Self::new(1)?;
+        if exponent.value == zero() {
+            return Some(one);
+        }
+        let mut precise_guess = one.clone();
+        let mut term = precise_guess.clone();
+        let (x_minus_a, x_minus_a_negative) = self.unsigned_sub(&precise_guess);
+        let exponent_plus_one = exponent.checked_add(&one)?;
+        let mut negative = false;
+        for k in 1..max_iterations {
+            let k = Self::new(k)?;
+            let (current_exponent, current_exponent_negative) = exponent_plus_one.unsigned_sub(&k);
+            term = term.checked_mul(&current_exponent)?;
+            term = term.checked_mul(&x_minus_a)?;
+            term = term.checked_div(&k)?;
+            if term.value < Self::precision() {
+                break;
+            }
+            if x_minus_a_negative {
+                negative = !negative;
+            }
+            if current_exponent_negative {
+                negative = !negative;
+            }
+            if negative {
+                precise_guess = precise_guess.checked_sub(&term)?;
+            } else {
+                precise_guess = precise_guess.checked_add(&term)?;
+            }
+        }
+        Some(precise_guess)
+    }
+
+    /// Get the power of a number, where the exponent is expressed as a fraction
+    /// (numerator / denominator)
+    pub fn checked_pow_fraction(&self, exponent: &Self) -> Option<Self> {
+        assert!(self.value >= Self::min_pow_base());
+        assert!(self.value <= Self::max_pow_base());
+        let whole_exponent = exponent.floor()?;
+        let precise_whole = self.checked_pow(whole_exponent.to_imprecise()?)?;
+        let (remainder_exponent, negative) = exponent.unsigned_sub(&whole_exponent);
+        assert!(!negative);
+        if remainder_exponent.value == U256::from(0) {
+            return Some(precise_whole);
+        }
+        let precise_remainder = self
+            .checked_pow_approximation(&remainder_exponent, Self::MAX_APPROXIMATION_ITERATIONS)?;
+        precise_whole.checked_mul(&precise_remainder)
+    }
+
+    /// Approximate the nth root of a number using Newton's method
+    /// https://en.wikipedia.org/wiki/Newton%27s_method
+    pub fn newtonian_root_approximation(&self, root: &Self, mut guess: Self) -> Option<Self> {
+        if root.value == zero() {
+            return None;
+        }
+        let one = Self::new(1)?;
+        let root_minus_one = root.checked_sub(&one)?;
+        let root_minus_one_whole = root_minus_one.to_imprecise()?;
+        let mut last_guess = guess.clone();
+        let precision = Self::precision();
+        for _ in 0..Self::MAX_APPROXIMATION_ITERATIONS {
+            // x_k+1 = ((n - 1) * x_k + A / (x_k ^ (n - 1))) / n
+            let first_term = root_minus_one.checked_mul(&guess)?;
+            let power = guess.checked_pow(root_minus_one_whole);
+            let second_term = match power {
+                Some(num) => self.checked_div(&num)?,
+                None => Self::new(0)?,
+            };
+            guess = first_term.checked_add(&second_term)?.checked_div(&root)?;
+            if last_guess.almost_eq(&guess, precision) {
+                break;
+            } else {
+                last_guess = guess.clone();
+            }
+        }
+        Some(guess)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn check_pow_approximation(base: U256, exponent: U256, expected: U256) {
+        let precision = U256::from(5_000_000); // correct to at least 3 decimal places
+        let base = PreciseNumber { value: base };
+        let exponent = PreciseNumber { value: exponent };
+        let root = base
+            .checked_pow_approximation(&exponent, PreciseNumber::MAX_APPROXIMATION_ITERATIONS)
+            .unwrap();
+        let expected = PreciseNumber { value: expected };
+        assert!(root.almost_eq(&expected, precision));
+    }
+
+    #[test]
+    fn test_root_approximation() {
+        let one = one();
+        // square root
+        check_pow_approximation(one / 4, one / 2, one / 2); // 1/2
+        check_pow_approximation(one * 11 / 10, one / 2, U256::from(1_0488088481u128)); // 1.0488088481
+
+        // 5th root
+        check_pow_approximation(one * 4 / 5, one * 2 / 5, U256::from(9146101038u128)); // 0.9146101038
+
+        // 10th root
+        check_pow_approximation(one / 2, one * 4 / 50, U256::from(9460576467u128));
+        // 0.9460576467
+    }
+
+    fn check_pow_fraction(base: U256, exponent: U256, expected: U256, precision: U256) {
+        let base = PreciseNumber { value: base };
+        let exponent = PreciseNumber { value: exponent };
+        let power = base.checked_pow_fraction(&exponent).unwrap();
+        let expected = PreciseNumber { value: expected };
+        assert!(power.almost_eq(&expected, precision));
+    }
+
+    #[test]
+    fn test_pow_fraction() {
+        let one = one();
+        let precision = U256::from(5_000_000); // correct to at least 3 decimal places
+        let less_precision = precision * 100; // correct to at least 1 decimal place
+        check_pow_fraction(one, one, one, precision);
+        check_pow_fraction(
+            one * 20 / 13,
+            one * 50 / 3,
+            U256::from(1312_5344847391u128),
+            precision,
+        ); // 1312.5344847391
+        check_pow_fraction(one * 2 / 7, one * 49 / 4, U256::from(2163), precision);
+        check_pow_fraction(
+            one * 5000 / 5100,
+            one / 9,
+            U256::from(9978021269u128),
+            precision,
+        ); // 0.99780212695
+           // results get less accurate as the base gets further from 1, so allow
+           // for a greater margin of error
+        check_pow_fraction(
+            one * 2,
+            one * 27 / 5,
+            U256::from(42_2242531447u128),
+            less_precision,
+        ); // 42.2242531447
+        check_pow_fraction(
+            one * 18 / 10,
+            one * 11 / 3,
+            U256::from(8_6297692905u128),
+            less_precision,
+        ); // 8.629769290
+    }
+
+    #[test]
+    fn test_newtonian_approximation() {
+        // square root
+        let test = PreciseNumber::new(9).unwrap();
+        let nth_root = PreciseNumber::new(2).unwrap();
+        let guess = test.checked_div(&nth_root).unwrap();
+        let root = test
+            .newtonian_root_approximation(&nth_root, guess)
+            .unwrap()
+            .to_imprecise()
+            .unwrap();
+        assert_eq!(root, 3); // actually 3
+
+        let test = PreciseNumber::new(101).unwrap();
+        let nth_root = PreciseNumber::new(2).unwrap();
+        let guess = test.checked_div(&nth_root).unwrap();
+        let root = test
+            .newtonian_root_approximation(&nth_root, guess)
+            .unwrap()
+            .to_imprecise()
+            .unwrap();
+        assert_eq!(root, 10); // actually 10.049875
+
+        let test = PreciseNumber::new(1_000_000_000).unwrap();
+        let nth_root = PreciseNumber::new(2).unwrap();
+        let guess = test.checked_div(&nth_root).unwrap();
+        let root = test
+            .newtonian_root_approximation(&nth_root, guess)
+            .unwrap()
+            .to_imprecise()
+            .unwrap();
+        assert_eq!(root, 31_623); // actually 31622.7766
+
+        // 5th root
+        let test = PreciseNumber::new(500).unwrap();
+        let nth_root = PreciseNumber::new(5).unwrap();
+        let guess = test.checked_div(&nth_root).unwrap();
+        let root = test
+            .newtonian_root_approximation(&nth_root, guess)
+            .unwrap()
+            .to_imprecise()
+            .unwrap();
+        assert_eq!(root, 3); // actually 3.46572422
+    }
+}

token-swap/program/src/curve/mod.rs

@@ -6,4 +6,5 @@ pub mod constant_price;
 pub mod constant_product;
 pub mod fees;
 pub mod math;
+pub mod offset;
 pub mod stable;

token-swap/program/src/curve/offset.rs

@@ -0,0 +1,289 @@
+//! The Uniswap invariant calculator with an extra offset
+
+use crate::{
+    curve::{
+        calculator::{
+            CurveCalculator, DynPack, SwapWithoutFeesResult, TradeDirection, TradingTokenResult,
+        },
+        constant_product::swap,
+        math::PreciseNumber,
+    },
+    error::SwapError,
+};
+use arrayref::{array_mut_ref, array_ref};
+use solana_program::{
+    program_error::ProgramError,
+    program_pack::{IsInitialized, Pack, Sealed},
+};
+
+/// Offset curve, uses ConstantProduct under the hood, but adds an offset to
+/// one side on swap calculations
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct OffsetCurve {
+    /// Amount to offset the token B liquidity account
+    pub token_b_offset: u64,
+}
+
+impl CurveCalculator for OffsetCurve {
+    /// Constant product swap ensures token a * (token b + offset) = constant
+    fn swap_without_fees(
+        &self,
+        source_amount: u128,
+        swap_source_amount: u128,
+        swap_destination_amount: u128,
+        trade_direction: TradeDirection,
+    ) -> Option<SwapWithoutFeesResult> {
+        let token_b_offset = self.token_b_offset as u128;
+        let swap_source_amount = match trade_direction {
+            TradeDirection::AtoB => swap_source_amount,
+            TradeDirection::BtoA => swap_source_amount.checked_add(token_b_offset)?,
+        };
+        let swap_destination_amount = match trade_direction {
+            TradeDirection::AtoB => swap_destination_amount.checked_add(token_b_offset)?,
+            TradeDirection::BtoA => swap_destination_amount,
+        };
+        swap(source_amount, swap_source_amount, swap_destination_amount)
+    }
+
+    /// The conversion for the offset curve needs to take into account the
+    /// offset
+    fn pool_tokens_to_trading_tokens(
+        &self,
+        pool_tokens: u128,
+        pool_token_supply: u128,
+        swap_token_a_amount: u128,
+        swap_token_b_amount: u128,
+    ) -> Option<TradingTokenResult> {
+        let token_b_offset = self.token_b_offset as u128;
+        let token_a_amount = pool_tokens
+            .checked_mul(swap_token_a_amount)?
+            .checked_div(pool_token_supply)?;
+        let token_b_amount = pool_tokens
+            .checked_mul(swap_token_b_amount.checked_add(token_b_offset)?)?
+            .checked_div(pool_token_supply)?;
+        Some(TradingTokenResult {
+            token_a_amount,
+            token_b_amount,
+        })
+    }
+
+    /// Get the amount of pool tokens for the given amount of token A and B,
+    /// taking into account the offset
+    fn trading_tokens_to_pool_tokens(
+        &self,
+        source_amount: u128,
+        swap_token_a_amount: u128,
+        swap_token_b_amount: u128,
+        pool_supply: u128,
+        trade_direction: TradeDirection,
+    ) -> Option<u128> {
+        let token_b_offset = self.token_b_offset as u128;
+        let swap_source_amount = match trade_direction {
+            TradeDirection::AtoB => swap_token_a_amount,
+            TradeDirection::BtoA => swap_token_b_amount.checked_add(token_b_offset)?,
+        };
+        let swap_source_amount = PreciseNumber::new(swap_source_amount)?;
+        let source_amount = PreciseNumber::new(source_amount)?;
+        let ratio = source_amount.checked_div(&swap_source_amount)?;
+        let one = PreciseNumber::new(1)?;
+        let two = PreciseNumber::new(2)?;
+        let base = one.checked_add(&ratio)?;
+        let guess = base.checked_div(&two)?;
+        let root = base
+            .newtonian_root_approximation(&two, guess)?
+            .checked_sub(&one)?;
+        let pool_supply = PreciseNumber::new(pool_supply)?;
+        pool_supply.checked_mul(&root)?.to_imprecise()
+    }
+
+    fn validate(&self) -> Result<(), SwapError> {
+        if self.token_b_offset == 0 {
+            Err(SwapError::InvalidCurve)
+        } else {
+            Ok(())
+        }
+    }
+
+    fn validate_supply(&self, token_a_amount: u64, _token_b_amount: u64) -> Result<(), SwapError> {
+        if token_a_amount == 0 {
+            return Err(SwapError::EmptySupply);
+        }
+        Ok(())
+    }
+
+    /// Offset curves can cause arbitrage opportunities if outside users are
+    /// allowed to deposit.  For example, in the offset curve, if there's swap
+    /// with 1 million of token A against an offset of 2 million token B,
+    /// someone else can deposit 1 million A and 2 million B for LP tokens.
+    /// The pool creator can then use their LP tokens to steal the 2 million B,
+    fn allows_deposits(&self) -> bool {
+        false
+    }
+}
+
+/// IsInitialized is required to use `Pack::pack` and `Pack::unpack`
+impl IsInitialized for OffsetCurve {
+    fn is_initialized(&self) -> bool {
+        true
+    }
+}
+impl Sealed for OffsetCurve {}
+impl Pack for OffsetCurve {
+    const LEN: usize = 8;
+    fn pack_into_slice(&self, output: &mut [u8]) {
+        (self as &dyn DynPack).pack_into_slice(output);
+    }
+
+    fn unpack_from_slice(input: &[u8]) -> Result<OffsetCurve, ProgramError> {
+        let token_b_offset = array_ref![input, 0, 8];
+        Ok(Self {
+            token_b_offset: u64::from_le_bytes(*token_b_offset),
+        })
+    }
+}
+
+impl DynPack for OffsetCurve {
+    fn pack_into_slice(&self, output: &mut [u8]) {
+        let token_b_offset = array_mut_ref![output, 0, 8];
+        *token_b_offset = self.token_b_offset.to_le_bytes();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::curve::calculator::test::check_pool_token_conversion;
+
+    #[test]
+    fn pack_curve() {
+        let token_b_offset = u64::MAX;
+        let curve = OffsetCurve { token_b_offset };
+
+        let mut packed = [0u8; OffsetCurve::LEN];
+        Pack::pack_into_slice(&curve, &mut packed[..]);
+        let unpacked = OffsetCurve::unpack(&packed).unwrap();
+        assert_eq!(curve, unpacked);
+
+        let mut packed = vec![];
+        packed.extend_from_slice(&token_b_offset.to_le_bytes());
+        let unpacked = OffsetCurve::unpack(&packed).unwrap();
+        assert_eq!(curve, unpacked);
+    }
+
+    #[test]
+    fn swap_no_offset() {
+        let swap_source_amount: u128 = 1_000;
+        let swap_destination_amount: u128 = 50_000;
+        let source_amount: u128 = 100;
+        let curve = OffsetCurve::default();
+        let result = curve
+            .swap_without_fees(
+                source_amount,
+                swap_source_amount,
+                swap_destination_amount,
+                TradeDirection::AtoB,
+            )
+            .unwrap();
+        assert_eq!(result.source_amount_swapped, source_amount);
+        assert_eq!(result.destination_amount_swapped, 4545);
+        let result = curve
+            .swap_without_fees(
+                source_amount,
+                swap_source_amount,
+                swap_destination_amount,
+                TradeDirection::BtoA,
+            )
+            .unwrap();
+        assert_eq!(result.source_amount_swapped, source_amount);
+        assert_eq!(result.destination_amount_swapped, 4545);
+    }
+
+    #[test]
+    fn swap_offset() {
+        let swap_source_amount: u128 = 1_000_000;
+        let swap_destination_amount: u128 = 0;
+        let source_amount: u128 = 100;
+        let token_b_offset = 1_000_000;
+        let curve = OffsetCurve { token_b_offset };
+        let result = curve
+            .swap_without_fees(
+                source_amount,
+                swap_source_amount,
+                swap_destination_amount,
+                TradeDirection::AtoB,
+            )
+            .unwrap();
+        assert_eq!(result.source_amount_swapped, source_amount);
+        assert_eq!(result.destination_amount_swapped, source_amount - 1);
+
+        let bad_result = curve.swap_without_fees(
+            source_amount,
+            swap_source_amount,
+            swap_destination_amount,
+            TradeDirection::BtoA,
+        );
+        assert!(bad_result.is_none());
+    }
+
+    #[test]
+    fn swap_a_to_b_max_offset() {
+        let swap_source_amount: u128 = 10_000_000;
+        let swap_destination_amount: u128 = 1_000;
+        let source_amount: u128 = 1_000;
+        let token_b_offset = u64::MAX;
+        let curve = OffsetCurve { token_b_offset };
+        let result = curve
+            .swap_without_fees(
+                source_amount,
+                swap_source_amount,
+                swap_destination_amount,
+                TradeDirection::AtoB,
+            )
+            .unwrap();
+        assert_eq!(result.source_amount_swapped, source_amount);
+        assert_eq!(result.destination_amount_swapped, 1_844_489_958_375_117);
+    }
+
+    #[test]
+    fn swap_b_to_a_max_offset() {
+        let swap_source_amount: u128 = 10_000_000;
+        let swap_destination_amount: u128 = 1_000;
+        let source_amount: u128 = u64::MAX.into();
+        let token_b_offset = u64::MAX;
+        let curve = OffsetCurve { token_b_offset };
+        let result = curve
+            .swap_without_fees(
+                source_amount,
+                swap_source_amount,
+                swap_destination_amount,
+                TradeDirection::BtoA,
+            )
+            .unwrap();
+        assert_eq!(result.source_amount_swapped, 18_373_104_376_818_475_561);
+        assert_eq!(result.destination_amount_swapped, 499);
+    }
+
+    #[test]
+    fn pool_token_conversion() {
+        let tests: &[(u64, u128, u128, u128)] = &[
+            (10_000, 1_000_000, 1, 100_000),
+            (10, 1_000, 100, 100),
+            (1_251, 30, 1_288, 100_000),
+            (1_000_251, 1_000, 1_288, 100_000),
+            (1_000_000_000_000, 212, 10_000, 100_000),
+        ];
+        for (token_b_offset, swap_token_a_amount, swap_token_b_amount, token_a_amount) in
+            tests.iter()
+        {
+            let curve = OffsetCurve {
+                token_b_offset: *token_b_offset,
+            };
+            check_pool_token_conversion(
+                &curve,
+                *swap_token_a_amount,
+                *swap_token_b_amount,
+                *token_a_amount,
+            );
+        }
+    }
+}

token-swap/program/src/error.rs

@@ -88,6 +88,9 @@ pub enum SwapError {
     /// The provided curve parameters are invalid
     #[error("The provided curve parameters are invalid")]
     InvalidCurve,
+    /// The operation cannot be performed on the given curve
+    #[error("The operation cannot be performed on the given curve")]
+    UnsupportedCurveOperation,
 }
 impl From<SwapError> for ProgramError {
     fn from(e: SwapError) -> Self {

token-swap/program/src/processor.rs

@@ -448,6 +448,9 @@ impl Processor {
             return Err(ProgramError::IncorrectProgramId);
         }
         let token_swap = SwapInfo::unpack(&swap_info.data.borrow())?;
+        if !token_swap.swap_curve.calculator.allows_deposits() {
+            return Err(SwapError::UnsupportedCurveOperation.into());
+        }
         if *authority_info.key != Self::authority_id(program_id, swap_info.key, token_swap.nonce)? {
             return Err(SwapError::InvalidProgramAddress.into());
         }
@@ -475,13 +478,16 @@ impl Processor {
         let pool_mint = Self::unpack_mint(pool_mint_info, &token_swap.token_program_id)?;
         let pool_token_amount = to_u128(pool_token_amount)?;
         let pool_mint_supply = to_u128(pool_mint.supply)?;
+        let new_pool_mint_supply = pool_mint_supply
+            .checked_add(pool_token_amount)
+            .ok_or(SwapError::CalculationFailure)?;
 
         let calculator = token_swap.swap_curve.calculator;
 
         let results = calculator
             .pool_tokens_to_trading_tokens(
                 pool_token_amount,
-                pool_mint_supply,
+                new_pool_mint_supply,
                 to_u128(token_a.amount)?,
                 to_u128(token_b.amount)?,
             )
@@ -612,35 +618,41 @@ impl Processor {
         if token_a_amount < minimum_token_a_amount {
             return Err(SwapError::ExceededSlippage.into());
         }
-        if token_a_amount == 0 {
+        if token_a_amount == 0 && token_a.amount != 0 {
             return Err(SwapError::ZeroTradingTokens.into());
         }
         let token_b_amount = to_u64(results.token_b_amount)?;
         if token_b_amount < minimum_token_b_amount {
             return Err(SwapError::ExceededSlippage.into());
         }
-        if token_b_amount == 0 {
+        if token_b_amount == 0 && token_b.amount != 0 {
             return Err(SwapError::ZeroTradingTokens.into());
         }
 
-        Self::token_transfer(
+        let token_a_amount = std::cmp::min(token_a.amount, token_a_amount);
-            swap_info.key,
+        if token_a_amount > 0 {
-            token_program_info.clone(),
+            Self::token_transfer(
-            token_a_info.clone(),
+                swap_info.key,
-            dest_token_a_info.clone(),
+                token_program_info.clone(),
-            authority_info.clone(),
+                token_a_info.clone(),
-            token_swap.nonce,
+                dest_token_a_info.clone(),
-            token_a_amount,
+                authority_info.clone(),
-        )?;
+                token_swap.nonce,
-        Self::token_transfer(
+                token_a_amount,
-            swap_info.key,
+            )?;
-            token_program_info.clone(),
+        }
-            token_b_info.clone(),
+        let token_b_amount = std::cmp::min(token_b.amount, token_b_amount);
-            dest_token_b_info.clone(),
+        if token_b_amount > 0 {
-            authority_info.clone(),
+            Self::token_transfer(
-            token_swap.nonce,
+                swap_info.key,
-            token_b_amount,
+                token_program_info.clone(),
-        )?;
+                token_b_info.clone(),
+                dest_token_b_info.clone(),
+                authority_info.clone(),
+                token_swap.nonce,
+                token_b_amount,
+            )?;
+        }
         if withdraw_fee > 0 {
             Self::token_transfer(
                 swap_info.key,
@@ -793,6 +805,9 @@ impl PrintProgramError for SwapError {
             SwapError::InvalidCurve => {
                 msg!("Error: The provided curve parameters are invalid")
             }
+            SwapError::UnsupportedCurveOperation => {
+                msg!("Error: The operation cannot be performed on the given curve")
+            }
         }
     }
 }
@@ -812,7 +827,7 @@ mod tests {
         curve::calculator::{CurveCalculator, INITIAL_SWAP_POOL_AMOUNT},
         curve::{
             base::CurveType, constant_price::ConstantPriceCurve,
-            constant_product::ConstantProductCurve,
+            constant_product::ConstantProductCurve, offset::OffsetCurve,
         },
         instruction::{deposit, initialize, swap, withdraw},
     };
@@ -2014,6 +2029,53 @@ mod tests {
             accounts.initialize_swap().unwrap();
         }
 
+        // create invalid offset swap
+        {
+            let token_b_offset = 0;
+            let fees = Fees {
+                trade_fee_numerator,
+                trade_fee_denominator,
+                owner_trade_fee_numerator,
+                owner_trade_fee_denominator,
+                owner_withdraw_fee_numerator,
+                owner_withdraw_fee_denominator,
+                host_fee_numerator,
+                host_fee_denominator,
+            };
+            let swap_curve = SwapCurve {
+                curve_type: CurveType::Offset,
+                calculator: Box::new(OffsetCurve { token_b_offset }),
+            };
+            let mut accounts =
+                SwapAccountInfo::new(&user_key, fees, swap_curve, token_a_amount, token_b_amount);
+            assert_eq!(
+                Err(SwapError::InvalidCurve.into()),
+                accounts.initialize_swap()
+            );
+        }
+
+        // create valid offset swap
+        {
+            let token_b_offset = 10;
+            let fees = Fees {
+                trade_fee_numerator,
+                trade_fee_denominator,
+                owner_trade_fee_numerator,
+                owner_trade_fee_denominator,
+                owner_withdraw_fee_numerator,
+                owner_withdraw_fee_denominator,
+                host_fee_numerator,
+                host_fee_denominator,
+            };
+            let swap_curve = SwapCurve {
+                curve_type: CurveType::Offset,
+                calculator: Box::new(OffsetCurve { token_b_offset }),
+            };
+            let mut accounts =
+                SwapAccountInfo::new(&user_key, fees, swap_curve, token_a_amount, token_b_amount);
+            accounts.initialize_swap().unwrap();
+        }
+
         // wrong owner key in constraint
         {
             let new_key = Pubkey::new_unique();
@@ -2290,9 +2352,11 @@ mod tests {
         let mut accounts =
             SwapAccountInfo::new(&user_key, fees, swap_curve, token_a_amount, token_b_amount);
 
-        let deposit_a = token_a_amount / 10;
+        // depositing 10% of the current pool amount means that our share will
-        let deposit_b = token_b_amount / 10;
+        // be 1 / 11 of the final pool amount
         let pool_amount = INITIAL_SWAP_POOL_AMOUNT / 10;
+        let deposit_a = token_a_amount / 11;
+        let deposit_b = token_b_amount / 11;
 
         // swap not initialized
         {
@@ -2813,6 +2877,7 @@ mod tests {
                 )
             );
         }
+
         // correctly deposit
         {
             let (
@@ -3873,12 +3938,20 @@ mod tests {
             token_b_amount,
         );
         check_valid_swap_curve(
-            fees,
+            fees.clone(),
             CurveType::ConstantPrice,
             Box::new(ConstantPriceCurve {}),
             token_a_amount,
             token_b_amount,
         );
+        let token_b_offset = 10_000_000_000;
+        check_valid_swap_curve(
+            fees,
+            CurveType::Offset,
+            Box::new(OffsetCurve { token_b_offset }),
+            token_a_amount,
+            token_b_amount,
+        );
     }
 
     #[test]
@@ -3913,12 +3986,20 @@ mod tests {
             token_b_amount,
         );
         check_valid_swap_curve(
-            fees,
+            fees.clone(),
             CurveType::ConstantPrice,
             Box::new(ConstantPriceCurve {}),
             token_a_amount,
             token_b_amount,
         );
+        let token_b_offset = 1;
+        check_valid_swap_curve(
+            fees,
+            CurveType::Offset,
+            Box::new(OffsetCurve { token_b_offset }),
+            token_a_amount,
+            token_b_amount,
+        );
     }
 
     #[test]
@@ -4715,4 +4796,235 @@ mod tests {
             );
         }
     }
+
+    #[test]
+    fn test_overdraw_offset_curve() {
+        let trade_fee_numerator = 1;
+        let trade_fee_denominator = 10;
+        let owner_trade_fee_numerator = 1;
+        let owner_trade_fee_denominator = 30;
+        let owner_withdraw_fee_numerator = 1;
+        let owner_withdraw_fee_denominator = 30;
+        let host_fee_numerator = 10;
+        let host_fee_denominator = 100;
+
+        let token_a_amount = 1_000_000_000;
+        let token_b_amount = 0;
+        let fees = Fees {
+            trade_fee_numerator,
+            trade_fee_denominator,
+            owner_trade_fee_numerator,
+            owner_trade_fee_denominator,
+            owner_withdraw_fee_numerator,
+            owner_withdraw_fee_denominator,
+            host_fee_numerator,
+            host_fee_denominator,
+        };
+
+        let token_b_offset = 2_000_000;
+        let swap_curve = SwapCurve {
+            curve_type: CurveType::Offset,
+            calculator: Box::new(OffsetCurve { token_b_offset }),
+        };
+        let user_key = Pubkey::new_unique();
+        let swapper_key = Pubkey::new_unique();
+
+        let mut accounts =
+            SwapAccountInfo::new(&user_key, fees, swap_curve, token_a_amount, token_b_amount);
+
+        accounts.initialize_swap().unwrap();
+
+        let swap_token_a_key = accounts.token_a_key;
+        let swap_token_b_key = accounts.token_b_key;
+        let initial_a = 500_000;
+        let initial_b = 1_000;
+
+        let (
+            token_a_key,
+            mut token_a_account,
+            token_b_key,
+            mut token_b_account,
+            _pool_key,
+            _pool_account,
+        ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
+
+        // swap a to b way, fails, there's no liquidity
+        let a_to_b_amount = initial_a;
+        let minimum_token_b_amount = 0;
+
+        assert_eq!(
+            Err(SwapError::ZeroTradingTokens.into()),
+            accounts.swap(
+                &swapper_key,
+                &token_a_key,
+                &mut token_a_account,
+                &swap_token_a_key,
+                &swap_token_b_key,
+                &token_b_key,
+                &mut token_b_account,
+                a_to_b_amount,
+                minimum_token_b_amount,
+            )
+        );
+
+        // swap b to a, succeeds at offset price
+        let b_to_a_amount = initial_b;
+        let minimum_token_a_amount = 0;
+        accounts
+            .swap(
+                &swapper_key,
+                &token_b_key,
+                &mut token_b_account,
+                &swap_token_b_key,
+                &swap_token_a_key,
+                &token_a_key,
+                &mut token_a_account,
+                b_to_a_amount,
+                minimum_token_a_amount,
+            )
+            .unwrap();
+
+        // try a to b again, succeeds due to new liquidity
+        accounts
+            .swap(
+                &swapper_key,
+                &token_a_key,
+                &mut token_a_account,
+                &swap_token_a_key,
+                &swap_token_b_key,
+                &token_b_key,
+                &mut token_b_account,
+                a_to_b_amount,
+                minimum_token_b_amount,
+            )
+            .unwrap();
+
+        // try a to b again, fails due to no more liquidity
+        assert_eq!(
+            Err(SwapError::ZeroTradingTokens.into()),
+            accounts.swap(
+                &swapper_key,
+                &token_a_key,
+                &mut token_a_account,
+                &swap_token_a_key,
+                &swap_token_b_key,
+                &token_b_key,
+                &mut token_b_account,
+                a_to_b_amount,
+                minimum_token_b_amount,
+            )
+        );
+
+        // Try to deposit, fails because deposits are not allowed for offset
+        // curve swaps
+        {
+            let initial_a = 100;
+            let initial_b = 100;
+            let pool_amount = 100;
+            let (
+                token_a_key,
+                mut token_a_account,
+                token_b_key,
+                mut token_b_account,
+                pool_key,
+                mut pool_account,
+            ) = accounts.setup_token_accounts(&user_key, &swapper_key, initial_a, initial_b, 0);
+            assert_eq!(
+                Err(SwapError::UnsupportedCurveOperation.into()),
+                accounts.deposit(
+                    &swapper_key,
+                    &token_a_key,
+                    &mut token_a_account,
+                    &token_b_key,
+                    &mut token_b_account,
+                    &pool_key,
+                    &mut pool_account,
+                    pool_amount,
+                    initial_a,
+                    initial_b,
+                )
+            );
+        }
+    }
+
+    #[test]
+    fn test_withdraw_all_offset_curve() {
+        let trade_fee_numerator = 1;
+        let trade_fee_denominator = 10;
+        let owner_trade_fee_numerator = 1;
+        let owner_trade_fee_denominator = 30;
+        let owner_withdraw_fee_numerator = 0;
+        let owner_withdraw_fee_denominator = 30;
+        let host_fee_numerator = 10;
+        let host_fee_denominator = 100;
+
+        let token_a_amount = 1_000_000_000;
+        let token_b_amount = 10;
+        let fees = Fees {
+            trade_fee_numerator,
+            trade_fee_denominator,
+            owner_trade_fee_numerator,
+            owner_trade_fee_denominator,
+            owner_withdraw_fee_numerator,
+            owner_withdraw_fee_denominator,
+            host_fee_numerator,
+            host_fee_denominator,
+        };
+
+        let token_b_offset = 2_000_000;
+        let swap_curve = SwapCurve {
+            curve_type: CurveType::Offset,
+            calculator: Box::new(OffsetCurve { token_b_offset }),
+        };
+        let total_pool = swap_curve.calculator.new_pool_supply();
+        let user_key = Pubkey::new_unique();
+        let withdrawer_key = Pubkey::new_unique();
+
+        let mut accounts =
+            SwapAccountInfo::new(&user_key, fees, swap_curve, token_a_amount, token_b_amount);
+
+        accounts.initialize_swap().unwrap();
+
+        let (
+            token_a_key,
+            mut token_a_account,
+            token_b_key,
+            mut token_b_account,
+            _pool_key,
+            _pool_account,
+        ) = accounts.setup_token_accounts(&user_key, &withdrawer_key, 0, 0, 0);
+
+        let pool_key = accounts.pool_token_key;
+        let mut pool_account = accounts.pool_token_account.clone();
+
+        // Withdraw takes all tokens for A and B.
+        // The curve's calculation for token B will say to transfer
+        // `token_b_offset + token_b_amount`, but only `token_b_amount` will be
+        // moved.
+        accounts
+            .withdraw(
+                &user_key,
+                &pool_key,
+                &mut pool_account,
+                &token_a_key,
+                &mut token_a_account,
+                &token_b_key,
+                &mut token_b_account,
+                total_pool.try_into().unwrap(),
+                0,
+                0,
+            )
+            .unwrap();
+
+        let token_a = spl_token::state::Account::unpack(&token_a_account.data).unwrap();
+        assert_eq!(token_a.amount, token_a_amount);
+        let token_b = spl_token::state::Account::unpack(&token_b_account.data).unwrap();
+        assert_eq!(token_b.amount, token_b_amount);
+        let swap_token_a =
+            spl_token::state::Account::unpack(&accounts.token_a_account.data).unwrap();
+        assert_eq!(swap_token_a.amount, 0);
+        let swap_token_b =
+            spl_token::state::Account::unpack(&accounts.token_b_account.data).unwrap();
+        assert_eq!(swap_token_b.amount, 0);
+    }
 }