HyVM example: fetch multiple balances on-chain, in one call

main.ts

import * as ethers from 'ethers';

export const ETH_ADDRESS: HexString = '0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee';

setImmediate(async () => {
    const provider = ethers.getDefaultProvider();

    const who: HexString  = '0x945f803f01F443616546d1F31466c0E7ACfF36f7';
    
    // fetch how much is owned by the the above address, of the below tokens
    const balances = await fetchMultipleBalances(provider, who, [
        ETH_ADDRESS,
        '0xdAC17F958D2ee523a2206206994597C13D831ec7', // USDT
        '0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48', // USDC
        // add addresses here !
    ]);

    // print result
    console.log(balances);
});


// ======================= IMPLEMENTATION

export type HexString = `0x${string}`;
async function fetchMultipleBalances(
    provider: ethers.providers.Provider,
    who: HexString,
    forTokens: HexString[],
): Promise<[HexString, ethers.BigNumber][]> {
    const bytecode = fetchBalancesBytecode(who, forTokens);


    const result = await provider.call({
        to: '0xdb4516887ccd9a593e390f6a43e34494a524a551', // HyVM
        data: bytecode,
    });

    const [decoded] = new ethers.utils.AbiCoder().decode(['uint256[]'], result) as [ethers.BigNumber[]];
    if (decoded?.length !== forTokens.length) {
        throw new Error('Should not happen: HyVM call did not return the right balances length');
    }
    return forTokens.map((t, i) => [t, decoded[i]]);
}


function fetchBalancesBytecode(userAddress: HexString, tokens: HexString[]) {

    const ops: string[] = [];
    const memStart = 0x40;

    // store header in memory, so this respects return data encoding of solidity
    ops.push(
        // store array location
        ...push(0x20),
        ...push(memStart),
        mstore,
        // push array len
        ...push(tokens.length),
        ...push(memStart + 0x20),
        mstore,
    );

    // init mem size at 2 words (which have been written by the header above)
    let memSize = 0x40;

    // fetch all balances
    for (const token of tokens) {
        // this will be the place to store our balance
        const whereToStore = memStart + memSize;

        if (token === ETH_ADDRESS) {
            // get ETH balance
            ops.push(
                // --- push whose balance to get
                ...pushAddress(userAddress),
                // --- balance
                balance,
                // --- store it
                ...push(whereToStore),
                mstore,
            );
        } else {
            // ERC20
            ops.push(
                // ===== write args to memory
                // --- store balanceOf() selector
                ...pushSig('70a08231'),
                ...push(0), // push1 0 (=where to store)
                mstore,
                // --- store whose balance to get at 0x4 (after sig)
                ...pushAddress(userAddress),
                ...push(0x04), // push1 4 (=where to store)
                mstore,

                // ===== call
                // -- retSize
                ...push(0x20),
                // -- retOffset => where we will store the op result
                ...push(whereToStore),
                // -- argSize
                ...push(0x24), // args are 0x24 bytes length (selector + address)
                // -- argOffset
                ...push(0),
                // -- contract to call
                ...pushAddress(token),
                // -- gas
                gas, // gas
                // => staticcall !
                staticcall,

                // ===== pop result (ignore failures)
                pop,
            );
        }
        memSize += 0x20;
    }

    // return result
    ops.push(
        // push return size
        ...push(memSize),
        ...push(memStart),
        // return
        'F3',
    );
    return '0x' + ops.join('');
}

const mstore = '52';
const balance = '31';
const gas = '5A';
const staticcall = 'FA';
const pop = '50';

function push(num: number) {
    const str = num.toString(16);
    const nBytes = Math.ceil(str.length / 2);
    if (num < 0 || nBytes > 32) {
        throw new Error('invalid number');
    }
    return [(0x60 + nBytes - 1).toString(16), str.padStart(nBytes * 2, '0')];
}

function pushAddress(_address: HexString) {
    const address = _address.substring('0x'.length);
    if (address.length !== 40) {
        throw new Error('Invalid address');
    }
    return [
        '73', // push20
        address, // push20 address
    ];
}

function pushSig(sig: string) {
    if (!/^[\da-fA-F]{8}$/.test(sig)) {
        throw new Error('Invalid signature');
    }
    return [
        '7F', // push32
        sig + '00000000000000000000000000000000000000000000000000000000',
    ];
}