pragma circom 2.0.0;

/////////////////////////////////////////////////////////////////////////////////////
/////////////////////// templates from the circomlib ////////////////////////////////
////////////////// Copy-pasted here for easy reference //////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////

/*
 * outputs `a` AND `b`
 */
template AND() {
    signal input a;
    signal input b;
    signal output out;

    out <== a*b;
}

/*
 * outputs `a` OR `b`
 */
template OR() {
    signal input a;
    signal input b;
    signal output out;

    out <== a + b - a*b;
}

/*
 * `out` = `cond` ? `L` : `R`
 */
template IfThenElse() {
    signal input cond;
    signal input L;
    signal input R;
    signal output out;

    out <== cond * (L - R) + R;
}

/*
 * (`outL`, `outR`) = `sel` ? (`R`, `L`) : (`L`, `R`)
 */
template Switcher() {
    signal input sel;
    signal input L;
    signal input R;
    signal output outL;
    signal output outR;

    signal aux;

    aux <== (R-L)*sel;
    outL <==  aux + L;
    outR <== -aux + R;
}

/*
 * Decomposes `in` into `b` bits, given by `bits`.
 * Least significant bit in `bits[0]`.
 * Enforces that `in` is at most `b` bits long.
 */
template Num2Bits(b) {
    signal input in;
    signal output bits[b];

    for (var i = 0; i < b; i++) {
        bits[i] <-- (in >> i) & 1;
        bits[i] * (1 - bits[i]) === 0;
    }
    var sum_of_bits = 0;
    for (var i = 0; i < b; i++) {
        sum_of_bits += (2 ** i) * bits[i];
    }
    sum_of_bits === in;
}

/*
 * Reconstructs `out` from `b` bits, given by `bits`.
 * Least significant bit in `bits[0]`.
 */
template Bits2Num(b) {
    signal input bits[b];
    signal output out;
    var lc = 0;

    for (var i = 0; i < b; i++) {
        lc += (bits[i] * (1 << i));
    }
    out <== lc;
}

/*
 * Checks if `in` is zero and returns the output in `out`.
 */
template IsZero() {
    signal input in;
    signal output out;

    signal inv;

    inv <-- in!=0 ? 1/in : 0;

    out <== -in*inv +1;
    in*out === 0;
}

/*
 * Checks if `in[0]` == `in[1]` and returns the output in `out`.
 */
template IsEqual() {
    signal input in[2];
    signal output out;

    component isz = IsZero();

    in[1] - in[0] ==> isz.in;

    isz.out ==> out;
}

/*
 * Checks if `in[0]` < `in[1]` and returns the output in `out`.
 */
template LessThan(n) {
    assert(n <= 252);
    signal input in[2];
    signal output out;

    component n2b = Num2Bits(n+1);

    n2b.in <== in[0]+ (1<<n) - in[1];

    out <== 1-n2b.bits[n];
}

/////////////////////////////////////////////////////////////////////////////////////
///////////////////////// templates for this lab ////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////

/*
 * outputs `out` = 1 if `in` is at most `b` bits long, and 0 otherwise.
 */
template CheckBitLength(b) {
    signal input in;
    signal output out;

    // (from `Num2Bits`)
    signal bits[b];
    
    for (var i = 0; i < b; i++) {
        bits[i] <-- (in >> i) & 1;
        bits[i] * (1 - bits[i]) === 0;
    }
    var sum_of_bits = 0;
    for (var i = 0; i < b; i++) {
        sum_of_bits += (1 << i) * bits[i];
    }

    // in - sum of bits ?= 0
    component is_zero = IsZero();
    is_zero.in <== in - sum_of_bits;

    out <== is_zero.out;
}


/*
 * Enforces the well-formedness of an exponent-mantissa pair (e, m), which is defined as follows:
 * if `e` is zero, then `m` must be zero
 * else, `e` must be at most `k` bits long, and `m` must be in the range [2^p, 2^p+1)
 */
