miden-processor 0.22.1

use crate::{
    Felt, ONE, ZERO,
    field::Field,
    operation::OperationError,
    processor::{Processor, StackInterface},
    tracer::OperationHelperRegisters,
};

#[cfg(test)]
mod tests;

// FIELD OPERATIONS
// ================================================================================================

/// Pops two elements off the stack, adds them together, and pushes the result back onto the
/// stack.
#[inline(always)]
pub(super) fn op_add<P>(processor: &mut P) -> Result<OperationHelperRegisters, OperationError>
where
    P: Processor,
{
    pop2_applyfn_push(processor, |a, b| a + b)?;
    Ok(OperationHelperRegisters::Empty)
}

/// Pops an element off the stack, computes its additive inverse, and pushes the result back
/// onto the stack.
#[inline(always)]
pub(super) fn op_neg<P: Processor>(processor: &mut P) -> OperationHelperRegisters {
    let element = processor.stack().get(0);
    processor.stack_mut().set(0, -element);
    OperationHelperRegisters::Empty
}

/// Pops two elements off the stack, multiplies them, and pushes the result back onto the
/// stack.
#[inline(always)]
pub(super) fn op_mul<P>(processor: &mut P) -> Result<OperationHelperRegisters, OperationError>
where
    P: Processor,
{
    pop2_applyfn_push(processor, |a, b| a * b)?;
    Ok(OperationHelperRegisters::Empty)
}

/// Pops an element off the stack, computes its multiplicative inverse, and pushes the result
/// back onto the stack.
///
/// # Errors
/// Returns an error if the value on the top of the stack is ZERO.
#[inline(always)]
pub(super) fn op_inv<P: Processor>(
    processor: &mut P,
) -> Result<OperationHelperRegisters, OperationError> {
    let top = processor.stack_mut().get_mut(0);
    if (*top) == ZERO {
        return Err(OperationError::DivideByZero);
    }
    *top = top.inverse();
    Ok(OperationHelperRegisters::Empty)
}

/// Pops an element off the stack, adds ONE to it, and pushes the result back onto the stack.
#[inline(always)]
pub(super) fn op_incr<P: Processor>(processor: &mut P) -> OperationHelperRegisters {
    *processor.stack_mut().get_mut(0) += ONE;
    OperationHelperRegisters::Empty
}

/// Pops two elements off the stack, computes their boolean AND, and pushes the result back
/// onto the stack.
///
/// # Errors
/// Returns an error if either of the two elements on the top of the stack is not a binary
/// value.
#[inline(always)]
pub(super) fn op_and<P>(processor: &mut P) -> Result<OperationHelperRegisters, OperationError>
where
    P: Processor,
{
    pop2_applyfn_push_op(processor, |a, b| {
        assert_binary(b)?;
        assert_binary(a)?;

        if a == ONE && b == ONE { Ok(ONE) } else { Ok(ZERO) }
    })?;
    Ok(OperationHelperRegisters::Empty)
}

/// Pops two elements off the stack, computes their boolean OR, and pushes the result back
/// onto the stack.
///
/// # Errors
/// Returns an error if either of the two elements on the top of the stack is not a binary
/// value.
#[inline(always)]
pub(super) fn op_or<P>(processor: &mut P) -> Result<OperationHelperRegisters, OperationError>
where
    P: Processor,
{
    pop2_applyfn_push_op(processor, |a, b| {
        assert_binary(b)?;
        assert_binary(a)?;

        if a == ONE || b == ONE { Ok(ONE) } else { Ok(ZERO) }
    })?;
    Ok(OperationHelperRegisters::Empty)
}

/// Pops an element off the stack, computes its boolean NOT, and pushes the result back onto
/// the stack.
///
/// # Errors
/// Returns an error if the value on the top of the stack is not a binary value.
#[inline(always)]
pub(super) fn op_not<P: Processor>(
    processor: &mut P,
) -> Result<OperationHelperRegisters, OperationError> {
    let top = processor.stack_mut().get_mut(0);
    if *top == ZERO {
        *top = ONE;
    } else if *top == ONE {
        *top = ZERO;
    } else {
        return Err(OperationError::NotBinaryValue { value: *top });
    }
    Ok(OperationHelperRegisters::Empty)
}

/// Pops two elements off the stack and compares them. If the elements are equal, pushes ONE
/// onto the stack, otherwise pushes ZERO onto the stack.
#[inline(always)]
pub(super) fn op_eq<P>(processor: &mut P) -> Result<OperationHelperRegisters, OperationError>
where
    P: Processor,
{
    let b = processor.stack().get(0);
    let a = processor.stack().get(1);

    // Directly manipulate the stack instead of using pop2_applyfn_push() since we need
    // to return user op helpers, which makes the abstraction less suitable here.
    processor.stack_mut().decrement_size()?;

    let result = if a == b { ONE } else { ZERO };
    processor.stack_mut().set(0, result);

    Ok(OperationHelperRegisters::Eq { stack_second: a, stack_first: b })
}

/// Pops an element off the stack and compares it to ZERO. If the element is ZERO, pushes ONE
/// onto the stack, otherwise pushes ZERO onto the stack.
#[inline(always)]
pub(super) fn op_eqz<P: Processor>(processor: &mut P) -> OperationHelperRegisters {
    let top = processor.stack_mut().get_mut(0);
    let old_top = *top;

    if old_top == ZERO {
        *top = ONE;
    } else {
        *top = ZERO;
    };

    OperationHelperRegisters::Eqz { top: old_top }
}

