#include <stdio.h>
#include <stdlib.h>
#include <math.h>

struct Element {
    int value;
    int flag;
};

#if defined(VARIANT) && VARIANT == 0
void sum_if_flagged(struct Element *elements, int size, int *sum) {
    asm volatile(
        "beqz %[size], end%=\n\t"
        "start%=:"

        "lw a1, 4(%[elements])\n\t"
        "beqz a1, loopend%=\n\t"
        "lw a0, 0(%[elements])\n\t"
        "add %[sum], %[sum], a0\n\t"

        "loopend%=:"
        "addi %[size], %[size], -1\n\t"
        "addi %[elements], %[elements], 8\n\t"
        "bnez %[size], start%=\n"
        "end%=:"
        : [sum] "+r" (*sum)
        : [elements] "r" (elements), [size] "r" (size)
        : "a0", "a1"
    );
}
#elif defined(VARIANT) && VARIANT == 1
void sum_if_flagged(struct Element *elements, int size, int *sum) {
    asm volatile(
        "beqz %[size], end%=\n\t"
        "start%=:"

        "lw a0, 0(%[elements])\n\t"
        "lw a1, 4(%[elements])\n\t"
        "sltu a2, x0, a1\n\t"
        "mul a0, a0, a2\n\t"
        "add %[sum], %[sum], a0\n\t"

        "loopend%=:"
        "addi %[size], %[size], -1\n\t"
        "addi %[elements], %[elements], 8\n\t"
        "bnez %[size], start%=\n"
        "end%=:"
        : [sum] "+r" (*sum)
        : [elements] "r" (elements), [size] "r" (size)
        : "a0", "a1", "a2"
    );
}
#elif defined(VARIANT) && VARIANT == 2
void sum_if_flagged(struct Element *elements, int size, int *sum) {
    asm volatile(
        "beqz %[size], end%=\n\t"
        "start%=:"

        "lw a1, 4(%[elements])\n\t"
        "bms problematicbranch%=,branchfree%=,22\n"
        "problematicbranch%=:\n\t"
        "beqz a1, loopendtop%=\n\t"
        "lw a0, 0(%[elements])\n\t"
        "add %[sum], %[sum], a0\n\t"

        "loopendtop%=:"
        "addi %[size], %[size], -1\n\t"
        "addi %[elements], %[elements], 8\n\t"
        "bnez %[size], start%=\n"
        "j end%=\n"

        "branchfree%=:\n\t"
        "lw a0, 0(%[elements])\n\t"
        "sltu a2, x0, a1\n\t"
        "mul a0, a0, a2\n\t"
        "add %[sum], %[sum], a0\n\t"
        "bjr a0, problematicbranch%=, loopendtop%=\n"
        
        "loopend%=:"
        "addi %[size], %[size], -1\n\t"
        "addi %[elements], %[elements], 8\n\t"
        "bnez %[size], start%=\n"
        "end%=:"
        : [sum] "+r" (*sum)
        : [elements] "r" (elements), [size] "r" (size)
        : "a0", "a1", "a2"
    );
}
#endif


int main() {
    struct Element elements[ELEMENTS];

    #if defined(PREDICTABLE) && PREDICTABLE == 0
    srand(0);
    for (int i = 0; i < ELEMENTS; i++) {
        elements[i] = (Element) {i, rand() % 2};
    }
    #elif defined(PREDICTABLE) && PREDICTABLE == 1
    struct Element nonRandomElements[ELEMENTS];
    for (int i = 0; i < ELEMENTS; i++) {
        elements[i] = (Element) {i, i % 2};
    }
    #elif defined(PREDICTABLE) && PREDICTABLE == 2
    // Flags get less random as the index increases
    for (int i = 0; i < ELEMENTS; i++) {
        int pred = i % 2;
        int random = rand() % 2;
        if (i < ELEMENTS / 30) {
            pred = random;
        }
        elements[i] = (Element) {i, pred};
    }
    #endif

    int sum = 0;

    long long cyclesStart;
    asm volatile("csrr %0, cycle" : "=r" (cyclesStart));

    sum_if_flagged(elements, ELEMENTS, &sum);

    long long cyclesEnd;
    asm volatile("csrr %0, cycle" : "=r" (cyclesEnd));

    printf("Sum: %d\n", sum);
    printf("Cycles: %lld\n", cyclesEnd - cyclesStart);

    return 0;
}