# Constant value for the depth at which leaves sit
const.LEAF_DEPTH=64

# SET
# =================================================================================================

#! Inserts or removes a value associated with the given key. The leaf to which we're inserting is
#! guaranteed to be empty.
#!
#! Inputs:  
#!   Operand stack: [V, K, R, ...]
#!
#! Outputs: 
#!   Operand stack: [V_old, R_new, ...]
#!
#! Cycles
#!   Insert empty value: X cycles
#!   Insert non-empty value: X cycles
proc.set_empty_leaf
    # Check if we're inserting the empty value (X cycles)
    padw eqw
    #=> [V == ZERO, ZERO, V, K, R]

    if.true
        # Inserting an empty value; this is a no-op (4 cycles)
        dropw
        #=> [V (=ZERO), K, R, ...]

        # Prepare stack: verify that the leaf is actually empty
        # (X cycles)
        movupw.2 swapw dup.8 movdn.4 push.LEAF_DEPTH movdn.4
        #=> [V (=ZERO), depth, K[3], R, K, ...]

        # (1 cycle)
        mtree_verify
        #=> [V (=ZERO), depth, K[3], R, K, ...]
        
        # Prepare stack for return (X cycles)
        movup.4 drop movup.4 drop movupw.2 dropw
        #=> [V (=ZERO), R, ...]
    else
        # Inserting a non-empty value (4 cycles)
        dropw
        #=> [V, K, R, ...]

        # Update advice map
        adv.insert_hdword
        #=> [V, K, R, ...]

        # Compute hash([K, V]); the new node value (NV)
        # (21 cycles)
        dupw.1 swapw hmerge
        # => [NV, K, R]

        # Prepare stack for `mtree_set` (5 cycles)
        movupw.2 dup.8 push.LEAF_DEPTH
        #=> [depth, K[3], R, NV, K]

        # Insert node in Merkle store (29 cycles)
        mtree_set
        #=> [V_in_leaf, R_new, K]

        # Check that V_in_leaf is indeed empty (15 cycles)
        padw assert_eqw
        #=> [R_new, K]

        # Prepare stack for return (9 cycles)
        swapw dropw padw
        #=> [ZERO, R_new]
    end
end

#! Inserts a value at the given key. The leaf to which we're inserting is
#! guaranteed to hold a single key-value pair (provided on the advice stack).
#!
#! Inputs: 
#!   Operand stack: [V, K, R, ...]
#!   Advice stack:  [K_in_leaf, V_in_leaf]
#!
#! Outputs:
#!   Operand stack: [V_old, R_new, ...]
#!
#! Cycles:
#!   Leaf single after insertion: X cycles
#!   Leaf multiple after insertion: unimplemented
proc.insert_single_leaf
    # Push the leaf pre-image on stack
    # (X cycles)
    adv_push.8
    # => [V_in_leaf, K_in_leaf, V, K, R]

    # Check if the key stored in the leaf is the same as K
    # (X cycles)
    movupw.3 movupw.2 eqw
    # => [K_in_leaf==K, K_in_leaf, K, V_in_leaf, V, R]

    if.true
        # Leaf stays a "single" variant

        # (4 cycles)
        dropw 
        # => [K, V_in_leaf, V, R]

        # Update advice map (3 cycles)
        movupw.2 adv.insert_hdword
        # => [V, K, V_in_leaf, R]

        # Compute hash([K, V]); the new node value (NV)
        # (X cycles)
        dupw.1 swapw hmerge
        # => [NV, K, V_in_leaf, R]

        # Prepare stack to update Merkle store
        # (X cycles)
        movupw.3 dup.8 push.LEAF_DEPTH
        # => [depth, K[3], R, NV, K, V_in_leaf]

        # Update Merkle store (29 cycles)
        mtree_set
        # => [NV_old, R_new, K, V_in_leaf]

        # Confirm that claimed `V_in_leaf` from advice provider is correct by checking if 
        # `[K, V_in_leaf]` hashes to `NV_old`
        # (33 cycles)
        movupw.2 dupw.3 hmerge assert_eqw
        # => [R_new, V_in_leaf]

        # Clean up stack for return
        # (1 cycle)
        swapw
        # => [V_in_leaf, R_new]
    else
        # Leaf becomes a Multiple kv-pair case
        # TODO (fail for now)
        push.1 assertz
    end
end

