optimize reentrancy check & fix test

contracts/main/TTDex.ligo

@@ -1,7 +1,8 @@
 #define FA2_STANDARD_ENABLED
 #define TOKEN_TO_TOKEN_ENABLED
-#include "../partials/TTDex.ligo"
+#include "../partials/ITTDex.ligo"
 #include "../partials/TTMethodDex.ligo"
+#include "../partials/TTDex.ligo"
 
 (* DexFA2 - Contract for exchanges for XTZ - FA2 token pair *)
 function main (const p : full_action; const s : full_dex_storage) : full_return is

contracts/partials/ITTDex.ligo

@@ -4,7 +4,7 @@
 
 (* record that represents account shares *)
 type account_info is record [
-  balance           : nat; (* liquid tokens *)
+  balance           : nat; (* LP tokens *)
   allowances        : set (address); (* accounts allowed to act on behalf of the user *)
 ]
 
@@ -23,8 +23,8 @@ type token_type is
 | Fa2
 
 type pair_info is record [
-  token_a_pool        : nat; (* tez reserves in the pool *)
-  token_b_pool        : nat; (* token reserves in the pool *)
+  token_a_pool        : nat; (* token A reserves in the pool *)
+  token_b_pool        : nat; (* token B reserves in the pool *)
   total_supply        : nat; (* total shares count *)
 ]
 
