add: optimized get token balances (#104)

core/definitions/src/platform.ts

@@ -7,7 +7,8 @@ import { NativeAddress } from "./address";
 import { WormholeMessageId } from "./attestation";
 import { ChainContext } from "./chain";
 import { RpcConnection } from "./rpc";
-import { ChainsConfig, SignedTx, TokenId, TxHash } from "./types";
+import { AnyAddress, Balances, ChainsConfig, TokenId, TxHash } from "./types";
+import { SignedTx } from "./types";
 import { UniversalAddress } from "./universalAddress";
 
 export interface PlatformUtils<P extends PlatformName> {
@@ -21,14 +22,20 @@ export interface PlatformUtils<P extends PlatformName> {
   getDecimals(
     chain: ChainName,
     rpc: RpcConnection<P>,
-    token: NativeAddress<P> | UniversalAddress | "native",
+    token: AnyAddress,
   ): Promise<bigint>;
   getBalance(
     chain: ChainName,
     rpc: RpcConnection<P>,
     walletAddr: string,
-    token: NativeAddress<P> | UniversalAddress | "native",
+    token: AnyAddress,
   ): Promise<bigint | null>;
+  getBalances(
+    chain: ChainName,
+    rpc: RpcConnection<P>,
+    walletAddress: string,
+    tokens: AnyAddress[],
+  ): Promise<Balances>;
   getCurrentBlock(rpc: RpcConnection<P>): Promise<number>;
 
   // Platform interaction utils

core/definitions/src/protocols/cctp.ts

@@ -5,12 +5,12 @@ import {
   deserializeLayout,
   uint8ArrayToHexByteString,
 } from "@wormhole-foundation/sdk-base";
-import { ChainAddress, UniversalOrNative } from "../address";
+import { ChainAddress } from "../address";
 import { CircleMessageId } from "../attestation";
 import { universalAddressItem } from "../layout-items";
 import "../payloads/connect";
 import { RpcConnection } from "../rpc";
-import { TokenId } from "../types";
+import { AnyAddress, TokenId } from "../types";
 import { UnsignedTransaction } from "../unsignedTransaction";
 import { keccak256 } from "../utils";
 
@@ -55,7 +55,7 @@ export type CircleTransferMessage = {
 export interface AutomaticCircleBridge<P extends PlatformName> {
   transfer(
     token: ChainAddress,
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     recipient: ChainAddress,
     amount: bigint,
     nativeGas?: bigint,
@@ -66,13 +66,13 @@ export interface AutomaticCircleBridge<P extends PlatformName> {
 // https://github.com/circlefin/evm-cctp-contracts
 export interface CircleBridge<P extends PlatformName> {
   redeem(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     message: string,
     attestation: string,
   ): AsyncGenerator<UnsignedTransaction>;
   transfer(
     token: ChainAddress,
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     recipient: ChainAddress,
     amount: bigint,
   ): AsyncGenerator<UnsignedTransaction>;

core/definitions/src/protocols/core.ts

@@ -1,5 +1,5 @@
 import { PlatformName } from "@wormhole-foundation/sdk-base";
-import { UniversalOrNative } from "../address";
+import { AnyAddress } from "../types";
 import { UnsignedTransaction } from "../unsignedTransaction";
 import { RpcConnection } from "../rpc";
 
@@ -15,7 +15,7 @@ export function supportsWormholeCore<P extends PlatformName>(
 
 export interface WormholeCore<P extends PlatformName> {
   publishMessage(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     message: string | Uint8Array
   ): AsyncGenerator<UnsignedTransaction>;
   // TODO: parseTransactionDetails

core/definitions/src/protocols/ibc.ts

@@ -4,10 +4,10 @@ import {
   PlatformName,
   toChainId,
 } from "@wormhole-foundation/sdk-base";
-import { ChainAddress, NativeAddress, UniversalOrNative } from "../address";
+import { ChainAddress, NativeAddress } from "../address";
 import { IbcMessageId, WormholeMessageId } from "../attestation";
 import { RpcConnection } from "../rpc";
-import { TokenId, TxHash } from "../types";
+import { AnyAddress, TokenId, TxHash } from "../types";
 import { UnsignedTransaction } from "../unsignedTransaction";
 
 // Configuration for a transfer through the Gateway
@@ -209,9 +209,9 @@ export function supportsIbcBridge<P extends PlatformName>(
 export interface IbcBridge<P extends PlatformName> {
   //alternative naming: initiateTransfer
   transfer(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     recipient: ChainAddress,
-    token: UniversalOrNative<P> | "native",
+    token: AnyAddress,
     amount: bigint,
     payload?: Uint8Array,
   ): AsyncGenerator<UnsignedTransaction>;

core/definitions/src/protocols/tokenBridge.ts

@@ -1,6 +1,6 @@
 import { PlatformName } from "@wormhole-foundation/sdk-base";
-import { UniversalOrNative, NativeAddress, ChainAddress } from "../address";
-import { TokenId } from "../types";
+import { NativeAddress, ChainAddress } from "../address";
+import { AnyAddress, TokenId } from "../types";
 import { VAA } from "../vaa";
 import { UnsignedTransaction } from "../unsignedTransaction";
 import "../payloads/tokenBridge";
@@ -36,9 +36,9 @@ export function supportsAutomaticTokenBridge<P extends PlatformName>(
 
 export interface TokenBridge<P extends PlatformName> {
   // checks a native address to see if its a wrapped version
-  isWrappedAsset(nativeAddress: UniversalOrNative<P>): Promise<boolean>;
+  isWrappedAsset(nativeAddress: AnyAddress): Promise<boolean>;
   // returns the original asset with its foreign chain
-  getOriginalAsset(nativeAddress: UniversalOrNative<P>): Promise<TokenId>;
+  getOriginalAsset(nativeAddress: AnyAddress): Promise<TokenId>;
   // returns the wrapped version of the native asset
   getWrappedNative(): Promise<NativeAddress<P>>;
 
@@ -53,24 +53,24 @@ export interface TokenBridge<P extends PlatformName> {
   ): Promise<boolean>;
   //signer required:
   createAttestation(
-    token_to_attest: UniversalOrNative<P>,
-    payer?: UniversalOrNative<P>
+    token_to_attest: AnyAddress,
+    payer?: AnyAddress
   ): AsyncGenerator<UnsignedTransaction>;
   submitAttestation(
     vaa: VAA<"AttestMeta">,
-    payer?: UniversalOrNative<P>
+    payer?: AnyAddress
   ): AsyncGenerator<UnsignedTransaction>;
   //alternative naming: initiateTransfer
   transfer(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     recipient: ChainAddress,
-    token: UniversalOrNative<P> | "native",
+    token: AnyAddress,
     amount: bigint,
     payload?: Uint8Array
   ): AsyncGenerator<UnsignedTransaction>;
   //alternative naming: completeTransfer
   redeem(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     vaa: VAA<"Transfer"> | VAA<"TransferWithPayload">,
     unwrapNative?: boolean //default: true
   ): AsyncGenerator<UnsignedTransaction>;
@@ -79,15 +79,15 @@ export interface TokenBridge<P extends PlatformName> {
 
 export interface AutomaticTokenBridge<P extends PlatformName> {
   transfer(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     recipient: ChainAddress,
-    token: UniversalOrNative<P> | "native",
+    token: AnyAddress,
     amount: bigint,
     relayerFee: bigint,
     nativeGas?: bigint
   ): AsyncGenerator<UnsignedTransaction>;
   redeem(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     vaa: VAA<"TransferWithPayload">
   ): AsyncGenerator<UnsignedTransaction>;
   getRelayerFee(

core/definitions/src/relayer.ts

@@ -1,19 +1,19 @@
 import { Chain, PlatformName } from "@wormhole-foundation/sdk-base";
-import { UniversalOrNative } from "./address";
+import { AnyAddress } from "./types";
 
 export interface Relayer<P extends PlatformName> {
   relaySupported(chain: Chain): boolean;
   getRelayerFee(
     sourceChain: Chain,
     destChain: Chain,
-    tokenId: UniversalOrNative<P>,
+    tokenId: AnyAddress,
   ): Promise<bigint>;
   // TODO: What should this be named?
   // I don't think it should return an UnisgnedTransaction
   // rather it should take some signing callbacks and
   // a ref to track the progress
   startTransferWithRelay(
-    token: UniversalOrNative<P> | "native",
+    token: AnyAddress,
     amount: bigint,
     toNativeToken: string,
     sendingChain: Chain,
@@ -24,13 +24,13 @@ export interface Relayer<P extends PlatformName> {
   ): Promise<any>;
   calculateNativeTokenAmt(
     destChain: Chain,
-    tokenId: UniversalOrNative<P>,
+    tokenId: AnyAddress,
     amount: bigint,
     walletAddress: string,
   ): Promise<bigint>;
   calculateMaxSwapAmount(
     destChain: Chain,
-    tokenId: UniversalOrNative<P>,
+    tokenId: AnyAddress,
     walletAddress: string,
   ): Promise<bigint>;
 }

core/definitions/src/testing/mocks/platform.ts

@@ -19,6 +19,8 @@ import {
   nativeIsRegistered,
   NativeAddress,
   UniversalAddress,
+  AnyAddress,
+  Balances,
 } from "../..";
 import { MockRpc } from "./rpc";
 import { MockChain } from "./chain";
@@ -84,6 +86,15 @@ export class MockPlatform<P extends PlatformName> implements Platform<P> {
     throw new Error("Method not implemented.");
   }
 
+  getBalances(
+    chain: ChainName,
+    rpc: RpcConnection<PlatformName>,
+    walletAddress: string,
+    tokens: AnyAddress[],
+  ): Promise<Balances> {
+    throw new Error("method not implemented");
+  }
+
   getChain(chain: ChainName): ChainContext<P> {
     if (chain in this.conf) return new MockChain<P>(this.conf[chain]!);
     throw new Error("No configuration available for chain: " + chain);

core/definitions/src/testing/mocks/tokenBridge.ts

@@ -1,11 +1,10 @@
 import { PlatformName } from "@wormhole-foundation/sdk-base";
 import {
+  AnyAddress,
   ChainAddress,
   NativeAddress,
   RpcConnection,
   TokenBridge,
-  TokenId,
-  UniversalOrNative,
   UnsignedTransaction,
   VAA,
 } from "../..";
@@ -19,10 +18,10 @@ import {
 export class MockTokenBridge<P extends PlatformName> implements TokenBridge<P> {
   constructor(readonly rpc: RpcConnection<P>) {}
 
-  isWrappedAsset(token: UniversalOrNative<P>): Promise<boolean> {
+  isWrappedAsset(token: AnyAddress): Promise<boolean> {
     throw new Error("Method not implemented.");
   }
-  getOriginalAsset(token: UniversalOrNative<P>): Promise<ChainAddress> {
+  getOriginalAsset(token: AnyAddress): Promise<ChainAddress> {
     throw new Error("Method not implemented.");
   }
   hasWrappedAsset(original: ChainAddress): Promise<boolean> {
@@ -36,9 +35,7 @@ export class MockTokenBridge<P extends PlatformName> implements TokenBridge<P> {
   ): Promise<boolean> {
     throw new Error("Method not implemented.");
   }
-  createAttestation(
-    address: UniversalOrNative<P>,
-  ): AsyncGenerator<UnsignedTransaction> {
+  createAttestation(address: AnyAddress): AsyncGenerator<UnsignedTransaction> {
     throw new Error("Method not implemented.");
   }
   submitAttestation(
@@ -47,16 +44,16 @@ export class MockTokenBridge<P extends PlatformName> implements TokenBridge<P> {
     throw new Error("Method not implemented.");
   }
   transfer(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     recipient: ChainAddress,
-    token: "native" | UniversalOrNative<P>,
+    token: "native" | AnyAddress,
     amount: bigint,
     payload?: Uint8Array | undefined,
   ): AsyncGenerator<UnsignedTransaction> {
     throw new Error("Method not implemented.");
   }
   redeem(
-    sender: UniversalOrNative<P>,
+    sender: AnyAddress,
     vaa: VAA<"Transfer"> | VAA<"TransferWithPayload">,
     unwrapNative?: boolean | undefined,
   ): AsyncGenerator<UnsignedTransaction> {

core/definitions/src/types.ts

@@ -15,6 +15,14 @@ export type SequenceId = bigint;
 
 export type SignedTx = any;
 
+export type AnyAddress =
+  | NativeAddress<PlatformName>
+  | UniversalAddress
+  | string
+  | number
+  | Uint8Array
+  | number[];
+
 export type TokenId = ChainAddress;
 export function isTokenId(thing: TokenId | any): thing is TokenId {
   return (
@@ -23,6 +31,10 @@ export function isTokenId(thing: TokenId | any): thing is TokenId {
   );
 }
 
+export type Balances = {
+  [key: string]: BigInt | null;
+};
+
 export interface Signer {
   chain(): ChainName;
   address(): string;

core/definitions/src/universalAddress.ts

@@ -8,6 +8,7 @@ import { Address, NativeAddress, toNative } from "./address";
 
 export class UniversalAddress implements Address {
   static readonly byteSize = 32;
+  private readonly type = "Universal";
 
   private readonly address: Uint8Array;
 
@@ -45,7 +46,7 @@ export class UniversalAddress implements Address {
   }
 
   equals(other: UniversalAddress): boolean {
-    if (other instanceof UniversalAddress) {
+    if (UniversalAddress.instanceof(other)) {
       return other.toString() === this.toString();
     }
     return false;
@@ -54,4 +55,8 @@ export class UniversalAddress implements Address {
   static isValidAddress(address: string) {
     return isHexByteString(address, UniversalAddress.byteSize);
   }
+
+  static instanceof(address: any) {
+    return address.type === "Universal";
+  }
 }

platforms/cosmwasm/__tests__/unit/platform.test.ts

@@ -2,8 +2,6 @@ import { expect, test } from "@jest/globals";
 import {
   chains,
   chainConfigs,
-  DEFAULT_NETWORK,
-  chainToPlatform,
 } from "@wormhole-foundation/connect-sdk";
 import { CosmwasmPlatform } from "../../src/platform";