MiMCMerkle.sol helper contract for RollupNC.sol

MiMCMerkle.sol

pragma solidity >=0.4.21;

contract MiMC {

    function MiMCpe7(uint256,uint256) public pure returns(uint256) {}

}

contract MiMCMerkle{

    MiMC public mimc;
    uint public IV = 15021630795539610737508582392395901278341266317943626182700664337106830745361;
    // hashes for empty tree of depth 16
    uint[5] public zeroCache = [
        17400342990847699622034895903486521563192531922107760411846337521891653711537, //H0 = empty leaf
        6113825327972579408082802166670133747202624653742570870320185423954556080212,  //H1 = hash(H0, H0)
        6180012883826996691682233524035352980520561433337754209809143632670877151717,  //H2 = hash(H1, H1)
        20633846227573655562891472654875498275532732787736199734105126629336915134506, //...and so on
        19963324565646943143661364524780633879811696094118783241060299022396942068715
    ];

    constructor(
        address _mimcContractAddr
    ) public {
        mimc = MiMC(_mimcContractAddr);
    }

    function getRootFromProof2(
        uint256 _leaf,
        uint256[2] memory _position,
        uint256[2] memory _proof
    ) public view returns(uint) {
        uint256[2] memory root;

        uint r = IV;

        // if leaf is left sibling
        if (_position[0] == 0){
            root[0] = mimc.MiMCpe7(mimc.MiMCpe7(r, _leaf), _proof[0]);
        }
        // if leaf is right sibling
        else if (_position[0] == 1){
            root[0] = mimc.MiMCpe7(mimc.MiMCpe7(r, _proof[0]), _leaf);
        }

        for (uint i = 1; i < _proof.length; i++){
            // if leaf is left sibling
            if (_position[i] == 0){
                root[i] = mimc.MiMCpe7(mimc.MiMCpe7(r, root[i - 1]), _proof[i]);
            }
            // if leaf is right sibling
            else if (_position[i] == 1){
                root[i] = mimc.MiMCpe7(mimc.MiMCpe7(r, _proof[i]), root[i - 1]);
            }
        }
        // return (_claimedRoot == root[root.length - 1]);
        return root[root.length - 1];
    }

    function hashBalance(uint[5] memory array) public view returns(uint){
        //[pubkey_x, pubkey_y, balance, nonce, token_type]
        uint r = IV;
        for (uint i = 0; i < array.length; i++){
            r = mimc.MiMCpe7(r, array[i]);
        }
        return r;
    }

    function hashTx(uint[8] memory array) public view returns(uint){
        //[from_x, from_y,from_index, to_x, to_y, amt, token_type]
        uint r = IV;
        for (uint i = 0; i < array.length; i++){
            r = mimc.MiMCpe7(r, array[i]);
        }
        return r;
    }

    function hashPair(uint[2] memory array) public view returns(uint){
        uint r = IV;
        for (uint i = 0; i < array.length; i++){
            r = mimc.MiMCpe7(r, array[i]);
        }
        return r;
    }

    function hashHeight2Tree(uint[4] memory array) public view returns(uint){

        uint[2] memory level1;
        for (uint i = 0; i < level1.length; i++){
            level1[i] = hashPair([array[2*i], array[2*i + 1]]);
        }
        uint level2;
        level2 = hashPair([level1[0], level1[1]]);
        return level2;
    }

}

RollupNC.sol

pragma solidity ^0.5.0;

import "../build/Update_verifier.sol";
import "../build/Withdraw_verifier.sol";

contract IMiMC {
    function MiMCpe7(uint256,uint256) public pure returns(uint256) {}
}

contract IMiMCMerkle {

    uint[16] public zeroCache;
    function getRootFromProof2(
        uint256,
        uint256[2] memory,
        uint256[2] memory
    ) public view returns(uint) {}
    function hashBalance(uint[5] memory) public view returns(uint){}
    function hashTx(uint[8] memory) public view returns(uint) {}
    function hashPair(uint[2] memory) public view returns(uint){}
    function hashHeight2Tree(uint[4] memory) public view returns(uint){}
}

