Small readability edits

noir-projects/noir-contracts/contracts/amm_contract/src/state.nr → noir-projects/noir-contracts/contracts/amm_contract/src/config.nr

@@ -1,26 +1,26 @@
 use dep::aztec::protocol_types::{address::AztecAddress, traits::{Deserialize, Serialize}};
 
-global STATE_LENGTH: u32 = 3;
+global CONFIG_LENGTH: u32 = 3;
 
 /// We store the tokens of the pool in a struct such that to load it from SharedImmutable asserts only a single
 /// merkle proof.
 /// (Once we actually do the optimization. WIP in https://github.com/AztecProtocol/aztec-packages/pull/8022).
-pub struct State {
+pub struct Config {
     pub token0: AztecAddress,
     pub token1: AztecAddress,
     pub liquidity_token: AztecAddress,
 }
 
 // Note: I could not get #[derive(Serialize)] to work so I had to implement it manually.
-impl Serialize<STATE_LENGTH> for State {
-    fn serialize(self: Self) -> [Field; STATE_LENGTH] {
+impl Serialize<CONFIG_LENGTH> for Config {
+    fn serialize(self: Self) -> [Field; CONFIG_LENGTH] {
         [self.token0.to_field(), self.token1.to_field(), self.liquidity_token.to_field()]
     }
 }
 
-impl Deserialize<STATE_LENGTH> for State {
-    fn deserialize(fields: [Field; STATE_LENGTH]) -> Self {
-        State {
+impl Deserialize<CONFIG_LENGTH> for Config {
+    fn deserialize(fields: [Field; CONFIG_LENGTH]) -> Self {
+        Self {
             token0: AztecAddress::from_field(fields[0]),
             token1: AztecAddress::from_field(fields[1]),
             liquidity_token: AztecAddress::from_field(fields[2]),

noir-projects/noir-contracts/contracts/amm_contract/src/lib.nr

@@ -1,61 +1,63 @@
-/// Given some amount of an asset and pair reserves, returns an equivalent amount of the other asset.
-/// copy of https://github.com/Uniswap/v2-periphery/blob/0335e8f7e1bd1e8d8329fd300aea2ef2f36dd19f/contracts/libraries/UniswapV2Library.sol#L36
-fn get_quote(amountA: U128, reserveA: U128, reserveB: U128) -> U128 {
-    assert(amountA > U128::zero(), "INSUFFICIENT_AMOUNT");
-    assert((reserveA > U128::zero()) & (reserveB > U128::zero()), "INSUFFICIENT_LIQUIDITY");
-    (amountA * reserveB) / reserveA
-}
-
-/// Given an input amount of an asset and pair reserves, returns the maximum output amount of the other asset.
-/// copy of https://github.com/Uniswap/v2-periphery/blob/0335e8f7e1bd1e8d8329fd300aea2ef2f36dd19f/contracts/libraries/UniswapV2Library.sol#L43
-pub fn get_amount_out(amount_in: U128, reserve_in: U128, reserve_out: U128) -> U128 {
+/// Given an input amount of an asset and pair balances, returns the maximum output amount of the other asset.
+pub fn get_amount_out(amount_in: U128, balance_in: U128, balance_out: U128) -> U128 {
+    // Based on Uniswap v2: https://github.com/Uniswap/v2-periphery/blob/0335e8f7e1bd1e8d8329fd300aea2ef2f36dd19f/contracts/libraries/UniswapV2Library.sol#L43
     assert(amount_in > U128::zero(), "INSUFFICIENT_INPUT_AMOUNT");
-    assert((reserve_in > U128::zero()) & (reserve_out > U128::zero()), "INSUFFICIENT_LIQUIDITY");
+    assert((balance_in > U128::zero()) & (balance_out > U128::zero()), "INSUFFICIENT_LIQUIDITY");
     let amount_in_with_fee = amount_in * U128::from_integer(997);
-    let numerator = amount_in_with_fee * reserve_out;
-    let denominator = reserve_in * U128::from_integer(1000) + amount_in_with_fee;
+    let numerator = amount_in_with_fee * balance_out;
+    let denominator = balance_in * U128::from_integer(1000) + amount_in_with_fee;
     numerator / denominator
 }
 
-/// Given an output amount of an asset and pair reserves, returns a required input amount of the other asset.
-/// copy of https://github.com/Uniswap/v2-periphery/blob/0335e8f7e1bd1e8d8329fd300aea2ef2f36dd19f/contracts/libraries/UniswapV2Library.sol#L53
-pub fn get_amount_in(amount_out: U128, reserve_in: U128, reserve_out: U128) -> U128 {
+/// Given an output amount of an asset and pair balances, returns a required input amount of the other asset.
+pub fn get_amount_in(amount_out: U128, balance_in: U128, balance_out: U128) -> U128 {
+    // Based on Uniswap v2: https://github.com/Uniswap/v2-periphery/blob/0335e8f7e1bd1e8d8329fd300aea2ef2f36dd19f/contracts/libraries/UniswapV2Library.sol#L53
     assert(amount_out > U128::zero(), "INSUFFICIENT_OUTPUT_AMOUNT");
-    assert((reserve_in > U128::zero()) & (reserve_out > U128::zero()), "INSUFFICIENT_LIQUIDITY");
-    let numerator = reserve_in * amount_out * U128::from_integer(1000);
-    let denominator = (reserve_out - amount_out) * U128::from_integer(997);
+    assert((balance_in > U128::zero()) & (balance_out > U128::zero()), "INSUFFICIENT_LIQUIDITY");
+    let numerator = balance_in * amount_out * U128::from_integer(1000);
+    let denominator = (balance_out - amount_out) * U128::from_integer(997);
     (numerator / denominator) + U128::from_integer(1)
 }
 
-/// Given the desired amounts and reserves of token0 and token1 returns the optimal amount of token0 and token1 to be added to the pool.
+/// Given the desired amounts and balances of token0 and token1 returns the optimal amount of token0 and token1 to be added to the pool.
 pub fn get_amounts_to_add(
-    amount0_desired: U128,
-    amount1_desired: U128,
+    amount0_max: U128,
+    amount1_max: U128,
     amount0_min: U128,
     amount1_min: U128,
-    reserve0: U128,
-    reserve1: U128,
+    balance0: U128,
+    balance1: U128,
 ) -> (U128, U128) {
-    let mut amount0 = amount0_desired;
-    let mut amount1 = amount1_desired;
-    if ((reserve0 != U128::zero()) | (reserve1 != U128::zero())) {
+    if (balance0 == U128::zero()) | (balance1 == U128::zero()) {
+        // The token balances should only be zero when initializing the pool. In this scenario there is no prior ratio
+        // to follow so we simply
+        (amount0_max, amount1_max)
+    } else {
         // First calculate the optimal amount of token1 based on the desired amount of token0.
-        let amount1_optimal = get_quote(amount0_desired, reserve0, reserve1);
-        if (amount1_optimal <= amount1_desired) {
+        let amount1_optimal = get_equivalent_amount(amount0_max, balance0, balance1);
+        if (amount1_optimal <= amount1_max) {
             // Revert if the optimal amount of token1 is less than the desired amount of token1.
             assert(amount1_optimal >= amount1_min, "INSUFFICIENT_1_AMOUNT");
-            amount0 = amount0_desired;
-            amount1 = amount1_optimal;
+            (amount0_max, amount1_optimal)
         } else {
             // We got more amount of token1 than desired so we try repeating the process but this time by quoting
             // based on token1.
-            let amount0_optimal = get_quote(amount1_desired, reserve1, reserve0);
-            assert(amount0_optimal <= amount0_desired);
+            let amount0_optimal = get_equivalent_amount(amount1_max, balance1, balance0);
+            assert(amount0_optimal <= amount0_max);
             assert(amount0_optimal >= amount0_min, "INSUFFICIENT_0_AMOUNT");
-            amount0 = amount0_optimal;
-            amount1 = amount1_desired;
+
+            (amount0_optimal, amount1_max)
         }
     }
+}
+
+/// Given some amount of an asset and pair balances, returns an equivalent amount of the other asset. Tokens should be
+/// added and removed from the Pool respecting this ratio.
+fn get_equivalent_amount(amount0: U128, balance0: U128, balance1: U128) -> U128 {
+    // Based on Uniswap v2:  https://github.com/Uniswap/v2-periphery/blob/0335e8f7e1bd1e8d8329fd300aea2ef2f36dd19f/contracts/libraries/UniswapV2Library.sol#L36
+    assert((balance0 > U128::zero()) & (balance1 > U128::zero()), "INSUFFICIENT_LIQUIDITY");
 
-    (amount0, amount1)
+    // This is essentially the Rule of Three, since we're computing proportional ratios. Note we divide at the end to
+    // avoid introducing too much error due to truncation.
+    (amount0 * balance1) / balance0
 }