zingen/codegen/
code.rs

1//! Table for the code section.
2
3use indexmap::IndexMap;
4
5/// Code section for EVM.
6#[derive(Clone, Default, Debug)]
7pub struct Code {
8    /// The offset of the code section
9    offset: usize,
10    /// Function table.
11    funcs: IndexMap<ExtFunc, usize>,
12}
13
14impl Code {
15    /// Create a new code section.
16    pub fn new() -> Self {
17        Self {
18            offset: 0,
19            funcs: Default::default(),
20        }
21    }
22
23    /// Get the functions in the code section.
24    pub fn funcs(&self) -> Vec<ExtFunc> {
25        self.funcs.keys().cloned().collect()
26    }
27
28    /// Shift the code section.
29    pub fn shift(&mut self, offset: u16) {
30        tracing::trace!("shift code section by 0x{:x} bytes.", offset);
31        let offset = offset as usize;
32        self.offset += offset;
33        self.funcs.values_mut().for_each(|pc| *pc += offset);
34    }
35
36    /// Add a function to the code section.
37    pub fn try_add_func(&mut self, func: ExtFunc) {
38        if self.funcs.contains_key(&func) {
39            return;
40        }
41
42        let bytecode = func.bytecode.clone();
43        let len = bytecode.len();
44        self.funcs.insert(func, self.offset);
45        self.offset += len;
46    }
47
48    /// Get the current offset of the code section.
49    pub fn offset(&self) -> usize {
50        self.offset
51    }
52
53    /// Get the offset of a function.
54    pub fn offset_of(&self, func: &ExtFunc) -> Option<u16> {
55        self.funcs.get(func).and_then(|i| (*i).try_into().ok())
56    }
57
58    /// Get the bytecode of the code section.
59    pub fn finish(&self) -> Vec<u8> {
60        let mut code = Vec::new();
61        for func in self.funcs.keys() {
62            tracing::trace!("add function to code section: {:?}", func);
63            code.extend(func.bytecode.clone());
64        }
65        code
66    }
67}
68
69/// External function in code section.
70#[derive(PartialEq, Eq, Debug, Clone, Hash)]
71pub struct ExtFunc {
72    /// Stack input.
73    pub stack_out: u8,
74    /// Stack output.
75    pub stack_in: u8,
76    /// The bytecode of the external function.
77    pub bytecode: Vec<u8>,
78}