Add individual balances and allowances to OrderRelevantState

3 files changed, 74 insertions, 3 deletions

diff --git a/packages/order-utils/src/asset_data_utils.ts b/packages/order-utils/src/asset_data_utils.ts
index 2adaf31b8..18c1f93b6 100644
--- a/packages/order-utils/src/asset_data_utils.ts
+++ b/packages/order-utils/src/asset_data_utils.ts
@@ -1,4 +1,4 @@
-import { AssetData, AssetProxyId, ERC20AssetData, ERC721AssetData, MultiAssetData } from '@0x/types';
+import { AssetProxyId, ERC20AssetData, ERC721AssetData, MultiAssetData, SingleAssetData } from '@0x/types';
 import { AbiEncoder, BigNumber } from '@0x/utils';
 import { MethodAbi } from 'ethereum-types';
 import * as _ from 'lodash';
@@ -225,7 +225,7 @@ export const assetDataUtils = {
      * @param assetData Hex encoded assetData string to decode
      * @return Either a ERC20 or ERC721 assetData object
      */
-    decodeAssetDataOrThrow(assetData: string): AssetData {
+    decodeAssetDataOrThrow(assetData: string): SingleAssetData | MultiAssetData {
         const assetProxyId = assetDataUtils.decodeAssetProxyId(assetData);
         switch (assetProxyId) {
             case AssetProxyId.ERC20:
diff --git a/packages/order-utils/src/index.ts b/packages/order-utils/src/index.ts
index aae439e68..e098f23c3 100644
--- a/packages/order-utils/src/index.ts
+++ b/packages/order-utils/src/index.ts
@@ -34,7 +34,7 @@ export {
     OrderRelevantState,
     OrderState,
     ECSignature,
-    AssetData,
+    SingleAssetData,
     ERC20AssetData,
     ERC721AssetData,
     MultiAssetData,
diff --git a/packages/order-utils/src/order_state_utils.ts b/packages/order-utils/src/order_state_utils.ts
index fe0d6c773..9e3e228ba 100644
--- a/packages/order-utils/src/order_state_utils.ts
+++ b/packages/order-utils/src/order_state_utils.ts
@@ -1,15 +1,21 @@
 import {
+    AssetProxyId,
     ExchangeContractErrs,
+    MultiAssetData,
+    ObjectMap,
     OrderRelevantState,
     OrderState,
     OrderStateInvalid,
     OrderStateValid,
     SignedOrder,
+    SingleAssetData,
 } from '@0x/types';
 import { BigNumber } from '@0x/utils';
+import * as _ from 'lodash';
 
 import { AbstractBalanceAndProxyAllowanceFetcher } from './abstract/abstract_balance_and_proxy_allowance_fetcher';
 import { AbstractOrderFilledCancelledFetcher } from './abstract/abstract_order_filled_cancelled_fetcher';
+import { assetDataUtils } from './asset_data_utils';
 import { orderHashUtils } from './order_hash';
 import { OrderValidationUtils } from './order_validation_utils';
 import { RemainingFillableCalculator } from './remaining_fillable_calculator';
@@ -18,7 +24,9 @@ import { utils } from './utils';
 interface SidedOrderRelevantState {
     isMakerSide: boolean;
     traderBalance: BigNumber;
+    traderIndividualBalances: ObjectMap<BigNumber>;
     traderProxyAllowance: BigNumber;
+    traderIndividualProxyAllowances: ObjectMap<BigNumber>;
     traderFeeBalance: BigNumber;
     traderFeeProxyAllowance: BigNumber;
     filledTakerAssetAmount: BigNumber;
@@ -121,7 +129,9 @@ export class OrderStateUtils {
         const sidedOrderRelevantState = {
             isMakerSide: true,
             traderBalance: orderRelevantState.makerBalance,
+            traderIndividualBalances: orderRelevantState.makerIndividualBalances,
             traderProxyAllowance: orderRelevantState.makerProxyAllowance,
+            traderIndividualProxyAllowances: orderRelevantState.makerIndividualProxyAllowances,
             traderFeeBalance: orderRelevantState.makerFeeBalance,
             traderFeeProxyAllowance: orderRelevantState.makerFeeProxyAllowance,
             filledTakerAssetAmount: orderRelevantState.filledTakerAssetAmount,
@@ -165,7 +175,9 @@ export class OrderStateUtils {
 
         const orderRelevantState = {
             makerBalance: sidedOrderRelevantState.traderBalance,
+            makerIndividualBalances: sidedOrderRelevantState.traderIndividualBalances,
             makerProxyAllowance: sidedOrderRelevantState.traderProxyAllowance,
+            makerIndividualProxyAllowances: sidedOrderRelevantState.traderIndividualProxyAllowances,
             makerFeeBalance: sidedOrderRelevantState.traderFeeBalance,
             makerFeeProxyAllowance: sidedOrderRelevantState.traderFeeProxyAllowance,
             filledTakerAssetAmount: sidedOrderRelevantState.filledTakerAssetAmount,
@@ -236,10 +248,12 @@ export class OrderStateUtils {
         const isAssetZRX = assetData === zrxAssetData;
 
         const traderBalance = await this._balanceAndProxyAllowanceFetcher.getBalanceAsync(assetData, traderAddress);
+        const traderIndividualBalances = await this._getAssetBalancesAsync(assetData, traderAddress);
         const traderProxyAllowance = await this._balanceAndProxyAllowanceFetcher.getProxyAllowanceAsync(
             assetData,
             traderAddress,
         );
+        const traderIndividualProxyAllowances = await this._getAssetProxyAllowancesAsync(assetData, traderAddress);
         const traderFeeBalance = await this._balanceAndProxyAllowanceFetcher.getBalanceAsync(
             zrxAssetData,
             traderAddress,
@@ -278,7 +292,9 @@ export class OrderStateUtils {
         const sidedOrderRelevantState = {
             isMakerSide,
             traderBalance,
+            traderIndividualBalances,
             traderProxyAllowance,
+            traderIndividualProxyAllowances,
             traderFeeBalance,
             traderFeeProxyAllowance,
             filledTakerAssetAmount,
@@ -287,4 +303,59 @@ export class OrderStateUtils {
         };
         return sidedOrderRelevantState;
     }
+    private async _getAssetBalancesAsync(
+        assetData: string,
+        traderAddress: string,
+        initialBalances: ObjectMap<BigNumber> = {},
+    ): Promise<ObjectMap<BigNumber>> {
+        const decodedAssetData = assetDataUtils.decodeAssetDataOrThrow(assetData);
+        let balances: ObjectMap<BigNumber> = { ...initialBalances };
+        switch (decodedAssetData.assetProxyId) {
+            case AssetProxyId.ERC20:
+            case AssetProxyId.ERC721:
+                const balance = await this._balanceAndProxyAllowanceFetcher.getBalanceAsync(assetData, traderAddress);
+                const tokenAddress = (decodedAssetData as SingleAssetData).tokenAddress;
+                balances[tokenAddress] = _.isUndefined(initialBalances[tokenAddress])
+                    ? balance
+                    : balances[tokenAddress].add(balance);
+                break;
+            case AssetProxyId.MultiAsset:
+                for (const assetDataElement of (decodedAssetData as MultiAssetData).nestedAssetData) {
+                    balances = await this._getAssetBalancesAsync(assetDataElement, traderAddress, balances);
+                }
+                break;
+            default:
+                throw new Error(`Proxy with id ${decodedAssetData.assetProxyId} not supported`);
+        }
+        return balances;
+    }
+    private async _getAssetProxyAllowancesAsync(
+        assetData: string,
+        traderAddress: string,
+        initialAllowances: ObjectMap<BigNumber> = {},
+    ): Promise<ObjectMap<BigNumber>> {
+        const decodedAssetData = assetDataUtils.decodeAssetDataOrThrow(assetData);
+        let allowances: ObjectMap<BigNumber> = { ...initialAllowances };
+        switch (decodedAssetData.assetProxyId) {
+            case AssetProxyId.ERC20:
+            case AssetProxyId.ERC721:
+                const allowance = await this._balanceAndProxyAllowanceFetcher.getProxyAllowanceAsync(
+                    assetData,
+                    traderAddress,
+                );
+                const tokenAddress = (decodedAssetData as SingleAssetData).tokenAddress;
+                allowances[tokenAddress] = _.isUndefined(initialAllowances[tokenAddress])
+                    ? allowance
+                    : allowances[tokenAddress].add(allowance);
+                break;
+            case AssetProxyId.MultiAsset:
+                for (const assetDataElement of (decodedAssetData as MultiAssetData).nestedAssetData) {
+                    allowances = await this._getAssetBalancesAsync(assetDataElement, traderAddress, allowances);
+                }
+                break;
+            default:
+                throw new Error(`Proxy with id ${decodedAssetData.assetProxyId} not supported`);
+        }
+        return allowances;
+    }
 }