Add lp.

contract/src/error.rs

@@ -5,5 +5,13 @@ pub const PROMISE_TOO_MANY_RESULTS: &str = "Cross contract call returned more th
 pub const PROMISE_WRONG_VALUE_RECEIVED: &str = "Cross contract call returned invalid value.";
 pub const PROMISE_CALL_FAILED: &str = "Cross contract call failed.";
 
+pub const INVALID_TOKEN_DECIMAL: &str = "Token decimal is either too small (<1) or too big (>24).";
+pub const INTERNAL_OVERFLOW_ERROR: &str = "There is an internal error when calculating due to overflow.";
+
+pub const AMM_NOT_FUNCTIONAL_YET: &str = "This AMM contract is not yet fully functional — most likely because liquidity has not been provided.";
 pub const TOKEN_METADATA_NOT_INITIALISED: &str = "Token metadata has not been initialised.";
-pub const TOKEN_ZERO_BALANCE: &str = "Token has 0 balance, most likely liquidity has not been provided.";
+
+pub const INVALID_TOKEN_RECEIVER_MESSAGE: &str = "Invalid fungible token token receiver message, should be either 'lp_deposit' or 'swap'.";
+pub const INVALID_LP_DEPOSIT_SENDER: &str = "lp_deposit sender is not the owner of this AMM.";
+pub const INVALID_LP_DEPOSIT_TOKEN: &str = "lp_deposit token is not one of two tokens set on this AMM.";
+pub const INVALID_LP_DEPOSIT_AMOUNT: &str = "lp_deposit amount cannot be zero.";

contract/src/lib.rs

@@ -1,27 +1,44 @@
 mod error;
 mod token;
 
-use near_contract_standards::fungible_token::metadata::FungibleTokenMetadata;
+use near_contract_standards::fungible_token::{
+    metadata::FungibleTokenMetadata,
+    receiver::FungibleTokenReceiver,
+};
 use near_sdk::{
-    AccountId, PanicOnDefault, Promise, PromiseResult,
+    AccountId, BorshStorageKey, PanicOnDefault, Promise, PromiseOrValue, PromiseResult,
     borsh::{self, BorshDeserialize, BorshSerialize},
     env,
-    log,
-    near_bindgen,
+    json_types::U128,
+    log, near_bindgen,
     serde::{Deserialize, Serialize},
-    serde_json,
+    serde_json, Balance,
+    store::Vector,
 };
+use std::fmt;
 
 use error::*;
 use token::{Token, TokenMetadata, ext_fungible_token};
 
+#[derive(Serialize, Deserialize)]
+#[serde(crate = "near_sdk::serde")]
+pub struct ContractMetadata {
+    tokens: Vec<TokenMetadata>,
+    ratio: u128,
+}
+
+#[derive(BorshStorageKey, BorshSerialize)]
+pub(crate) enum StorageKey {
+    Token,
+}
+
 #[near_bindgen]
 #[derive(BorshDeserialize, BorshSerialize, PanicOnDefault)]
 pub struct Contract {
     owner_address: AccountId,
-    token_a: Token,
-    token_b: Token,
+    tokens: Vector<Token>,
     constant_product: u128,
+    functional: bool,
 }
 
 #[near_bindgen]