template CheckWellFormedness(k, p) {
    signal input e;
    signal input m;

    // check if `e` is zero
    component is_e_zero = IsZero();
    is_e_zero.in <== e;

    // Case I: `e` is zero
    //// `m` must be zero
    component is_m_zero = IsZero();
    is_m_zero.in <== m;

    // Case II: `e` is nonzero
    //// `e` is `k` bits
    component check_e_bits = CheckBitLength(k);
    check_e_bits.in <== e;
    //// `m` is `p`+1 bits with the MSB equal to 1
    //// equivalent to check `m` - 2^`p` is in `p` bits
    component check_m_bits = CheckBitLength(p);
    check_m_bits.in <== m - (1 << p);

    // choose the right checks based on `is_e_zero`
    component if_else = IfThenElse();
    if_else.cond <== is_e_zero.out;
    if_else.L <== is_m_zero.out;
    //// check_m_bits.out * check_e_bits.out is equivalent to check_m_bits.out AND check_e_bits.out
    if_else.R <== check_m_bits.out * check_e_bits.out;

    // assert that those checks passed
    if_else.out === 1;
}

/*
 * Right-shifts `x` by `shift` bits to output `y`, where `shift` is a public circuit parameter.
 */
template RightShift(b, shift) {
    assert(shift < b);
    signal input x;
    signal output y;
    signal r;

    // right shift
    y <-- x \ 2 ** shift;

    // remainder
    r <-- x % 2 ** shift;

    // y * 2^shift - x + remainder == 0
    y * 2 ** shift - x + r === 0;

    // x < 2^b
    component less_than = LessThan(shift);
    less_than.in[0] <== x;
    less_than.in[1] <== 2 ** b;

    less_than.out === 1;
}


/*
 * Rounds the input floating-point number and checks to ensure that rounding does not make the mantissa unnormalized.
 * Rounding is necessary to prevent the bitlength of the mantissa from growing with each successive operation.
 * The input is a normalized floating-point number (e, m) with precision `P`, where `e` is a `k`-bit exponent and `m` is a `P`+1-bit mantissa.
 * The output is a normalized floating-point number (e_out, m_out) representing the same value with a lower precision `p`.
 */
template RoundAndCheck(k, p, P) {
    signal input e;
    signal input m;
    signal output e_out;
    signal output m_out;
    assert(P > p);

    // check if no overflow occurs
    component if_no_overflow = LessThan(P+1);
    if_no_overflow.in[0] <== m;
    if_no_overflow.in[1] <== (1 << (P+1)) - (1 << (P-p-1));
    signal no_overflow <== if_no_overflow.out;

    var round_amt = P-p;
    // Case I: no overflow
    // compute (m + 2^{round_amt-1}) >> round_amt
    var m_prime = m + (1 << (round_amt-1));
    component right_shift = RightShift(round_amt);
    right_shift.x <== m_prime;
    var m_out_1 = right_shift.y;
    var e_out_1 = e;

    // Case II: overflow
    var e_out_2 = e + 1;
    var m_out_2 = (1 << p);

    // select right output based on no_overflow
    component if_else[2];
    for (var i = 0; i < 2; i++) {
        if_else[i] = IfThenElse();
        if_else[i].cond <== no_overflow;
    }
    if_else[0].L <== e_out_1;
    if_else[0].R <== e_out_2;
    if_else[1].L <== m_out_1;
    if_else[1].R <== m_out_2;
    e_out <== if_else[0].out;
    m_out <== if_else[1].out;
}

/*
 * Left-shifts `x` by `shift` bits to output `y`.
 * Enforces 0 <= `shift` < `shift_bound`.
 * If `skip_checks` = 1, then we don't care about the output and the `shift_bound` constraint is not enforced.
 */