/// Computes a single turn of exp accumulation for the given inputs. The top 4 elements in the
/// stack are arranged as follows (from the top):
/// - 0: least significant bit of the exponent in the previous trace if there's an expacc call,
///   otherwise ZERO,
/// - 1: base of the exponentiation; i.e. `b` in `b^a`,
/// - 2: accumulated result of the exponentiation so far,
/// - 3: the exponent; i.e. `a` in `b^a`.
///
/// It is expected that `Expacc` is called at least `num_exp_bits` times, where `num_exp_bits`
/// is the number of bits needed to represent `exp`. The initial call to `Expacc` should set the
/// stack as [0, base, 1, exponent]. The subsequent call will set the stack either as
/// - [0, base^2, acc, exp/2], or
/// - [1, base^2, acc * base, exp/2],
///
/// depending on the least significant bit of the exponent.
///
/// Expacc is based on the observation that the exponentiation of a number can be computed by
/// repeatedly squaring the base and multiplying those powers of the base by the accumulator,
/// for the powers of the base which correspond to the exponent's bits which are set to 1.
///
/// For example, take b^5 = (b^2)^2 * b. Over the course of 3 iterations (5 = 101b), the
/// algorithm will compute b, b^2 and b^4 (placed in `base_acc`). Hence, we want to multiply
/// `base_acc` in `result_acc` when `base_acc = b` and when `base_acc = b^4`, which occurs on
/// the first and third iterations (corresponding to the `1` bits in the binary representation
/// of 5).
#[inline(always)]
pub(super) fn op_expacc<P: Processor>(processor: &mut P) -> OperationHelperRegisters {
    let old_base = processor.stack().get(1);
    let old_acc = processor.stack().get(2);
    let old_exp_int = processor.stack().get(3).as_canonical_u64();

    // Compute new exponent.
    let new_exp = Felt::new(old_exp_int >> 1);

    // Compute new accumulator. We update the accumulator only when the least significant bit of
    // the exponent is 1.
    let exp_lsb = old_exp_int & 1;
    let acc_update_val = if exp_lsb == 1 { old_base } else { ONE };
    let new_acc = old_acc * acc_update_val;

    // Compute the new base.
    let new_base = old_base * old_base;

    processor.stack_mut().set(0, Felt::new(exp_lsb));
    processor.stack_mut().set(1, new_base);
    processor.stack_mut().set(2, new_acc);
    processor.stack_mut().set(3, new_exp);

    OperationHelperRegisters::Expacc { acc_update_val }
}

/// Gets the top four values from the stack [b0, b1, a0, a1] (low coefficient at lower
/// index/closer to top), where a = (a0, a1) and b = (b0, b1) are elements of the
/// extension field, and outputs the product c = (c0, c1) where c0 = a0 * b0 + 7 * a1 * b1
/// and c1 = a0 * b1 + a1 * b0. The extension field is defined by the irreducible polynomial
/// x² - 7. It leaves b0, b1 in the first and second positions on the stack, sets c0 and c1
/// to the third and fourth positions, and leaves the rest of the stack unchanged.
#[inline(always)]
pub(super) fn op_ext2mul<P: Processor>(processor: &mut P) -> OperationHelperRegisters {
    const SEVEN: Felt = Felt::new(7);
    // get_word returns [s0, s1, s2, s3] where s0 is top of stack
    // Stack layout: s0=b0, s1=b1, s2=a0, s3=a1
    let [b0, b1, a0, a1]: [Felt; 4] = processor.stack().get_word(0).into();

    /* top 2 elements remain unchanged */

    let b0_times_a0 = b0 * a0;
    let b1_times_a1 = b1 * a1;
    let c1 = (b0 + b1) * (a1 + a0) - b0_times_a0 - b1_times_a1;
    let c0 = b0_times_a0 + SEVEN * b1_times_a1;
    processor.stack_mut().set(2, c0); // c0 (low) at position 2
    processor.stack_mut().set(3, c1); // c1 (high) at position 3
    OperationHelperRegisters::Empty
}

// HELPERS
// ----------------------------------------------------------------------------------------------

/// Pops the top two elements from the stack, applies the given infallible function to them,
/// and pushes the result back onto the stack.
///
/// The size of the stack is decremented by 1.
#[inline(always)]
fn pop2_applyfn_push<P>(
    processor: &mut P,
    f: impl FnOnce(Felt, Felt) -> Felt,
) -> Result<(), OperationError>
where
    P: Processor,
{
    let b = processor.stack().get(0);
    let a = processor.stack().get(1);

    processor.stack_mut().decrement_size()?;

    processor.stack_mut().set(0, f(a, b));
    Ok(())
}

/// Pops the top two elements from the stack, applies the given function to them, and pushes the
/// result back onto the stack. The function returns an `OperationError` on failure.
///
/// The size of the stack is decremented by 1.
#[inline(always)]
fn pop2_applyfn_push_op<P>(
    processor: &mut P,
    f: impl FnOnce(Felt, Felt) -> Result<Felt, OperationError>,
) -> Result<(), OperationError>
where
    P: Processor,
{
    let b = processor.stack().get(0);
    let a = processor.stack().get(1);

    processor.stack_mut().decrement_size()?;

    processor.stack_mut().set(0, f(a, b)?);

    Ok(())
}

/// Asserts that the given value is a binary value (0 or 1).
#[inline(always)]
fn assert_binary(value: Felt) -> Result<(), OperationError> {
    if value != ZERO && value != ONE {
        Err(OperationError::NotBinaryValue { value })
    } else {
        Ok(())
    }
}