@@ -30,25 +47,64 @@ impl Contract {
     pub fn new(owner_id: AccountId, token_a_id: AccountId, token_b_id: AccountId) -> Self {
         assert_ne!(owner_id, env::current_account_id(), "{}", OWNER_CANNOT_BE_CONTRACT_ACCOUNT_ITSELF);
         assert_ne!(token_a_id, token_b_id, "{}", DUPLICATE_TOKENS);
+
         Self::get_token_metadata(token_a_id.clone());
         Self::get_token_metadata(token_b_id.clone());
+
+        let mut tokens = Vector::new(StorageKey::Token);
+        tokens.push(Token::new(token_a_id));
+        tokens.push(Token::new(token_b_id));
+
         Self {
             owner_address: owner_id,
-            token_a: Token::new(token_a_id),
-            token_b: Token::new(token_b_id),
+            tokens,
             constant_product: 0,
+            functional: false,
         }
     }
 
     pub fn get_metadata(&self) -> ContractMetadata {
-        let ratio = self.token_a.get_balance() as f64 / self.token_b.get_balance() as f64;
+        assert!(self.functional, "{}", AMM_NOT_FUNCTIONAL_YET);
+        let ratio = self.tokens
+            .iter()
+            .fold(1u128, |acc, token| {
+                token.get_canonical_balance().checked_div(acc).expect(INTERNAL_OVERFLOW_ERROR)
+            });
         ContractMetadata {
-            token_a: self.token_a.get_metadata(),
-            token_b: self.token_b.get_metadata(),
+            tokens: self.tokens.iter().map(|token| token.get_metadata()).collect(),
             ratio,
         }
     }
 
+    fn deposit(&mut self, token_in: AccountId, amount: Balance) {
+        self.tokens
+            .iter_mut()
+            .find(|token| token.check_address(&token_in))
+            .expect(INVALID_LP_DEPOSIT_TOKEN)
+            .set_balance(amount);
+
+        match self.functional {
+            false => {
+                if self.tokens.iter().all(|token| token.get_balance() > 0) {
+                    self.set_constant_product();
+                    log!("Turning on the engine!");
+                    self.functional = true;            
+                }
+            },
+            true => {
+                self.set_constant_product();
+            }
+        };
+    }
+
+    fn set_constant_product(&mut self) {
+        self.constant_product = self.tokens
+            .iter()
+            .fold(1u128, |acc, token| {
+                token.get_canonical_balance().checked_mul(acc).expect(INTERNAL_OVERFLOW_ERROR)
+            })
+    }
+
     fn get_token_metadata(token_id: AccountId) -> Promise {
         ext_fungible_token::ext(token_id.clone())
             .ft_metadata()
@@ -61,18 +117,16 @@ impl Contract {
     #[private]
     pub fn post_fungible_token_metadata(&mut self, token_id: AccountId) {
         assert_eq!(env::promise_results_count(), 1, "{}", PROMISE_TOO_MANY_RESULTS);
-        log!("Received callback from ft_metadata cross contract call!");
+        log!("Received callback from {}'s ft_metadata cross contract call!", token_id);
         match env::promise_result(0) {
             PromiseResult::NotReady => unreachable!(),
             PromiseResult::Successful(value) => {
                 if let Ok(metadata) = serde_json::from_slice::<FungibleTokenMetadata>(&value) {
-                    if token_id.as_str() == self.token_a.get_address().as_str() {
-                        self.token_a.set_metadata(metadata);
-                    } else if token_id.as_str() == self.token_b.get_address().as_str() {
-                        self.token_b.set_metadata(metadata);
-                    } else {
-                        env::panic_str(PROMISE_WRONG_VALUE_RECEIVED);
-                    }
+                    self.tokens
+                        .iter_mut()
+                        .find(|token| token.check_address(&token_id))
+                        .expect(PROMISE_WRONG_VALUE_RECEIVED)
+                        .set_metadata(metadata);
                 } else {
                     env::panic_str(PROMISE_WRONG_VALUE_RECEIVED);
                 }
@@ -84,8 +138,43 @@ impl Contract {
 
 #[derive(Serialize, Deserialize)]
 #[serde(crate = "near_sdk::serde")]
-pub struct ContractMetadata {
-    token_a: TokenMetadata,
-    token_b: TokenMetadata,
-    ratio: f64,
+enum FungibleTokenReceiverMessage {
+    LPDeposit,
+    Swap,
+}
+
+impl fmt::Display for FungibleTokenReceiverMessage {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            FungibleTokenReceiverMessage::LPDeposit => write!(f, "lp_deposit"),
+            FungibleTokenReceiverMessage::Swap => write!(f, "swap"),
+        }
+    }
+}
+
+#[near_bindgen]
+impl FungibleTokenReceiver for Contract {
+    fn ft_on_transfer(
+        &mut self,
+        sender_id: AccountId,
+        amount: U128,
+        msg: String,
+    ) -> PromiseOrValue<U128> {
+        let token_in = env::predecessor_account_id();
+        log!("Received token {} from {}!", token_in, sender_id);
+
+        if msg == FungibleTokenReceiverMessage::LPDeposit.to_string() {
+            assert_eq!(sender_id, self.owner_address, "{}", INVALID_LP_DEPOSIT_SENDER);
+            assert!(
+                self.tokens.iter().any(|token| token.check_address(&token_in)),
+                "{}",
+                INVALID_LP_DEPOSIT_TOKEN,
+            );
+            assert!(amount > U128(0), "{}", INVALID_LP_DEPOSIT_AMOUNT);
+            self.deposit(token_in, amount.0);
+            PromiseOrValue::Value(U128(0))
+        } else {
+            env::panic_str(INVALID_TOKEN_RECEIVER_MESSAGE);
+        }
+    }
 }

contract/src/token.rs

@@ -9,6 +9,9 @@ use near_sdk::{
 
 use crate::error::*;
 
+pub const MIN_DECIMAL: u8 = 1;
+pub const MAX_DECIMAL: u8 = 24;
+
 #[derive(BorshDeserialize, BorshSerialize)]
 pub struct Token {
     address: AccountId,
@@ -30,6 +33,11 @@ impl Token {
     }
 
     pub fn set_metadata(&mut self, metadata: FungibleTokenMetadata) {
+        assert!(
+            metadata.decimals >= MIN_DECIMAL && metadata.decimals <= MAX_DECIMAL,
+            "{}",
+            INVALID_TOKEN_DECIMAL
+        );
         log!("Set token {} name: {}, ticker: {}, decimal: {}", self.address, metadata.name, metadata.symbol, metadata.decimals);
 
         self.name = Some(metadata.name);
@@ -54,9 +62,26 @@ impl Token {
     }
 
     pub fn get_balance(&self) -> Balance {
-        assert!(self.balance > 0, "{}", TOKEN_ZERO_BALANCE);
         self.balance
     }
+
+    pub fn set_balance(&mut self, amount: Balance) {
+        self.balance = self.balance.checked_add(amount).expect(INTERNAL_OVERFLOW_ERROR);
+    }
+
+    pub fn check_address(&self, address: &AccountId) -> bool {
+        self.address.as_str() == address.as_str()
+    }
+
+    pub fn get_canonical_balance(&self) -> Balance {
+        assert!(self.decimal.is_some(), "{}", TOKEN_METADATA_NOT_INITIALISED);
+
+        let factor = 10u128
+            .checked_pow((MAX_DECIMAL - self.decimal.unwrap()) as u32)
+            .expect(INTERNAL_OVERFLOW_ERROR);
+
+        self.balance.checked_mul(factor).expect(INTERNAL_OVERFLOW_ERROR)
+    }
 }
 
 #[derive(Serialize, Deserialize)]