1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

//! Precompiles support.

use primitive_types::U256;
pub use zkevm_opcode_defs::system_params::{
    ECRECOVER_INNER_FUNCTION_PRECOMPILE_ADDRESS, KECCAK256_ROUND_FUNCTION_PRECOMPILE_ADDRESS,
    SECP256R1_VERIFY_PRECOMPILE_ADDRESS, SHA256_ROUND_FUNCTION_PRECOMPILE_ADDRESS,
};
use zksync_vm2_interface::CycleStats;

pub use self::legacy::LegacyPrecompiles;
use crate::heap::Heap;

mod legacy;

/// Provides access to the input memory for a precompile call.
#[derive(Debug, Clone)]
pub struct PrecompileMemoryReader<'a, const IN_WORDS: bool = false> {
    heap: &'a Heap,
    offset: u32,
    len: u32,
}

impl<'a> PrecompileMemoryReader<'a> {
    pub(crate) fn new(heap: &'a Heap, offset: u32, len: u32) -> Self {
        Self { heap, offset, len }
    }

    /// Assumes that the input offset and length passed via ABI are measured in 32-byte words, rather than bytes.
    pub fn assume_offset_in_words(self) -> PrecompileMemoryReader<'a, true> {
        PrecompileMemoryReader {
            heap: self.heap,
            offset: self.offset * 32,
            len: self.len * 32,
        }
    }
}

/// Iterates over input bytes.
impl<const IN_WORDS: bool> Iterator for PrecompileMemoryReader<'_, IN_WORDS> {
    type Item = u8;

    fn next(&mut self) -> Option<Self::Item> {
        if self.len == 0 {
            return None;
        }

        // This assumes the offset never overflows
        let output = self.heap.read_byte(self.offset);
        self.offset += 1;
        self.len -= 1;
        Some(output)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.len as usize, Some(self.len as usize))
    }
}

impl ExactSizeIterator for PrecompileMemoryReader<'_> {
    fn len(&self) -> usize {
        self.len as usize
    }
}

/// Output of a precompile call returned from [`Precompiles::call_precompile()`].
#[derive(Debug, Default)]
pub struct PrecompileOutput {
    pub(crate) buffer: [U256; 2],
    pub(crate) len: u32,
    pub(crate) cycle_stats: Option<CycleStats>,
}

impl PrecompileOutput {
    /// Assigns cycle stats for this output.
    #[must_use]
    pub fn with_cycle_stats(mut self, stats: CycleStats) -> Self {
        self.cycle_stats = Some(stats);
        self
    }
}

impl From<U256> for PrecompileOutput {
    fn from(value: U256) -> Self {
        Self {
            buffer: [value, U256::zero()],
            len: 1,
            cycle_stats: None,
        }
    }
}

impl From<[U256; 2]> for PrecompileOutput {
    fn from(value: [U256; 2]) -> Self {
        Self {
            buffer: value,
            len: 2,
            cycle_stats: None,
        }
    }
}

/// Encapsulates precompiles used during VM execution.
pub trait Precompiles {
    /// Calls to a precompile.
    fn call_precompile(
        &self,
        address_low: u16,
        memory: PrecompileMemoryReader<'_>,
        aux_input: u64,
    ) -> PrecompileOutput;
}