Building zkRollup (Part — 4)

This article is in continuation of Part-1, Part -2 and Part-3. In this part, I will be focusing on the on-chain contract code of ZK-rollup.

On-chain code consists of primarily 3 parts:

1. deposit()
2. forgeBatch()
3. withdraw()

Deposit

This is straight forward method; used to deposit the ETH against the EdDSA public key X and Y components.

ForgeBatch

This method in-house major functionality of zk-rollup. It has 3 major steps to perform before it can apply the state change to layer -1. Refer to the swim lane diagram in Part-2 of this series.

It first compares the on-chain state values against the public inputs of the circom circuit, if those values are found to be ok, then it moves on to verify the proof against the on-chain verifier contract.

We have built our on-chain verifier contract as per the steps given in part-1. We can refer here for more details on these steps.

After the successful verification of the SNARK proof with the public inputs, this method performs the on-chain state change. However in this state change process, if the EdDSA To and EdDSA From addresses are the same, it is considered to be a withdrawal request rather than a simple ETH transfer, and a new state is added to the withdrawal struct.

// SPDX-License-Identifier: GPL-3.0
pragma solidity ^0.8.0;
import "./PlonkVerifier.sol";
import "./PlonkWithdrawVerifier.sol";
// import "https://github.com/OpenZeppelin/openzeppelin-contracts/blob/master/contracts/security/ReentrancyGuard.sol";
import "./ReentrancyGuard.sol" ;

contract xponentDB is ReentrancyGuard {

    struct accountTrie {
        uint256 transaction_nonce;
        uint256 balance;
    }

    uint256 public transactionRoot;
    uint256 public transactionCount;
    uint256 public withdrawRoot;
    address verifierAddress ;
    address withdrawVerifierAddress ;
    mapping(uint256 => mapping(uint256 => accountTrie)) public accounts;
    mapping(uint256 => mapping(uint256 => uint256)) public withdraws;

    constructor(address _verifierAddress, address _withdrawVerifierAddress) {
        transactionRoot = 16676031662730057682516412636081941128209113782684013852209561292624791417706;
        withdrawRoot = 16676031662730057682516412636081941128209113782684013852209561292624791417706 ;
        verifierAddress = _verifierAddress ;
        withdrawVerifierAddress = _withdrawVerifierAddress ;
        transactionCount = 2 ;
    }

    function deposit(uint256 pub_x, uint256 pub_y) public payable {
        require(msg.value > 0, "BAD_AMOUNT");
        accounts[pub_x][pub_y].balance += msg.value;
    }

    /**
    * @dev user can directly withdraw from this method
    */
    function withdraw(uint[] calldata input, bytes calldata proof) external nonReentrant  {
        
        PlonkWithdrawVerifier withdrawVerifier = PlonkWithdrawVerifier(withdrawVerifierAddress);
        
        require( withdrawRoot == input[1], "ROOT_MISMATCHED");
        require( withdraws[input[3]][input[4]] >= input[5], "INSUFFICIENT_BALANCE");
        require(withdrawVerifier.verifyProof(proof, input),"INVALID_SNARK_PROOF");

        //State change begins
        withdrawRoot = input[2] ;
        withdraws[input[3]][input[4]] -= input[5] ;
        (bool success, ) = address(uint160((input[6]))).call{ value: input[5] }("");
        require(success, "SEND_FAILED");
    }

    /**
    * @dev coordinator will forge the batch here, 
    */
    function forgeBatch(
        uint[] calldata input,
        bytes calldata proof
    ) public payable {
        uint i =  0 ;
        uint j =  0 ;
        PlonkVerifier snarkVerifier = PlonkVerifier(verifierAddress);
        require( transactionRoot == input[1], "ROOT_MISMATCHED"); 

        for (uint t = 0 ; t < transactionCount ; t++){
            require( accounts[input[i+3]][input[i+4]].balance >= input[j+15] , "INSUFFICIENT_BALANCE");    
            i += 2 ; 
            j += 1 ; 
        }
        i = 0 ;
        j = 0 ;
        
        require(snarkVerifier.verifyProof(proof, input),"INVALID_SNARK_PROOF");

        // State change begins
        transactionRoot = input[2] ; // new root
        for (uint t = 0 ; t < transactionCount ; t++){
            accounts[input[i+3]][input[i+4]].transaction_nonce += 1  ; 
            accounts[input[i+3]][input[i+4]].balance  -= input[j+15] ;
            if( input[i+3] == input[i+7] && input[i+4] == input[i+8] ) {
                withdraws[input[i+3]][input[i+4]] += input[j+15] ;
            }else{
                accounts[input[i+7]][input[i+8]].balance  += input[j+15] ;
                accounts[input[i+7]][input[i+8]].transaction_nonce += 1  ;
            }
            i += 2 ; 
            j += 1 ; 
        }
    }
}

This sample code doesn't have any fixed co-ordinator public key and no provision of gas-related economies. Both these topics are for the later part of the discussion

Withdraw

This method is used to transfer the ETH out of the roll-up contract. Similar to ForgeBatch this method also needs to first check the Circom public input signals against the state values and if found ok then it verifies the proof again the withdraw verifier on-chain contract. After passing the proof it transfers the ETH to the “to” address present in the public signal of the circuit.

Withdraw is part of the complex process of any zk-rollup. I have tried to over-simplify this whole function just to make it work in this sample code. However, in the actual production-grade code there are several checks and balances that need to be put in place.

To know more about the withdrawal process please refer part-3

And that's it, we are done for now. I have presented you with a working prototype code for a zk-rollup off-chain and on-chain code. There are many other important aspects that I have left in this prototype e.g gas economics and revenue mode for the coordinator etc. I will try to cover some of these aspects in my upcoming stories.

Next Steps

This zk-rollup focuses to solve just the scalability part of Ethereum (layes1). In my future stories, I will be incorporating some of the major Defi use cases like DEX into zk-rollup.