#! Removes the provided key/value pair from the leaf. The leaf to which we're inserting is
#! guaranteed to hold a single key-value pair (provided on the advice stack). Hence, after the
#! operation, the leaf will be empty.
#!
#! Inputs:
#!   Operand stack: [V (=ZERO), K, R, ...]
#!   Advice stack:  [K_in_leaf, V_in_leaf]
#!
#! Outputs:
#!   Operand stack: [V_old, R_new, ...]
#!
#! Cycles: X 
proc.remove_single_leaf
    # Push the leaf pre-image on stack
    # (0 cycles)
    adv_push.8
    # => [V_in_leaf, K_in_leaf, V, K, R]

    # Check if the key stored in the leaf is the same as K
    # (X cycles)
    movupw.3 movupw.2 eqw
    # => [K_in_leaf==K, K_in_leaf, K, V_in_leaf, V, R]

    if.true
        # Keys match; we're removing the value associated with K

        # (4 cycles)
        dropw 
        # => [K, V_in_leaf, V, R]

        # Update advice map (3 cycles)
        movupw.2 adv.insert_hdword
        # => [V, K, V_in_leaf, R]

        # Prepare the stack for `mtree_set`
        # Note that the new node value will be the empty word, so we can use `V`
        # as the node value (since we confirmed that it's `ZERO`)
        # (7 cycles)
        movupw.3 dup.8 push.LEAF_DEPTH
        # => [depth, K[3], R, V, K, V_in_leaf]

        # (29 cycles)
        mtree_set
        # => [NV_old, R_new, K, V_in_leaf, ...] 

        # Confirm that hmerge([K, V_in_leaf]) = NV_old
        # (33 cycles)
        movupw.2 dupw.3 hmerge assert_eqw
        # => [R_new, V_in_leaf, ...] 

        # Cleanup stack for return (1 cycle)
        swapw
        # => [V_in_leaf, R_new, ...] 
    else
        # Keys don't match; this is a no-op
        # We need to ensure that hash([K_in_leaf, V_in_leaf]) = NV;
        # that is, we need to verify the advice provider's claims.
        # If all checks pass, we're done.

        # => [K_in_leaf, K, V_in_leaf, V, R]

        # We no longer need V, since we're not removing anything
        movupw.3 dropw
        # => [K_in_leaf, K, V_in_leaf, R]

        # Prepare stack for mtree_get
        movupw.3 dup.8 push.LEAF_DEPTH
        # => [depth, K[3], R, K_in_leaf, K, V_in_leaf]

        # Retrieve node value (NV) from merkle tree
        mtree_get
        # => [NV, R, K_in_leaf, K, V_in_leaf]

        # Cleanup stack (we no longer need K)
        movupw.3 dropw
        # => [NV, R, K_in_leaf, V_in_leaf]

        # Ensure that hash([K_in_leaf, V_in_leaf]) == NV
        movupw.2 movupw.3 hmerge assert_eqw
        # => [R]

        # Prepare stack for return
        padw
        # => [ZERO, R]
    end
end

#! Inserts or removes a value associated with the given key. The leaf to which we're inserting is
#! guaranteed to hold a single key-value pair (provided on the advice stack).
#!
#! Inputs:
#!   Operand stack: [V, K, R, ...]
#!   Advice stack:  [K_in_leaf, V_in_leaf]
#!
#! Outputs:
#!   Operand stack: [V_old, R_new, ...]
#! Cycles:
#!   Remove: X cycles
#!   Insert; leaf single after insertion: X cycles
#!   Insert; leaf multiple after insertion: unimplemented
proc.set_single_leaf
    # Check if we're inserting or removing a value
    # (X cycles)
    padw eqw
    # => [V==ZERO, ZERO, V, K, R, ...] 
    if.true
        # we're removing the value associated with K (if any)

        # (4 cycles)
        dropw 
        # => [V, K, R, ...] 

        # (X cycles)
        exec.remove_single_leaf
        # => [V_old, R_new]
    else
        # we're inserting the key/value pair

        # (4 cycles)
        dropw 
        # => [V, K, R, ...] 

        # (X cycles)
        exec.insert_single_leaf
        # => [V_old, R_new]
    end
end

