use eth_types::{bytecode, Bytecode, U256};
use rand::{rngs::ThreadRng, Rng};
pub struct Sha3CodeGen {
pub offset: usize,
pub size: usize,
data_len: usize,
rng: ThreadRng,
}
impl Sha3CodeGen {
pub fn mem_lt_size(offset: usize, size: usize) -> Self {
let mut rng = rand::thread_rng();
let data_len = offset
+ if size.gt(&0) {
rng.gen_range(0..size)
} else {
0
};
Self {
offset,
size,
data_len,
rng,
}
}
pub fn mem_eq_size(offset: usize, size: usize) -> Self {
let data_len = offset + size;
Self {
offset,
size,
data_len,
rng: rand::thread_rng(),
}
}
pub fn mem_gt_size(offset: usize, size: usize) -> Self {
let mut rng = rand::thread_rng();
let data_len = offset
+ size
+ if size.gt(&0) {
rng.gen_range(0..size)
} else {
0
};
Self {
offset,
size,
data_len,
rng,
}
}
pub fn mem_empty(offset: usize, size: usize) -> Self {
Self {
offset,
size,
data_len: 0,
rng: rand::thread_rng(),
}
}
fn rand_bytes(&mut self) -> Vec<u8> {
(0..self.data_len)
.map(|_| self.rng.gen())
.collect::<Vec<u8>>()
}
pub fn gen_sha3_code(&mut self) -> (Bytecode, Vec<u8>) {
let data = self.rand_bytes();
let mut memory = Vec::with_capacity(self.data_len);
let mut code = Bytecode::default();
for (i, mem_chunk) in data.chunks(32).enumerate() {
let mem_value = if mem_chunk.len() < 32 {
std::iter::repeat(0u8)
.take(32 - mem_chunk.len())
.chain(mem_chunk.to_vec())
.collect::<Vec<u8>>()
} else {
mem_chunk.to_vec()
};
memory.extend_from_slice(&mem_value);
code.op_mstore(32 * i, U256::from_big_endian(&mem_value));
}
let code_tail = bytecode! {
PUSH32(self.size)
PUSH32(self.offset)
SHA3
STOP
};
code.append(&code_tail);
(code, memory)
}
}