From 1f3e3de77c3cf9b965a32061beb153a07b6a4da9 Mon Sep 17 00:00:00 2001
From: Dmitry Vagner <dmitry.vagner@gmail.com>
Date: Wed, 22 Mar 2023 10:16:02 -0700
Subject: [PATCH] clean and generalize

---
 evm/src/cpu/kernel/tests/bls381.rs | 46 +++++++++++++-----------------
 1 file changed, 20 insertions(+), 26 deletions(-)

diff --git a/evm/src/cpu/kernel/tests/bls381.rs b/evm/src/cpu/kernel/tests/bls381.rs
index 0a35ebbb..6d412849 100644
--- a/evm/src/cpu/kernel/tests/bls381.rs
+++ b/evm/src/cpu/kernel/tests/bls381.rs
@@ -8,17 +8,20 @@ use crate::cpu::kernel::interpreter::{
 use crate::extension_tower::{Stack, BLS381};
 use crate::memory::segments::Segment::KernelGeneral;
 
-fn extract_stack(interpreter: Interpreter<'static>) -> Vec<U256> {
-    interpreter
-        .stack()
-        .iter()
-        .rev()
-        .cloned()
-        .collect::<Vec<U256>>()
-}
-
-fn combine_u256s(hi: U256, lo: U256) -> U512 {
-    U512::from(lo) + (U512::from(hi) << 256)
+fn run_and_return_bls(label: String, x: BLS381, y: BLS381) -> BLS381 {
+    let mut stack = x.on_stack();
+    stack.extend(y.on_stack());
+    let setup = InterpreterMemoryInitialization {
+        label,
+        stack,
+        segment: KernelGeneral,
+        memory: vec![],
+    };
+    let interpreter = run_interpreter_with_memory(setup).unwrap();
+    let output = interpreter.stack();
+    BLS381 {
+        val: U512::from(output[1]) + (U512::from(output[0]) << 256),
+    }
 }
 
 #[test]
@@ -27,22 +30,13 @@ fn test_bls_ops() -> Result<()> {
     let x: BLS381 = rng.gen::<BLS381>();
     let y: BLS381 = rng.gen::<BLS381>();
 
-    let mut stack = x.on_stack();
-    stack.extend(y.on_stack());
+    let output_add = run_and_return_bls("test_add_fp381".to_string(), x, y);
+    let output_mul = run_and_return_bls("test_mul_fp381".to_string(), x, y);
+    let output_sub = run_and_return_bls("test_sub_fp381".to_string(), x, y);
 
-    let setup = InterpreterMemoryInitialization {
-        label: "test_mul_fp381".to_string(),
-        stack,
-        segment: KernelGeneral,
-        memory: vec![],
-    };
-    let interpreter = run_interpreter_with_memory(setup).unwrap();
-    let output = extract_stack(interpreter);
-    println!("{:#?}", output);
-    let output_512 = combine_u256s(output[1], output[0]);
-    let expected = x * y;
-
-    assert_eq!(expected.val, output_512);
+    assert_eq!(output_add, x + y);
+    assert_eq!(output_mul, x * y);
+    assert_eq!(output_sub, x - y);
 
     Ok(())
 }