template LeftShift(shift_bound) {
    signal input x;
    signal input shift;
    signal input skip_checks;
    signal output y;

    // left shift
    y <-- x * 2 ** shift;

    // 0 <= shift <==> 0 < shift + 1
    component is_zero = IsZero();
    is_zero.in <== shift + 1;
    is_zero.out === 0;

    // shift < shift bound
    component less_than = LessThan(shift_bound);
    less_than.in[0] <== shift;
    less_than.in[1] <== shift_bound;
    less_than.out === 1 - skip_checks;
}

/*
 * Find the Most-Significant Non-Zero Bit (MSNZB) of `in`, where `in` is assumed to be non-zero value of `b` bits.
 * outputs the MSNZB as a one-hot vector `one_hot` of `b` bits, where `one_hot`[i] = 1 if MSNZB(`in`) = i and 0 otherwise.
 * The MSNZB is output as a one-hot vector to reduce the number of constraints in the subsequent `Normalize` template.
 * Enforces that `in` is non-zero as MSNZB(0) is undefined.
 * If `skip_checks` = 1, then we don't care about the output and the non-zero constraint is not enforced.
 */
template MSNZB(b) {
    signal input in;
    signal input skip_checks;
    signal output one_hot[b];
    
    // guard clause: in != 0; skip check
    component is_zero = IsZero();
    is_zero.in <== in;
    is_zero.out === skip_checks;

    // mask: one_bit[i] = (1 - in[i+1]) * one_bit[i+1]
    component in_bits = Num2Bits(b);
    in_bits.in <== in;

    signal one_bits[b]; one_bits[b-1] <== 1;
    for (var i = b-2; i >= 0; i--) {
        one_bits[i] <== (1 - in_bits.bits[i+1]) * one_bits[i+1];
    }

    // assign to output (current bits * mask)
    for (var i = 0; i < b; i++) {
        one_hot[i] <== in_bits.bits[i] * one_bits[i];
    }
}

/*
 * Normalizes the input floating-point number.
 * The input is a floating-point number with a `k`-bit exponent `e` and a `P`+1-bit *unnormalized* mantissa `m` with precision `p`, where `m` is assumed to be non-zero.
 * The output is a floating-point number representing the same value with exponent `e_out` and a *normalized* mantissa `m_out` of `P`+1-bits and precision `P`.
 * Enforces that `m` is non-zero as a zero-value can not be normalized.
 * If `skip_checks` = 1, then we don't care about the output and the non-zero constraint is not enforced.
 */
template Normalize(k, p, P) {
    signal input e;
    signal input m;
    signal input skip_checks;
    signal output e_out;
    signal output m_out;
    assert(P > p);

    // normalize exponent and mantissa
    // (m > 0; skip_checks implemented in MSNZB)
    component msnzb = MSNZB(P+1);
    msnzb.in <== m;
    msnzb.skip_checks <== skip_checks;

    var e_term = 0, m_factor = 0;
    for (var i = 0; i <= P; i++) {
        // compute the decimal value of the exponent
        e_term += msnzb.one_hot[i] * i;
        // compute the decimal value of the mantissa
        m_factor += msnzb.one_hot[i] * 2 ** (P-i);
    }

    // convert the decimal value of the exponent to the biased exponent value
    e_out <== e + e_term - p;
    // compute the final mantissa value
    m_out <== m_factor * m;
}

/*
 * Adds two floating-point numbers.
 * The inputs are normalized floating-point numbers with `k`-bit exponents `e` and `p`+1-bit mantissas `m` with scale `p`.
 * Does not assume that the inputs are well-formed and makes appropriate checks for the same.
 * The output is a normalized floating-point number with exponent `e_out` and mantissa `m_out` of `p`+1-bits and scale `p`.
 * Enforces that inputs are well-formed.
 */
