use starknet_types_core::qm31::QM31;
use crate::math_utils::qm31_pack_reduced;
use crate::types::relocatable::MaybeRelocatable;
use crate::types::{errors::math_errors::MathError, instruction::OpcodeExtension};
use crate::vm::errors::vm_errors::VirtualMachineError;
use crate::Felt252;
pub fn typed_add(
x: &MaybeRelocatable,
y: &MaybeRelocatable,
opcode_extension: OpcodeExtension,
) -> Result<MaybeRelocatable, VirtualMachineError> {
match opcode_extension {
OpcodeExtension::Stone => Ok(x.add(y)?),
OpcodeExtension::QM31Operation => {
if let (MaybeRelocatable::Int(num_x), MaybeRelocatable::Int(num_y)) = (x, y) {
let x_qm31 = QM31::unpack_from_felt(num_x)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let y_qm31 = QM31::unpack_from_felt(num_y)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let res = x_qm31 + y_qm31;
Ok(MaybeRelocatable::Int(qm31_pack_reduced(res)))
} else {
Err(VirtualMachineError::Math(MathError::RelocatableQM31Add(
Box::new((x.clone(), y.clone())),
)))
}
}
_ => Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(
"typed_add".to_owned().into_boxed_str(),
)),
}
}
pub fn typed_sub(
x: &MaybeRelocatable,
y: &MaybeRelocatable,
opcode_extension: OpcodeExtension,
) -> Result<MaybeRelocatable, VirtualMachineError> {
match opcode_extension {
OpcodeExtension::Stone => Ok(x.sub(y)?),
OpcodeExtension::QM31Operation => {
if let (MaybeRelocatable::Int(num_x), MaybeRelocatable::Int(num_y)) = (x, y) {
let x_qm31 = QM31::unpack_from_felt(num_x)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let y_qm31 = QM31::unpack_from_felt(num_y)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let res = x_qm31 - y_qm31;
Ok(MaybeRelocatable::Int(qm31_pack_reduced(res)))
} else {
Err(VirtualMachineError::Math(MathError::RelocatableQM31Sub(
Box::new((x.clone(), y.clone())),
)))
}
}
_ => Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(
"typed_sub".to_owned().into_boxed_str(),
)),
}
}
pub fn typed_mul(
x: &MaybeRelocatable,
y: &MaybeRelocatable,
opcode_extension: OpcodeExtension,
) -> Result<MaybeRelocatable, VirtualMachineError> {
if let (MaybeRelocatable::Int(num_x), MaybeRelocatable::Int(num_y)) = (x, y) {
match opcode_extension {
OpcodeExtension::Stone => Ok(MaybeRelocatable::Int(num_x * num_y)),
OpcodeExtension::QM31Operation => {
let x_qm31 = QM31::unpack_from_felt(num_x)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let y_qm31 = QM31::unpack_from_felt(num_y)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let res = x_qm31 * y_qm31;
Ok(MaybeRelocatable::Int(qm31_pack_reduced(res)))
}
_ => Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(
"typed_mul".to_owned().into_boxed_str(),
)),
}
} else {
Err(VirtualMachineError::ComputeResRelocatableMul(Box::new((
x.clone(),
y.clone(),
))))
}
}
pub fn typed_div(
x: &Felt252,
y: &Felt252,
opcode_extension: OpcodeExtension,
) -> Result<Felt252, VirtualMachineError> {
match opcode_extension {
OpcodeExtension::Stone => {
Ok(x.field_div(&y.try_into().map_err(|_| MathError::DividedByZero)?))
}
OpcodeExtension::QM31Operation => {
let x_qm31 = QM31::unpack_from_felt(x)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let y_qm31 = QM31::unpack_from_felt(y)
.map_err(|e| MathError::QM31UnreducedError(Box::new(e.0)))?;
let res = (x_qm31 / y_qm31).map_err(|_| MathError::DividedByZero)?;
Ok(qm31_pack_reduced(res))
}
_ => Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(
"typed_div".to_owned().into_boxed_str(),
)),
}
}
#[cfg(test)]
mod decoder_test {
use super::*;
use assert_matches::assert_matches;
#[test]
fn typed_add_blake() {
let a = &MaybeRelocatable::from(5);
let b = &MaybeRelocatable::from(6);
let error = typed_add(a, b, OpcodeExtension::Blake);
assert_matches!(
error,
Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(ref message)) if message.as_ref() == "typed_add"
);
}
#[test]
fn typed_sub_blake() {
let a = &MaybeRelocatable::from(7);
let b = &MaybeRelocatable::from(3);
let error = typed_sub(a, b, OpcodeExtension::Blake);
assert_matches!(
error,
Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(ref message)) if message.as_ref() == "typed_sub"
);
}
#[test]
fn relocatable_typed_sub_q31_operation() {
let a = &MaybeRelocatable::from((6, 8));
let b = &MaybeRelocatable::from(2);
let error = typed_sub(a, b, OpcodeExtension::QM31Operation);
assert_matches!(
error,
Err(VirtualMachineError::Math(MathError::RelocatableQM31Sub(bx))) if *bx ==
(MaybeRelocatable::from((6, 8)), MaybeRelocatable::from(2))
);
}
#[test]
fn typed_mul_blake_finalize() {
let a = &MaybeRelocatable::from(4);
let b = &MaybeRelocatable::from(9);
let error = typed_mul(a, b, OpcodeExtension::BlakeFinalize);
assert_matches!(
error,
Err(VirtualMachineError::InvalidTypedOperationOpcodeExtension(ref message)) if message.as_ref() == "typed_mul"
);
}
}