@@ -41,7 +41,7 @@ type token_pair is bytes
 
 (* record for the dex storage *)
 type dex_storage is record [
-  entered             : bool;
+  entered             : bool; (* reentrancy protection *)
   pairs_count         : nat; (* total shares count *)
   tokens              : big_map(nat, tokens_info); (* all the tokens list *)
   token_to_id         : big_map(token_pair, nat); (* all the tokens list *)
@@ -56,15 +56,15 @@ type swap_slice_type is record [
 ]
 
 type swap_side is record [
-  pool : nat;
-  token : address;
-  id: nat;
-  standard: token_type;
+  pool                    : nat;
+  token                   : address;
+  id                      : nat;
+  standard                : token_type;
 ]
 
 type swap_data is record [
-  to_: swap_side;
-  from_: swap_side;
+  to_                     : swap_side;
+  from_                   : swap_side;
 ]
 
 type internal_swap_type is record [
@@ -81,34 +81,34 @@ type internal_swap_type is record [
 type token_to_token_route_params is
   [@layout:comb]
   record [
-    swaps                 : list(swap_slice_type);
+    swaps                 : list(swap_slice_type); (* swap operations list*)
     amount_in             : nat; (* amount of tokens to be exchanged *)
-    min_amount_out        : nat; (* min amount of XTZ received to accept exchange *)
+    min_amount_out        : nat; (* min amount of tokens received to accept exchange *)
     receiver              : address; (* tokens receiver *)
   ]
 
 type invest_liquidity_params is
   [@layout:comb]
   record [
     pair            : tokens_info;
-    token_a_in      : nat; (* min amount of XTZ received to accept the divestment *)
-    token_b_in      : nat; (* min amount of tokens received to accept the divestment *)
+    token_a_in      : nat; (* min amount of tokens A invested  *)
+    token_b_in      : nat; (* min amount of tokens B invested *)
   ]
 
 type divest_liquidity_params is
   [@layout:comb]
   record [
     pair                 : tokens_info;
-    min_token_a_out      : nat; (* min amount of XTZ received to accept the divestment *)
-    min_token_b_out      : nat; (* min amount of tokens received to accept the divestment *)
+    min_token_a_out      : nat; (* min amount of tokens A received to accept the divestment *)
+    min_token_b_out      : nat; (* min amount of tokens B received to accept the divestment *)
     shares               : nat; (* amount of shares to be burnt *)
   ]
 
 type dex_action is
 | InitializeExchange          of invest_liquidity_params  (* sets initial liquidity *)
-| TokenToTokenRoutePayment    of token_to_token_route_params  (* exchanges XTZ to tokens and sends them to receiver *)
-| InvestLiquidity             of invest_liquidity_params  (* mints min shares after investing tokens and XTZ *)
-| DivestLiquidity             of divest_liquidity_params  (* burns shares and sends tokens and XTZ to the owner *)
+| TokenToTokenRoutePayment    of token_to_token_route_params  (* exchanges token to another token and sends them to receiver *)
+| InvestLiquidity             of invest_liquidity_params  (* mints min shares after investing tokens *)
+| DivestLiquidity             of divest_liquidity_params  (* burns shares and sends tokens to the owner *)
 
 type use_params is dex_action
 type get_reserves_params is record [

contracts/partials/TTDex.ligo

@@ -1,5 +1,3 @@
-#include "./ITTDex.ligo"
-
 (* Route exchange-specific action
 
 Due to the fabulous storage, gas and operation size limits
@@ -11,7 +9,10 @@ The function is responsible for fiding the appropriate method
 based on the argument type.
 
 *)
-[@inline] function call_dex (const p : dex_action; const this : address; const s : full_dex_storage) :  full_return is
+[@inline] function call_dex(
+  const p : dex_action;
+  const this : address;
+  const s : full_dex_storage) : full_return is
 block {
     const idx : nat = case p of
       | InitializeExchange(n) -> 0n
@@ -37,13 +38,21 @@ The function is responsible for fiding the appropriate method
 based on the provided index.
 
 *)
-[@inline] function call_token (const p : token_action; const this : address; const idx : nat; const s : full_dex_storage) :  full_return is
+[@inline] function call_token(
+  const p : token_action;
+  const this : address;
+  const idx : nat;
+  const s : full_dex_storage) : full_return is
 block {
   const res : return = case s.token_lambdas[idx] of
     Some(f) -> f(p, s.storage, this)
     | None -> (failwith("Dex/function-not-set") : return)
   end;
   s.storage := res.1;
+  res.0 := Tezos.transaction(
+      unit,
+      0mutez,
+      get_close_entrypoint(this)) # res.0;
 } with (res.0, s)
 
 [@inline] function close (const s : full_dex_storage) :  full_dex_storage is
@@ -58,7 +67,9 @@ block {
 } with s
 
 (* Return the reserves to the contracts. *)
-[@inline] function get_reserves (const params : get_reserves_params; const s : full_dex_storage) : full_return is
+[@inline] function get_reserves(
+  const params : get_reserves_params;
+  const s : full_dex_storage) : full_return is
 block {
   const pair : pair_info = case s.storage.pairs[params.token_id] of
     None -> record [
@@ -77,7 +88,10 @@ block {
   ], s)
 
 (* Set the dex function code to factory storage *)
-[@inline] function set_dex_function (const idx : nat; const f : dex_func; const s : full_dex_storage) : full_dex_storage is
+[@inline] function set_dex_function(
+  const idx : nat;
+  const f : dex_func;
+  const s : full_dex_storage) : full_dex_storage is
 block {
   case s.dex_lambdas[idx] of
     Some(n) -> failwith("Dex/function-set")
@@ -86,7 +100,10 @@ block {
 } with s
 
 (* Set the token function code to factory storage *)
-[@inline] function set_token_function (const idx : nat; const f : token_func; const s : full_dex_storage) : full_dex_storage is
+[@inline] function set_token_function(
+  const idx : nat;
+  const f : token_func;
+  const s : full_dex_storage) : full_dex_storage is
 block {
   case s.token_lambdas[idx] of
     Some(n) -> failwith("Dex/function-set")

contracts/partials/TTMethodDex.ligo

@@ -102,7 +102,7 @@ function wrap_fa12_transfer_trx(
   const value : nat) : transfer_type_fa12 is
   TransferTypeFA12(owner, (receiver, value))
 
-(* Helper function to get token contract *)
+(* Helper function to get fa2 token contract *)
 function get_fa2_token_contract(
   const token_address : address) : contract(transfer_type_fa2) is
   case (Tezos.get_entrypoint_opt(
@@ -112,7 +112,7 @@ function get_fa2_token_contract(
   | None -> (failwith("Dex/not-token") : contract(transfer_type_fa2))
   end;
 
-(* Helper function to get token contract *)
+(* Helper function to get fa1.2 token contract *)
 function get_fa12_token_contract(
   const token_address : address) : contract(transfer_type_fa12) is
   case (Tezos.get_entrypoint_opt(
@@ -122,7 +122,7 @@ function get_fa12_token_contract(
   | None -> (failwith("Dex/not-token") : contract(transfer_type_fa12))
   end;
 
-(* Helper function to get self reentrancy entrypoint *)
+(* Helper function to get the reentrancy entrypoint of the current contract *)
 function get_close_entrypoint(
   const contract_address : address) : contract(unit) is
   case (Tezos.get_entrypoint_opt(
@@ -131,6 +131,7 @@ function get_close_entrypoint(
   | None -> (failwith("Dex/no-close-entrypoint") : contract(unit))
   end;
 
+(* Helper function to transfer fa2 tokens *)
 function transfer_fa2(
   const sender_ : address;
   const receiver : address;
@@ -147,6 +148,7 @@ function transfer_fa2(
     get_fa2_token_contract(contract_address)
   );
 
+(* Helper function to transfer fa1.2 tokens *)
 function transfer_fa12(
   const sender_ : address;
   const receiver : address;
@@ -161,6 +163,7 @@ function transfer_fa12(
     get_fa12_token_contract(contract_address)
   );
 
+(* Helper function to transfer the asset based on its standard *)
 function typed_transfer(
   const sender_ : address;
   const receiver : address;
@@ -190,12 +193,7 @@ function initialize_exchange(
   const s : dex_storage;
   const this: address) :  return is
   block {
-    var operations: list(operation) := list[
-      Tezos.transaction(
-        unit,
-        0mutez,
-        get_close_entrypoint(this)
-    )];
+    var operations: list(operation) := list[];
     case p of
       InitializeExchange(params) -> {
         if s.entered
@@ -209,7 +207,7 @@ function initialize_exchange(
         if params.pair.token_a_address > params.pair.token_b_address
         then failwith("Dex/wrong-pair") else skip;
 
-        (* get par info*)
+        (* read pair info*)
         const res : (pair_info * nat) = get_pair(params.pair, s);
         const pair : pair_info = res.0;
         const token_id : nat = res.1;
@@ -225,9 +223,9 @@ function initialize_exchange(
         then failwith("Dex/non-zero-reserves") else skip;
         if pair.total_supply =/= 0n  (* no shares owned *)
         then failwith("Dex/non-zero-shares") else skip;
-        if params.token_a_in < 1n (* XTZ provided *)
+        if params.token_a_in < 1n (* no token A invested *)
         then failwith("Dex/no-token-a") else skip;
-        if params.token_b_in < 1n (* XTZ provided *)
+        if params.token_b_in < 1n (* no token B invested *)
         then failwith("Dex/no-token-b") else skip;
 
         (* update pool reserves *)
@@ -289,7 +287,7 @@ function internal_token_to_token_swap(
     if params.pair.token_a_address > params.pair.token_b_address
     then failwith("Dex/wrong-pair") else skip;
 
-    (* get par info*)
+    (* get pair info*)
     const res : (pair_info * nat) = get_pair(params.pair, tmp.s);
     const pair : pair_info = res.0;
     const token_id : nat = res.1;
@@ -300,6 +298,7 @@ function internal_token_to_token_swap(
     then failwith("Dex/not-launched") else skip;
     if tmp.amount_in = 0n (* non-zero amount of tokens exchanged *)
     then failwith ("Dex/zero-amount-in") else skip;
+    (* ensure the route is corresponding *)
     if swap.from_.token =/= tmp.token_address_in
     or swap.from_.id =/= tmp.token_id_in
     then failwith("Dex/wrong-route") else skip;
@@ -312,8 +311,8 @@ function internal_token_to_token_swap(
     (* calculate swapped token amount *)
     const out : nat = numerator / denominator;
 
-    (* ensure requirements *)
-    if out > swap.to_.pool / 3n (* the price impact isn't too high *)
+    (* ensure the price impact isn't too high *)
+    if out > swap.to_.pool / 3n
     then failwith("Dex/high-out") else skip;
 
     (* update pools amounts *)
@@ -349,11 +348,7 @@ function token_to_token_route(
   const s : dex_storage;
   const this : address) : return is
   block {
-    var operations: list(operation) := list[Tezos.transaction(
-      unit,
-      0mutez,
-      get_close_entrypoint(this)
-    )];
+    var operations: list(operation) := list[];
     case p of
     | InitializeExchange(n) -> skip
     | TokenToTokenRoutePayment(params) -> {
@@ -362,14 +357,14 @@ function token_to_token_route(
         else s.entered := True;
 
         (* validate input params *)
-        if List.size(params.swaps) < 1n (* non-zero amount of tokens exchanged *)
+        if List.size(params.swaps) < 1n (* the route is empty *)
         then failwith ("Dex/too-few-swaps") else skip;
         if params.amount_in = 0n (* non-zero amount of tokens exchanged *)
         then failwith ("Dex/zero-amount-in") else skip;
-        if params.min_amount_out = 0n (* non-zero amount of tokens exchanged *)
+        if params.min_amount_out = 0n (* non-zero amount of tokens to receive *)
         then failwith ("Dex/zero-min-amount-out") else skip;
 
-        (* collect the operations to execute *)
+        (* get the first exchange info *)
         const first_swap : swap_slice_type = case List.head_opt(params.swaps) of
           Some(swap) -> swap
         | None -> (failwith("Dex/zero-swaps") : swap_slice_type)
@@ -429,7 +424,7 @@ function token_to_token_route(
         (* update storage*)
         s := tmp.s;
 
-        (* add token transfer to user's account *)
+        (* add token transfer to user's account to operations *)
         const last_operation : operation = case tmp.operation of
           Some(o) -> o
         | None -> (failwith("Dex/too-few-swaps") : operation)
@@ -448,11 +443,7 @@ function invest_liquidity(
   const s : dex_storage;
   const this: address) : return is
   block {
-    var operations: list(operation) := list[Tezos.transaction(
-      unit,
-      0mutez,
-      get_close_entrypoint(this)
-    )];
+    var operations: list(operation) := list[];
     case p of
     | InitializeExchange(n) -> skip
     | TokenToTokenRoutePayment(n) -> skip
@@ -468,7 +459,7 @@ function invest_liquidity(
         if params.pair.token_a_address > params.pair.token_b_address
         then failwith("Dex/wrong-pair") else skip;
 
-        (* get par info*)
+        (* read pair info*)
         const res : (pair_info * nat) = get_pair(params.pair, s);
         const pair : pair_info = res.0;
         const token_id : nat = res.1;
@@ -487,8 +478,8 @@ function invest_liquidity(
           then shares_a_purchased
           else shares_b_purchased;
 
-        (* ensure *)
-        if shares_purchased = 0n (* purchsed shares satisfy required minimum *)
+        (* ensure purchsed shares satisfy required minimum *)
+        if shares_purchased = 0n
         then failwith("Dex/wrong-params") else skip;
 
         (* calculate tokens to be withdrawn *)
@@ -558,11 +549,7 @@ function divest_liquidity(
   const s : dex_storage;
   const this: address) :  return is
   block {
-    var operations: list(operation) := list[Tezos.transaction(
-      unit,
-      0mutez,
-      get_close_entrypoint(this)
-    )];
+    var operations: list(operation) := list[];
     case p of
     | InitializeExchange(token_amount) -> skip
     | TokenToTokenRoutePayment(n) -> skip
@@ -579,7 +566,7 @@ function divest_liquidity(
         if params.pair.token_a_address > params.pair.token_b_address
         then failwith("Dex/wrong-pair") else skip;
 
-        (* get par info*)
+        (* read pair info*)
         const res : (pair_info * nat) = get_pair(params.pair, s);
         const pair : pair_info = res.0;
         const token_id : nat = res.1;
@@ -628,7 +615,7 @@ function divest_liquidity(
         (* update storage *)
         s.pairs[token_id] := pair;
 
-        (* prepare operations with XTZ and tokens to user *)
+        (* prepare operations with tokens to user *)
         operations :=
           typed_transfer(
             this,