template FloatAdd(k, p) {
    signal input e[2];
    signal input m[2];
    signal output e_out;
    signal output m_out;
    
    // well-formedness
    component wellformedness_0 = CheckWellFormedness(k, p);
    wellformedness_0.e <== e[0];
    wellformedness_0.m <== m[0];

    component wellformedness_1 = CheckWellFormedness(k, p);
    wellformedness_1.e <== e[1];
    wellformedness_1.m <== m[1];

    // identify the number with the largest magnitude

    // concatenate the exponent and mantissa values by left-shifting the exponent value 
    // by (p+1) bits and then adding the mantissa value to it
    signal n0 <== e[0] * 2 ** (p+1) + m[0];
    signal n1 <== e[1] * 2 ** (p+1) + m[1];
    signal largest_e, largest_m, smallest_e, smallest_m;

    // compare these numbers
    component less_than = LessThan(k+p);
    less_than.in[0] <== n0;
    less_than.in[1] <== n1;

    // get the exponent of the largest number
    component switcher_e = Switcher(); 
    switcher_e.sel <== less_than.out;
    switcher_e.L <== e[0];
    switcher_e.R <== e[1];

    largest_e <== switcher_e.outL; smallest_e <== switcher_e.outR;

    // get the mantissa of the largest number
    component switcher_m = Switcher();
    switcher_m.sel <== less_than.out;
    switcher_m.L <== m[0];
    switcher_m.R <== m[1];

    largest_m <== switcher_m.outL; smallest_m <== switcher_m.outR;

    // subnormal float? Check whether:
    // - the exponent of the larger number is zero
    // OR
    // - the difference between the two exponents is greater than p+1
    signal is_subnormal; 
    signal diff_e <== largest_e - smallest_e;

    // (subnormal) case I: largest exponent == 0?
    component is_zero = IsZero();
    is_zero.in <== largest_e;

    // (subnormal) case II: diff exponents > p+1?
    component less_than_diff = LessThan(k);
    less_than_diff.in[0] <== p+1;
    less_than_diff.in[1] <== diff_e;

    // (subnormal) case I OR case II?
    component or = OR();
    or.a <== less_than_diff.out;
    or.b <== is_zero.out;
    is_subnormal <== or.out;

    component if_else_largest_m = IfThenElse();
    if_else_largest_m.cond <== is_subnormal;
    if_else_largest_m.L <== 1;
    if_else_largest_m.R <== largest_m;

    component if_else_smallest_e = IfThenElse();
    if_else_smallest_e.cond <== is_subnormal;
    if_else_smallest_e.L <== 1;
    if_else_smallest_e.R <== smallest_e;

    component if_else_diff_e = IfThenElse();
    if_else_diff_e.cond <== is_subnormal;
    if_else_diff_e.L <== 0;
    if_else_diff_e.R <== diff_e;

    // shift the mantissa and adjust the exponent

    // left-shift the larger of the two mantissas by the difference in exponents between the two numbers
    component m_largest_left_shift = LeftShift(p+2);
    m_largest_left_shift.x <== if_else_largest_m.out;
    m_largest_left_shift.shift <== if_else_diff_e.out;
    m_largest_left_shift.skip_checks <== 0;

    // normalize
    component normalize = Normalize(k, p, 2*p+1);
    normalize.e <== if_else_smallest_e.out;
    normalize.m <== m_largest_left_shift.y + smallest_m;
    normalize.skip_checks <== 0;

    // round
    component round = RoundAndCheck(k, p, 2*p+1);
    round.e <== normalize.e_out;
    round.m <== normalize.m_out;

    // if subnormal, return rounded; 
    // if not, return the largest number
    component if_else_e = IfThenElse();
    if_else_e.cond <== is_subnormal;
    if_else_e.L <== largest_e;
    if_else_e.R <== round.e_out;

    component if_else_m = IfThenElse();
    if_else_m.cond <== is_subnormal;
    if_else_m.L <== largest_m;
    if_else_m.R <== round.m_out;

    e_out <== if_else_e.out;
    m_out <== if_else_m.out;
}