#include <Library/UefiLib.h>
#include <Library/BaseLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#include <Uefi.h>

#include "utils.h"
#include "utilsx1.h"
#include "params.h"
#include "thash.h"
#include "address.h"

/*
 * Generate the entire Merkle tree, computing the authentication path for
 * leaf_idx, and the resulting root node using Merkle's TreeHash algorithm.
 * Expects the layer and tree parts of the tree_addr to be set, as well as the
 * tree type (i.e. SPX_ADDR_TYPE_HASHTREE or SPX_ADDR_TYPE_FORSTREE)
 *
 * This expects tree_addr to be initialized to the addr structures for the
 * Merkle tree nodes
 *
 * Applies the offset idx_offset to indices before building addresses, so that
 * it is possible to continue counting indices across trees.
 *
 * This works by using the standard Merkle tree building algorithm,
 */
VOID treehashx1(UINT8 *root, UINT8 *auth_path,
                const spx_ctx* ctx,
                UINT32 leaf_idx, UINT32 idx_offset,
                UINT32 tree_height,
                VOID (*gen_leaf)(
                   UINT8* /* Where to write the leaves */,
                   const spx_ctx* /* ctx */,
                   UINT32 idx, VOID *info),
                UINT32 tree_addr[8],
                VOID *info)
{
    /* This is where we keep the intermediate nodes */
    SPX_VLA(UINT8, stack, tree_height*SPX_N);

    UINT32 idx;
    UINT32 max_idx = (UINT32)((1 << tree_height) - 1);
    for (idx = 0;; idx++) {
        UINT8 current[2*SPX_N];   /* Current logical node is at */
            /* index[SPX_N].  We do this to minimize the number of copies */
            /* needed during a thash */
        gen_leaf( &current[SPX_N], ctx, idx + idx_offset,
                    info );

        /* Now combine the freshly generated right node with previously */
        /* generated left ones */
        UINT32 internal_idx_offset = idx_offset;
        UINT32 internal_idx = idx;
        UINT32 internal_leaf = leaf_idx;
        UINT32 h;     /* The height we are in the Merkle tree */
        for (h=0;; h++, internal_idx >>= 1, internal_leaf >>= 1) {

            /* Check if we hit the top of the tree */
            if (h == tree_height) {
                /* We hit the root; return it */
                CopyMem( root, &current[SPX_N], SPX_N );
                return;
            }

            /*
             * Check if the node we have is a part of the
             * authentication path; if it is, write it out
             */
            if ((internal_idx ^ internal_leaf) == 0x01) {
                CopyMem( &auth_path[ h * SPX_N ],
                        &current[SPX_N],
                        SPX_N );
            }

            /*
             * Check if we're at a left child; if so, stop going up the stack
             * Exception: if we've reached the end of the tree, keep on going
             * (so we combine the last 4 nodes into the one root node in two
             * more iterations)
             */
            if ((internal_idx & 1) == 0 && idx < max_idx) {
                break;
            }

            /* Ok, we're at a right node */
            /* Now combine the left and right logical nodes together */

            /* Set the address of the node we're creating. */
            internal_idx_offset >>= 1;
            set_tree_height(tree_addr, h + 1);
            set_tree_index(tree_addr, internal_idx/2 + internal_idx_offset );

            UINT8 *left = &stack[h * SPX_N];
            CopyMem( &current[0], left, SPX_N );
            thash( &current[1 * SPX_N],
                   &current[0 * SPX_N],
                   2, ctx, tree_addr);
        }

        /* We've hit a left child; save the current for when we get the */
        /* corresponding right right */
        CopyMem( &stack[h * SPX_N], &current[SPX_N], SPX_N);
    }
}
