Loading of token balances

src/main/swap/Main.ts

@@ -17,6 +17,8 @@ import { TokenState as TokenStateV1, deserializeTokenState as deserializeTokenSt
 import { TokenState as TokenStateV2, deserializeTokenState as deserializeTokenStateV2 } from "../abi/TokenV2";
 import { NETWORK } from "../constant";
 import { AvlClient } from "../client/AvlClient";
+import {TokenV1Contract  } from "../token/contract/TokenV1Contract";
+import {TokenV2Contract  } from "../token/contract/TokenV2Contract";
 
 // UI constants
 const TOKEN_LIST = <Element>document.querySelector("#token-list");
@@ -28,6 +30,7 @@ const EXCHANGE_RATE_LIST = <Element>document.querySelector("#rate-list");
 type ContractType = "token_v1" | "token_v2" | "swap_lock_v1"
 
 interface ContractState {
+    address: BlockchainAddress,
     latest_state: TokenStateV1 | TokenStateV2 | null;
     swaps: BlockchainAddress[] | null,
     type: ContractType | null,
@@ -37,6 +40,7 @@ interface Route {
     from: BlockchainAddress | null,
     to: BlockchainAddress | null,
     swaps: BlockchainAddress[],
+    user: BlockchainAddress | null,
 }
 
 const ROUTERS: BlockchainAddress[] = [BlockchainAddress.fromString("02f8eb18e09dfe6797880c952527747202560338bf")];
@@ -47,7 +51,7 @@ const TOKENS: Record<BlockchainAddressKey, ContractState> = {};
 
 const SWAPS: Record<BlockchainAddressKey, ContractState> = {};
 
-const CURRENT_ROUTE: Route = { from: null, to: null, swaps: [] };
+const CURRENT_ROUTE: Route = { from: null, to: null, swaps: [], user: null };
 
 const SHARDED_CLIENT: ShardedClient = new ShardedClient(NETWORK.node_base_url, NETWORK.network_shards);
 const AVL_CLIENT:AvlClient  = new AvlClient(NETWORK .node_base_url, NETWORK.network_shards);;
@@ -73,6 +77,13 @@ function address_to_url(address: BlockchainAddress): string {
     return `${NETWORK.browser_base_url}/contracts/${address.asString()}`;
 }
 
+function setInnerText(query: string, text: string) {
+    console.log(query, text);
+    for (const a of document.querySelectorAll(query)) {
+        (<HTMLElement>a).innerText = text;
+    }
+}
+
 function ui_show_route() {
     if (CURRENT_ROUTE.from == null || CURRENT_ROUTE.to == null) {
         return;
@@ -84,12 +95,8 @@ function ui_show_route() {
     const symbolFrom = infoFrom.latest_state == null ? "???" : infoFrom.latest_state.symbol;
     const symbolTo = infoTo.latest_state == null ? "???" : infoTo.latest_state.symbol;
 
-    for (const a of document.querySelectorAll("[data-input-token]")) {
+    setInnerText("[data-input-token]", symbolFrom);
-        (<HTMLElement>a).innerText = symbolFrom;
+    setInnerText("[data-output-token]", symbolTo);
-    }
-    for (const a of document.querySelectorAll("[data-output-token]")) {
-        (<HTMLElement>a).innerText = symbolTo;
-    }
 }
 
 function reroute() {
@@ -142,6 +149,7 @@ function ui_add_token(tokenState: TokenStateV1 | TokenStateV2, address: Blockcha
     spanSymbol.title = tokenState.name;
     spanSymbol.href = address_to_url(address);
     const spanAmount = document.createElement("div");
+    spanAmount.setAttribute("data-user-token-amount", address.asString());
     spanAmount.innerText = "Login to view";
     TOKEN_LIST.append(spanFrom);
     TOKEN_LIST.append(spanTo);
@@ -149,12 +157,12 @@ function ui_add_token(tokenState: TokenStateV1 | TokenStateV2, address: Blockcha
     TOKEN_LIST.append(spanAmount);
 }
 
-async function get_current_liquidity(swapAddress: BlockchainAddress): Promise<TokenBalance> {
+async function get_swap_deposit_balances(swapAddress: BlockchainAddress, account: BlockchainAddress): Promise<TokenBalance> {
     const SWAP_CONTRACT_BALANCES_TREE_ID = 0;
     const dataBuffer = await AVL_CLIENT.getContractStateAvlValue(
       swapAddress.asString(),
       SWAP_CONTRACT_BALANCES_TREE_ID,
-      swapAddress.asBuffer(),
+      account.asBuffer(),
     );
 
     if (dataBuffer  === undefined) {
@@ -170,6 +178,10 @@ async function get_current_liquidity(swapAddress: BlockchainAddress): Promise<To
     }
 }
 
+function get_current_liquidity(swapAddress: BlockchainAddress): Promise<TokenBalance> {
+    return get_swap_deposit_balances(swapAddress, swapAddress);
+}
+
 const RATE_DECIMALS = new BN(10000);
 
 async function get_current_exchange_rate(swapAddress: BlockchainAddress): Promise<BN> {
@@ -209,19 +221,45 @@ function update_swap_contract_info(){
     }
 }
 
-async function get_token_state(tokenAddress: BlockchainAddress): Promise<TokenStateV1 | TokenStateV2> {
+async function get_token_state(tokenAddress: BlockchainAddress): Promise<[TokenStateV1 | TokenStateV2, ContractType]> {
     try {
-        return await get_contract_state(tokenAddress, deserializeTokenStateV1);
+        return [await get_contract_state(tokenAddress, deserializeTokenStateV1), "token_v1"];
     } catch {
         // Pass
     }
 
-    return await get_contract_state(tokenAddress, deserializeTokenStateV2);
+    return [await get_contract_state(tokenAddress, deserializeTokenStateV2), "token_v2"];
+}
+
+async function get_token_balance(info: ContractState, account: BlockchainAddress): Promise<BN> {
+    console.log(info);
+    if (info.type == "token_v2") {
+        return await new TokenV2Contract(SHARDED_CLIENT, undefined).tokenBalance(info.address, account);
+    } else {
+        return await new TokenV1Contract(SHARDED_CLIENT, undefined).tokenBalance(info.address, account);
+    }
+}
+
+
+function onLogin(userAddress: BlockchainAddress | null) {
+    if (userAddress == null) {
+        return; // TODO;
+    }
+    CURRENT_ROUTE.user = userAddress;
+
+    // TODO: Move into own function
+    for (const key in TOKENS) {
+        const info = TOKENS[key];
+         setInnerText(`[data-user-token-amount="${key}"]`, "Loading...");
+        get_token_balance(info, userAddress).catch(e => 0).then(balance => {
+            setInnerText(`[data-user-token-amount="${key}"]`, balance.toString());
+        });
+    }
 }
 
 function setup() {
     // Setup wallet widget
-    setupWalletWidget();
+    setupWalletWidget(onLogin);
 
     // Setup routes
     for (let router of ROUTERS) {
@@ -230,18 +268,18 @@ function setup() {
             console.log(state);
 
             for (const swapInfo of state.swapContracts) {
-                TOKENS[swapInfo.tokenAAddress.asString()] = { latest_state: null, swaps: null, type: null };
+                TOKENS[swapInfo.tokenAAddress.asString()] = { latest_state: null, swaps: null, type: null, address: swapInfo.tokenAAddress };
-                TOKENS[swapInfo.tokenBAddress.asString()] = { latest_state: null, swaps: null, type: null };
+                TOKENS[swapInfo.tokenBAddress.asString()] = { latest_state: null, swaps: null, type: null, address: swapInfo.tokenBAddress };
-                SWAPS[swapInfo.swapAddress.asString()] = { latest_state: null, swaps: [swapInfo.tokenAAddress, swapInfo.tokenBAddress], type: null };
+                SWAPS[swapInfo.swapAddress.asString()] = { latest_state: null, swaps: [swapInfo.tokenAAddress, swapInfo.tokenBAddress], type: null, address: swapInfo.swapAddress };
                 ui_add_swap(swapInfo.swapAddress);
             }
 
             for (const tokenAddressStr in TOKENS) {
                 const tokenAddress = BlockchainAddress.fromString(tokenAddressStr);
-                get_token_state(tokenAddress).then(state => {
+                get_token_state(tokenAddress).then(stateAndType => {
-                    TOKENS[tokenAddress.asString()].type = "token_v1";
+                    TOKENS[tokenAddress.asString()].type = stateAndType[1];
-                    TOKENS[tokenAddress.asString()].latest_state = state;
+                    TOKENS[tokenAddress.asString()].latest_state = stateAndType[0];
-                    ui_add_token(state, tokenAddress);
+                    ui_add_token(stateAndType[0], tokenAddress);
                 });
             }
 

src/main/token/Main.ts

@@ -35,7 +35,7 @@ import { PutTransactionWasSuccessful } from "../client/TransactionData";
 import { NETWORK } from "../constant";
 
 // Setup wallet widget
-setupWalletWidget();
+setupWalletWidget(account => {});
 
 // Setup event listener that sends a transfer transaction to the contract.
 // This requires that a wallet has been connected.

src/main/token/WalletIntegration.ts

@@ -16,6 +16,7 @@
  *
  */
 
+import { BlockchainAddress } from "@partisiablockchain/abi-client";
 import { ConnectedWallet } from "../shared/ConnectedWallet";
 import { connectMetaMask } from "../shared/MetaMaskIntegration";
 import { connectMpcWallet } from "../shared/MpcWalletIntegration";
@@ -84,6 +85,7 @@ const handleWalletConnect = (connect: Promise<ConnectedWallet>) => {
       setVisibility("#ledger-connect", false);
       setVisibility("#private-key-connect", false);
       setVisibility("#wallet-disconnect", true);
+      ON_LOGIN(BlockchainAddress.fromString(userAccount.address));
     })
     .catch((error) => {
       console.error(error);
@@ -92,6 +94,7 @@ const handleWalletConnect = (connect: Promise<ConnectedWallet>) => {
       } else {
         setConnectionStatus("An error occurred trying to connect wallet: " + error);
       }
+      ON_LOGIN(null);
     });
 };
 
@@ -107,6 +110,7 @@ export const disconnectWalletClick = () => {
   setVisibility("#private-key-connect", true);
   setVisibility("#wallet-disconnect", false);
   setVisibility("#connection-link-ledger-validate", false);
+  ON_LOGIN(null);
 };
 
 function setConnectionStatus(status: string, address: string | undefined = undefined) {
@@ -130,7 +134,11 @@ const setVisibility = (selector: string, visible: boolean) => {
   }
 };
 
-export function setupWalletWidget() {
+let ON_LOGIN: (userAccount: BlockchainAddress | null) => void = (userAccount) => {};
+
+export function setupWalletWidget(onLogin: (userAccount: BlockchainAddress | null) => void) {
+    ON_LOGIN = onLogin;
+
     // Setup event listener to connect to the MPC wallet browser extension
     const connectWallet = <Element>document.querySelector("#wallet-connect-btn");
     connectWallet.addEventListener("click", connectMpcWalletClick);

src/main/token/contract/TokenV1Contract.ts

@@ -86,6 +86,20 @@ export class TokenV1Contract implements TokenContract {
     return this.getState(contractAddress);
   }
 
+  public tokenBalance(
+    contractAddress: BlockchainAddress,
+    owner: BlockchainAddress
+  ): Promise<BN> {
+      return this.tokenBalances(contractAddress, undefined).then(balances => {
+          balances.balances.forEach((amount, tokenOwner) => {
+              if (tokenOwner === owner) {
+                  return amount;
+              }
+          });
+          throw new Error("No balance for user: " + owner.asString());
+      });
+  }
+
   /*eslint @typescript-eslint/no-unused-vars: ["error", { "argsIgnorePattern": "^unused" }]*/
   public tokenBalances(
     contractAddress: BlockchainAddress,