contract ITokenRegistry {
    address public coordinator;
    uint256 public numTokens;
    mapping(address => bool) public pendingTokens;
    mapping(uint256 => address) public registeredTokens;
    modifier onlyCoordinator(){
        assert (msg.sender == coordinator);
        _;
    }
    function registerToken(address tokenContract) public {}
    function approveToken(address tokenContract) public onlyCoordinator{}
}

contract RollupNC is Update_verifier, Withdraw_verifier{

    IMiMC public mimc;
    IMiMCMerkle public mimcMerkle;
    ITokenRegistry public tokenRegistry;

    uint256 public currentRoot;
    address public coordinator;
    uint256[] public pendingDeposits;
    uint public queueNumber;
    uint public depositSubtreeHeight;
    uint256 public updateNumber;

    uint256 public BAL_DEPTH = 4;
    uint256 public TX_DEPTH = 2;

    // (queueNumber => [pubkey_x, pubkey_y, balance, nonce, token_type])
    mapping(uint256 => uint256) public deposits; //leaf idx => leafHash
    mapping(uint256 => uint256) public updates; //txRoot => update idx

    event RegisteredToken(uint tokenType, address tokenContract);
    event RequestDeposit(uint[2] pubkey, uint amount, uint tokenType);
    event UpdatedState(uint currentRoot, uint oldRoot, uint txRoot);
    event Withdraw(uint[3] accountInfo, address recipient, uint txRoot, uint[3] txInfo);

    constructor(
        address _mimcContractAddr,
        address _mimcMerkleContractAddr,
        address _tokenRegistryAddr
    ) public {
        mimc = IMiMC(_mimcContractAddr);
        mimcMerkle = IMiMCMerkle(_mimcMerkleContractAddr);
        tokenRegistry = ITokenRegistry(_tokenRegistryAddr);
        currentRoot = mimcMerkle.zeroCache(BAL_DEPTH);
        coordinator = msg.sender;
        queueNumber = 0;
        depositSubtreeHeight = 0;
        updateNumber = 0;
    }

    modifier onlyCoordinator(){
        assert(msg.sender == coordinator);
        _;
    }

    function updateState(
            uint[2] memory a,
            uint[2][2] memory b,
            uint[2] memory c,
            uint[3] memory input
        ) public onlyCoordinator {
        require(currentRoot == input[2], "input does not match current root");
        //validate proof
        require(update_verifyProof(a,b,c,input),
        "SNARK proof is invalid");
        // update merkle root
        currentRoot = input[0];
        updateNumber++;
        updates[input[1]] = updateNumber;
        emit UpdatedState(input[0], input[1], input[2]); //newRoot, txRoot, oldRoot
    }

    // user tries to deposit ERC20 tokens
    function deposit(
        uint[2] memory pubkey,
        uint amount,
        uint tokenType
    ) public payable {
        require(
            (amount > 0 && tokenType > 1) ||
            (msg.value > 0 && tokenType == 1) ||
            msg.sender == coordinator,
            "Deposit must be greater than 0."
        );
        require(
            tokenType == 0 ||
            tokenType == 1 ||
            tokenRegistry.registeredTokens(tokenType) != address(0),
        "tokenType is not registered.");
        uint depositHash = mimcMerkle.hashBalance(
            [pubkey[0], pubkey[1], amount, 0, tokenType]
        );
        pendingDeposits.push(depositHash);
        emit RequestDeposit(pubkey, amount, tokenType);
        queueNumber++;
        uint tmpDepositSubtreeHeight = 0;
        uint tmp = queueNumber;
        while(tmp % 2 == 0){
            pendingDeposits[pendingDeposits.length - 2] = mimcMerkle.hashPair(
                [pendingDeposits[pendingDeposits.length - 2],
                pendingDeposits[pendingDeposits.length - 1]]);
            removeDeposit(pendingDeposits.length - 1);
            tmp = tmp / 2;
            tmpDepositSubtreeHeight++;
        }
        if (tmpDepositSubtreeHeight > depositSubtreeHeight){
            depositSubtreeHeight = tmpDepositSubtreeHeight;
        }
    }



    // coordinator adds certain number of deposits to balance tree
    // coordinator must specify subtree index in the tree since the deposits
    // are being inserted at a nonzero height
    function processDeposits(
        uint[2] memory subtreePosition,
        uint[2] memory subtreeProof
    ) public onlyCoordinator returns(uint256){
        uint emptySubtreeRoot = mimcMerkle.zeroCache(2); //empty subtree of height 2
        require(currentRoot == mimcMerkle.getRootFromProof2(
            emptySubtreeRoot, subtreePosition, subtreeProof),
            "specified subtree is not empty");
        currentRoot = mimcMerkle.getRootFromProof2(
            pendingDeposits[0], subtreePosition, subtreeProof);
        removeDeposit(0);
        queueNumber = queueNumber - 2**depositSubtreeHeight;
        return currentRoot;
    }

    function withdraw(
        uint[3] memory accountInfo, //[pubkeyX, pubkeyY, index]
        uint[3] memory txInfo, //[nonce, amount, token_type_from]
        uint[2][2] memory positionAndProof, //[[position], [proof]]
        uint txRoot,
        address recipient,
        uint[2] memory a,
        uint[2][2] memory b,
        uint[2] memory c
    ) public{
        require(updates[txRoot] > 0, "txRoot must exist");
        uint txLeaf = mimcMerkle.hashTx([
            accountInfo[0], accountInfo[1], accountInfo[2],
            0, 0, //withdraw to zero address
            txInfo[0], txInfo[1], txInfo[2]
        ]);
        require(txRoot == mimcMerkle.getRootFromProof2(
            txLeaf, positionAndProof[0], positionAndProof[1]),
            "transaction does not exist in specified transactions root"
        );

        // message is hash of nonce and recipient address
        uint m = mimcMerkle.hashPair([txInfo[0], uint(recipient)]);
        require(withdraw_verifyProof(
            a, b, c, [accountInfo[0], accountInfo[1], m]),
            "eddsa signature is not valid");

        emit Withdraw(accountInfo, recipient, txRoot, txInfo);
    }

    //call methods on TokenRegistry contract

    function registerToken(
        address tokenContract
    ) public {
        tokenRegistry.registerToken(tokenContract);
    }

    function approveToken(
        address tokenContract
    ) public onlyCoordinator {
        tokenRegistry.approveToken(tokenContract);
        emit RegisteredToken(tokenRegistry.numTokens(),tokenContract);
    }

    // helper functions
    function removeDeposit(uint index) internal returns(uint[] memory) {
        require(index < pendingDeposits.length, "index is out of bounds");

        for (uint i = index; i<pendingDeposits.length-1; i++){
            pendingDeposits[i] = pendingDeposits[i+1];
        }
        delete pendingDeposits[pendingDeposits.length-1];
        pendingDeposits.length--;
        return pendingDeposits;
    }
}

TestToken.sol

pragma solidity ^0.5.0;

import "openzeppelin-solidity/contracts/token/ERC20/ERC20.sol";
import "openzeppelin-solidity/contracts/ownership/Ownable.sol";
/**
* @title TestToken is a basic ERC20 Token
*/
contract TestToken is ERC20, Ownable{

/**
    * @dev assign totalSupply to account creating this contract */
    constructor() public {

        _mint(msg.sender, 100000000000);

    }
}

TODOs

1. refactor hash functions to use a single hash with arbitrary length array input
2. integrate ERC20 token contract transfers