# The state root transition proof circuit

// Derive the coordinator's public key from
// their private key 
var derivedCoordinatorPubKey = eddsaDerivePubKey(coordinatorPrivKey)

// Ensure via a constraint that it matches the
// coordinator's public key given as an input
assert(derivedCoordinatorPubKey == coordinatorPubKey)

// Generate the ECDH key
var ecdhSharedKey = genEcdhKey(coordinatorPrivKey, encPubKey)

// Use the ECDH shared key to decrypt the message
var decryptedCommand = decrypt(ecdhSharedKey, message)

// Ensure via a constraint that the message has been correctly decrypted
assert(decryptedCommand == command)

var generatedMsgTreeRoot = generateMerkleRoot(
    msgTreePathElements,
    msgTreePathIndex,
    message
)

assert(generatedMsgTreeRoot, msgTreeRoot)

/***********************************
This function generates a state leaf
***********************************/
function generateStateLeaf(
    command,
    computedNewVoteOptionRoot,
    newVoiceCreditBalance
) => {

    return [
        command.newPublicKeyX,
        command.newPublicKeyY,
        computedNewVoteOptionRoot,
        newVoiceCreditBalance,
        command.nonce + 1
    ]
}

/*************************
// The main circuit logic:
**************************/

// Record in a variable that the new state leaf index is
// valid (i.e. it is leq to the maximum allowed value)
var validStateLeafIndex = newStateTreePathIndex <= maxStateLeafIndex && newStateTreePathIndex > 0

// Record in a variable if the signature is valid
var validSignature = verifyEddsa(signature, decryptedCommand, userPubKey)

// Record in a variable if the nonce is correct
var correctNonce = decryptedCommand.nonce == userCurrentLeaf.nonce + 1

// Prove that the user's current leaf is part of the
// Merkle tree. Note that this check is independent of
// the noOp flag. As such, the coordinator cannot
// create an invalid proof by tampering with the
// Merkle proof and setting noOp to true; the Merkle proof
// *must* be valid.
var x = generateMerkleRoot(
    userCurrentLeafPathElements,
    userCurrentLeafPathIndex,
    userCurrentLeaf
)

assert(x == currentStateRoot)

// Prove that the the current vote option weight (leaf) 
// input is correct by checking that it exists in the
// tree at the given index
var y = generateMerkleRoot(
    decryptedCommand.currentVoteOptionPathElements,
    decryptedCommand.voteOptionIndex,
    decryptedCommand.currentVoteWeight
)

assert(y == userCurrentLeaf.voteOptionTreeRoot)

// Record in a variable if the user has enough
// voice credits

var newVoiceCreditBalance = 
    userCurrentLeaf.voiceCreditBalance + 
    (decryptedCommand.currentVoteWeight ^ 2) -
    (decryptedCommand.newVoteWeight ^ 2)
    
var enoughVoiceCredits = newVoiceCreditBalance >= 0

// Record in a variable if the new leaf's vote option
// tree root is the correct result of updating the
// vote option leaf.
var computedNewVoteOptionRoot = 
    updateMerkleTree(
        command.voteOptionIndex,
        command.voteOptionPath,
        command.voteWeight,
        decryptedCommand.newVoteWeight
    )

assert(newStateLeaf.voteOptionRoot == computedNewVoteOptionRoot)
    newStateLeaf.voteOptionRoot == computedNewVoteOptionRoot
  
// Record in a variable if the vote option index is
// within a permissible range (0 to 2 ** vote option tree 
// depth, inclusive) where VOTE_OPTION_TREE_DEPTH is
// not an input, but rather hardcoded during the trusted 
// setup.
var validVoteOptionTreeIndex = command.voteOptionIndex < VOTE_OPTION_TREE_DEPTH

var newStateLeaf = generateStateLeaf(
    command,
    computedNewVoteOptionRoot,
    newVoiceCreditBalance
)   

if (enoughVoiceCredits &&
    correctNonce &&
    validSignature &&
    validStateLeafIndex &&
    validVoteOptionTreeIndex
):
    // Use a constraint to ensure that the no-op flag
    // is set to false
    assert(noOp == false)
    

    // Generate the new state root.
    var s = merkleTreeUpdate(
        newStateTreePathIndex
        newStateLeaf,
        currentStateRoot,
        newStateTreePathElements
    )

    // Update the leaf at index `0` to generate a new state 
    // root, and ensure via a constraint that it is equal 
    // to the new state root passed to the snark as an
    // input.

    var updatedStateRoot = merkleTreeUpdate(
        0,
        randomLeaf,
        s,
        newStateTreePathElementsToZero
    )

    assert(updatedStateRoot == newStateRoot)

else:
    // Use a constraint to ensure that the no-op flag
    // is set to true
    assert(noOp == true)