#! Inserts the specified value under the specified key in a Sparse Merkle Tree defined by the
#! specified root. If the insert is successful, the old value located under the specified key
#! is returned via the stack.
#!
#! If the VALUE is an empty word (i.e., [ZERO; 4]), the new state of the tree is guaranteed to
#! be equivalent to the state as if the updated value was never inserted.
#!
#! Inputs:
#!   Operand stack: [V, K, R, ...]
#! Outputs:
#!   Operand stack: [V_old, R_new, ...]
#!
#! Fails if the tree with the specified root does not exits in the VM's advice provider.
#!
#! Cycles
#!   Leaf empty
#!     removal: 74 cycles
#!     insertion: 133 cycles
#!   Leaf single
#!     removal: 227 cycles
#!     insertion (leaf remains single): 205
#!     insertion (leaf becomes multiple): unimplemented
#!   Leaf multiple
#!     unimplemented
export.set
    # Prepare stack for adv.push_mtnode
    # (X cycles)
    movupw.2 dup.8 push.LEAF_DEPTH
    # => [depth, leaf_index, R, V, K]

    # Push MT node on advice stack, cleanup operand stack, and then
    # push MT node on operand stack (NV)
    # (X cycles)
    adv.push_mtnode drop drop movdnw.2 adv_push.4
    # => [NV, V, K, R]

    # (X cycles)
    padw eqw
    # => [NV == ZERO, ZERO, NV, V, K, R]

    if.true
        # empty leaf

        # (8 cycles)
        dropw dropw
        #=> [V, K, R]

        # (insert empty value: X cycles)
        # (insert non-empty value: X cycles)
        exec.set_empty_leaf
    else
        # Single or Multiple leaf

        # (X cycles)
        dropw
        # => [NV, V, K, R]

        # Retrieve leaf pre-image on advice stack, and push leaf size on stack
        # Note: the rest of the leaf pre-image will be pulled out later
        # (4 cycles)
        adv.push_mapvaln dropw adv_push.1
        # => [leaf_size, V, K, R]

        # Leaf size will be a multiple of 8 (each kv-pair in a leaf is 8 elements)
        # (3 cycles)
        dup eq.8
        # => [is_single_kv_pair, leaf_size, V, K, R]

        if.true
            #  Single kv-pair case

            # (1 cycle)
            drop
            # => [V, K, R]

            # (remove key/value: X cycles)
            # (insert; leaf single after insertion: X cycles)
            exec.set_single_leaf
        else
            # Multiple kv-pair case
            # TODO (fail for now)
            push.1 assertz
        end
    end
end

# GET
# =================================================================================================

#! Returns the value located under the specified key in the Sparse Merkle Tree defined by the
#! specified root.
#!
#! If no values had been previously inserted under the specified key, an empty word (i.e.,
#! [ZERO; 4]) is returned.
#!
#! Inputs:
#!   Operand stack: [K, R, ...]
#!
#! Outputs:
#!   Operand stack: [V, R, ...]
#!
#! Fails if the tree with the specified root does not exits in the VM's advice provider.
#!
#! Cycles
#!   Leaf empty: 48 cycles
#!   Leaf single: 99 cycles
#!   Leaf multiple: unimplemented
export.get
    # Prepare for `mtree_get`
    # (6 cycles)
    dupw.1 dup.4 push.LEAF_DEPTH
    # => [depth, K[3], R, K, R]

    # Retrieve node value from merkle store
    # (14 cycles)
    mtree_get swapw dropw
    # => [NV, K, R]

    # Check if value is empty; if so, return empty value
    # (19 cycles)
    padw eqw
    # => [NV == 0, ZERO, V, K, R]

    if.true
        # Return empty value
        # (9 cycles)
        dropw swapw dropw
        # => [NV, R]
    else
        # Drop extra ZERO word
        # (4 cycles)
        dropw
        # => [NV, K, R]

        # Get leaf pre-image from advice map. Push the leaf preimage size on the stack
        # (0 cycles)
        adv.push_mapvaln adv_push.1
        # => [leaf_size, NV, K, R]

        # Leaf size will be a multiple of 8 (each kv-pair in a leaf is 8 elements)
        # (3 cycles)
        dup eq.8
        # => [is_single_kv_pair, leaf_size, NV, K, R]

        if.true
            # Single kv-pair case

            # Push leaf pre-image on stack (single K-V pair)
            # (1 cycle)
            drop adv_push.8
            # => [V, K, NV, K, R]

            # Confirm that the key stored in the leaf is as expected
            # (18 cycles)
            movupw.3 dupw.2 assert_eqw
            # => [V, K, NV, R]

            # Duplicate V to return it after hash check
            # (7 cycles)
            dupw movdnw.3
            # => [V, K, NV, V, R]

            # Hash leaf preimage and ensure that it equals node value
            # (27 cycles)
            hmerge assert_eqw
            # => [V, R]
        else
            # Multiple kv-pair case
            # TODO (fail for now)
            push.1 assertz
        end
    end
end