zingen/codegen/
function.rs

1//! Code generation implementation.
2use crate::{
3    backtrace::Backtrace,
4    control::ControlStack,
5    jump::JumpTable,
6    local::{LocalSlot, LocalSlotType, Locals},
7    masm::MacroAssembler,
8    validator::ValidateThenVisit,
9    wasm::Env,
10    Buffer, Error, Result,
11};
12use opcodes::ShangHai as OpCode;
13use wasmparser::{FuncType, FuncValidator, LocalsReader, OperatorsReader, ValidatorResources};
14use zabi::Abi;
15
16/// The code generation abstraction.
17pub struct Function {
18    /// Abi of this function,
19    pub abi: Option<Abi>,
20    /// The backtrace.
21    pub backtrace: Backtrace,
22    /// Control stack frames.
23    pub control: ControlStack,
24    /// WASM environment.
25    pub env: Env,
26    /// The defined locals for a function.
27    pub locals: Locals,
28    /// The macro assembler.
29    pub masm: MacroAssembler,
30    /// The jump table.
31    pub table: JumpTable,
32    /// The function type.
33    pub ty: FuncType,
34    /// If this function is the main function.
35    pub is_main: bool,
36}
37
38impl Function {
39    /// Create a new code generator.
40    pub fn new(env: Env, ty: FuncType, abi: Option<Abi>, is_main: bool) -> Result<Self> {
41        let is_external = abi.is_some();
42        let mut codegen = Self {
43            abi,
44            backtrace: Backtrace::default(),
45            control: ControlStack::default(),
46            env,
47            ty,
48            locals: Default::default(),
49            masm: Default::default(),
50            table: Default::default(),
51            is_main,
52        };
53
54        if is_main {
55            return Ok(codegen);
56        }
57
58        // post process program counter and stack pointer.
59        if is_external {
60            // codegen.masm.increment_sp(1)?;
61            tracing::debug!("<External function>");
62            codegen.masm._jumpdest()?;
63        } else {
64            // Mock the stack frame for the callee function
65            //
66            // STACK: [ PC ]
67            tracing::debug!("<Internal function>");
68            codegen.masm.increment_sp(1)?;
69            codegen.masm._jumpdest()?;
70        }
71
72        Ok(codegen)
73    }
74
75    /// Emit function locals
76    ///
77    /// 1. the function parameters.
78    /// 2. function body locals.
79    ///
80    /// NOTE: we don't care about the original offset of the locals.
81    /// bcz we will serialize the locals to an index map anyway.
82    pub fn emit_locals(
83        &mut self,
84        locals: &mut LocalsReader<'_>,
85        validator: &mut FuncValidator<ValidatorResources>,
86    ) -> Result<()> {
87        let mut sp = if self.is_main { 0 } else { 1 };
88
89        // Define locals in function parameters.
90        for param in self.ty.params() {
91            self.locals
92                .push(LocalSlot::new(*param, LocalSlotType::Parameter, sp));
93            sp += 1;
94        }
95
96        // Define locals in function body.
97        //
98        // Record the offset for validation.
99        while let Ok((count, val)) = locals.read() {
100            for _ in 0..count {
101                // Define locals.
102                self.locals
103                    .push(LocalSlot::new(val, LocalSlotType::Variable, sp));
104
105                sp += 1;
106            }
107
108            let validation_offset = locals.original_position();
109            validator.define_locals(validation_offset, count, val)?;
110        }
111
112        tracing::trace!("{:?}", self.locals);
113        Ok(())
114    }
115
116    /// Emit function operators.
117    pub fn emit_operators(
118        &mut self,
119        ops: &mut OperatorsReader<'_>,
120        validator: &mut FuncValidator<ValidatorResources>,
121    ) -> Result<()> {
122        while !ops.eof() {
123            let offset = ops.original_position();
124            let mut validate_then_visit = ValidateThenVisit(validator.visitor(offset), self);
125            ops.visit_operator(&mut validate_then_visit)???;
126        }
127
128        if (self.abi.is_some() || self.is_main)
129            && self.masm.buffer().last() != Some(&OpCode::RETURN.into())
130        {
131            self._end()?;
132        }
133
134        Ok(())
135    }
136
137    /// Finish code generation.
138    pub fn finish(self, jump_table: &mut JumpTable, pc: u16) -> Result<Buffer> {
139        let sp = self.masm.sp();
140        if !self.is_main && self.abi.is_none() && self.masm.sp() != self.ty.results().len() as u16 {
141            return Err(Error::StackNotBalanced(sp));
142        }
143
144        jump_table.merge(self.table, pc)?;
145        Ok(self.masm.buffer().into())
146    }
147}