From 4a42ddb2038371f3d67fd5f24ecb80cee821f8b2 Mon Sep 17 00:00:00 2001
From: Dmitry Vagner <dmitry.vagner@gmail.com>
Date: Thu, 27 Apr 2023 17:15:25 -0700
Subject: [PATCH] on stack

---
 evm/src/cpu/kernel/tests/bn254.rs | 10 +++-----
 evm/src/curve_pairings.rs         | 10 +++++++-
 evm/src/extension_tower.rs        | 42 +++++++++++++++++++++++++------
 3 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/evm/src/cpu/kernel/tests/bn254.rs b/evm/src/cpu/kernel/tests/bn254.rs
index 1f2f12e3..51966f59 100644
--- a/evm/src/cpu/kernel/tests/bn254.rs
+++ b/evm/src/cpu/kernel/tests/bn254.rs
@@ -1,5 +1,3 @@
-use std::mem::transmute;
-
 use anyhow::Result;
 use ethereum_types::U256;
 use rand::Rng;
@@ -8,7 +6,7 @@ use crate::cpu::kernel::interpreter::{
     run_interpreter_with_memory, Interpreter, InterpreterMemoryInitialization,
 };
 use crate::cpu::kernel::tests::u256ify;
-use crate::curve_pairings::{final_exponent, gen_fp12_sparse, miller_loop, Curve, CyclicGroup};
+use crate::curve_pairings::{final_exponent, gen_fp12_sparse, miller_loop, Curve};
 use crate::extension_tower::{FieldExt, Fp12, Fp2, Fp6, Stack, BN254};
 use crate::memory::segments::Segment::BnPairing;
 
@@ -210,10 +208,8 @@ fn test_bn_miller() -> Result<()> {
     let p: Curve<BN254> = rng.gen::<Curve<BN254>>();
     let q: Curve<Fp2<BN254>> = rng.gen::<Curve<Fp2<BN254>>>();
 
-    let p_stack: [U256; 2] = unsafe { transmute(p) };
-    let q_stack: [U256; 4] = unsafe { transmute(q) };
-    let mut input = p_stack.to_vec();
-    input.extend(q_stack);
+    let mut input = p.on_stack();
+    input.extend(q.on_stack());
 
     let setup = InterpreterMemoryInitialization {
         label: "bn254_miller".to_string(),
diff --git a/evm/src/curve_pairings.rs b/evm/src/curve_pairings.rs
index 5f3fc7fa..fc369406 100644
--- a/evm/src/curve_pairings.rs
+++ b/evm/src/curve_pairings.rs
@@ -5,7 +5,7 @@ use rand::distributions::Standard;
 use rand::prelude::Distribution;
 use rand::Rng;
 
-use crate::extension_tower::{FieldExt, Fp12, Fp2, Fp6, BN254};
+use crate::extension_tower::{FieldExt, Fp12, Fp2, Fp6, Stack, BN254};
 
 #[derive(Debug, Copy, Clone, PartialEq)]
 pub struct Curve<T>
@@ -25,6 +25,14 @@ impl<T: FieldExt> Curve<T> {
     }
 }
 
+impl<T: FieldExt + Stack> Curve<T> {
+    pub fn on_stack(self) -> Vec<U256> {
+        let mut stack = self.x.on_stack();
+        stack.extend(self.y.on_stack());
+        stack
+    }
+}
+
 impl<T> Curve<T>
 where
     T: FieldExt,
diff --git a/evm/src/extension_tower.rs b/evm/src/extension_tower.rs
index 2c81b035..0e654c88 100644
--- a/evm/src/extension_tower.rs
+++ b/evm/src/extension_tower.rs
@@ -1,5 +1,4 @@
 use std::fmt::Debug;
-use std::mem::transmute;
 use std::ops::{Add, Div, Mul, Neg, Sub};
 
 use ethereum_types::{U256, U512};
@@ -1227,22 +1226,51 @@ pub trait Stack {
     fn on_stack(self) -> Vec<U256>;
 }
 
+impl Stack for BN254 {
+    fn on_stack(self) -> Vec<U256> {
+        vec![self.val]
+    }
+}
+
 impl Stack for BLS381 {
     fn on_stack(self) -> Vec<U256> {
         vec![self.lo(), self.hi()]
     }
 }
 
-impl Stack for Fp6<BN254> {
+impl<T> Stack for Fp2<T>
+where
+    T: FieldExt + Stack,
+{
     fn on_stack(self) -> Vec<U256> {
-        let f: [U256; 6] = unsafe { transmute(self) };
-        f.into_iter().collect()
+        let mut stack = self.re.on_stack();
+        stack.extend(self.im.on_stack());
+        stack
     }
 }
 
-impl Stack for Fp12<BN254> {
+impl<T> Stack for Fp6<T>
+where
+    T: FieldExt,
+    Fp2<T>: Adj + Stack,
+{
     fn on_stack(self) -> Vec<U256> {
-        let f: [U256; 12] = unsafe { transmute(self) };
-        f.into_iter().collect()
+        let mut stack = self.t0.on_stack();
+        stack.extend(self.t1.on_stack());
+        stack.extend(self.t2.on_stack());
+        stack
+    }
+}
+
+impl<T> Stack for Fp12<T>
+where
+    T: FieldExt,
+    Fp2<T>: Adj,
+    Fp6<T>: Stack,
+{
+    fn on_stack(self) -> Vec<U256> {
+        let mut stack = self.z0.on_stack();
+        stack.extend(self.z1.on_stack());
+        stack
     }
 }