dfwasm_compiler/
compiler.rs

1use std::{rc::Rc, vec};
2
3use dfwasm_template::{split_templates, Args, Block, Item, Template};
4use wasmparser::{ConstExpr, MemoryType, Operator, Parser, RecGroup};
5
6use crate::{
7    df_helper::{format_df_number_u64, var, DF_VAR_MEM_SIZE},
8    DFWasmError, DFWasmResult,
9};
10
11use super::{
12    df_helper::{format_df_number_i64, num, DF_FUNC_CALL_FUNC},
13    sections::compile_section,
14};
15
16/// An entry in the control stack when compiling.
17#[derive(Debug, Clone, PartialEq)]
18pub(crate) enum ControlStackEntry {
19    FunctionStart(usize),
20    Block(usize),
21    Loop(usize),
22    If(usize),
23    Else(usize),
24}
25
26/// The options for [`DFWasmCompiler`].
27#[derive(Default, Debug, Clone)]
28pub struct DFWasmCompilerOptions {
29    /// The module name. Useful for when compiling multiple modules.
30    pub module_name: Option<String>,
31    /// Whether or not to include debugger function calls.
32    pub debugger: bool,
33    /// Whether or not to call the debugger on nops.
34    pub skip_nop_debugger: bool,
35    /// Whether or not to limit the size of the templates (highly recommended).
36    /// Use this to set your plot size.
37    pub max_template_size: Option<usize>,
38    /// Configuration for batching data section memory setters. (i.e. combining multiple
39    /// set vars into one function call).
40    ///
41    /// If Some, enables batching with the specified number of bytes per batch.
42    /// Default max is 26 bytes, but 8 may be safer for massive plots.
43    ///
44    /// Using None disables batching.
45    pub batch_data_size: Option<usize>,
46    /// Whether or not to only include the module init function in the template.
47    /// Useful for debugging data section initialization.
48    pub only_include_module_init: bool,
49}
50
51/// The WASM to DiamondFire compiler.
52pub struct DFWasmCompiler<'a> {
53    /// The binary WASM file to compile
54    pub wasm: &'a [u8],
55    /// The options to use when compiling
56    pub options: DFWasmCompilerOptions,
57
58    /// The exported templates (also used as a working area for the compiler)
59    pub(crate) templates: Vec<Template>,
60
61    /// The currently defined memory type, if any.
62    pub(crate) memory_type: Option<MemoryType>,
63
64    /// The index of the start method, which is called when the module is loaded.
65    pub(crate) start_method: Option<usize>,
66
67    /// The stack of control flow operators (ifs, loops, blocks, etc)
68    pub(crate) control_stack: Vec<ControlStackEntry>,
69
70    /// A map front table index to the initialization expression for that table.
71    pub(crate) table_to_init_expr: Vec<Item>,
72
73    /// The current function being compiled. Also increased on imports.
74    pub(crate) function_counter: usize,
75
76    /// The current table being compiled. Also increased on imports.
77    pub(crate) table_counter: usize,
78
79    /// A map from function index to function signature.
80    pub(crate) function_to_type_signature: Vec<Rc<RecGroup>>,
81
82    /// The defined function signatures,
83    pub(crate) function_signatures: Vec<Rc<RecGroup>>,
84}
85
86impl<'a> DFWasmCompiler<'a> {
87    /// Compiles a WASM file into DiamondFire templates.
88    pub fn compile_wasm(
89        wasm: &'a [u8],
90        options: DFWasmCompilerOptions,
91    ) -> DFWasmResult<Vec<Template>> {
92        let this = Self {
93            wasm,
94            options,
95            templates: Vec::new(),
96            memory_type: None,
97            start_method: None,
98            control_stack: Vec::new(),
99            table_to_init_expr: Vec::new(),
100            function_counter: 0,
101            table_counter: 0,
102            function_to_type_signature: Vec::new(),
103            function_signatures: Vec::new(),
104        };
105        this.compile()
106    }
107
108    /// The main compile function. Goes through each section of the WASM file and compiles it into a template.
109    fn compile(mut self) -> DFWasmResult<Vec<Template>> {
110        // Create a new binary WASM parser (0 represents source offset)
111        let parser = Parser::new(0);
112
113        // Module template is what is called to initialize the module
114        let module_template_name = match &self.options.module_name {
115            Some(name) => format!("wasm.{}.init", name),
116            None => "wasm.module.init".to_string(),
117        };
118        let mut module_template = Template::start_function(module_template_name.clone());
119
120        // Iterate over every section in the WASM file
121        for section in parser.parse_all(self.wasm) {
122            let section = section?;
123            compile_section(&mut self, &mut module_template, section)?;
124        }
125
126        // Add the start method to the module template if it exists
127        if let Some(start_method) = self.start_method {
128            module_template.blocks.push(Block::CallFunction {
129                args: Args::with(vec![num(start_method), num(0), num(0)]),
130                func: DF_FUNC_CALL_FUNC.to_string(),
131            });
132        }
133
134        // Set the starting memory size
135        if let Some(memory_type) = self.memory_type {
136            module_template.set_var(
137                "=",
138                Args::with(vec![
139                    var(DF_VAR_MEM_SIZE),
140                    num(format_df_number_u64(memory_type.initial)),
141                ]),
142            );
143        }
144
145        if self.options.only_include_module_init {
146            self.templates.clear();
147        }
148
149        self.templates.push(module_template);
150
151        // Split the templates if they are too large
152        if let Some(max_template_size) = self.options.max_template_size {
153            Ok(split_templates(self.templates, max_template_size))
154        } else {
155            Ok(self.templates)
156        }
157    }
158
159    /// Get the current working template index.
160    pub(crate) fn get_current_template_idx(&self) -> Option<usize> {
161        self.control_stack.last().map(|entry| match entry {
162            ControlStackEntry::FunctionStart(template_idx)
163            | ControlStackEntry::Block(template_idx)
164            | ControlStackEntry::Loop(template_idx)
165            | ControlStackEntry::If(template_idx)
166            | ControlStackEntry::Else(template_idx) => *template_idx,
167        })
168    }
169
170    // Get the current working template.
171    pub(crate) fn get_current_template(&mut self) -> &mut Template {
172        let template_idx = self.get_current_template_idx().expect("No template");
173
174        self.templates
175            .get_mut(template_idx)
176            .expect("Template not found")
177    }
178
179    /// Evaluates a constant expression and returns the result as an Item.
180    /// As opposed to [`eval_const_expr_as_offset`], this is used for things like
181    /// table and global initializers.
182    pub(crate) fn eval_const_expr(&self, expr: &ConstExpr) -> DFWasmResult<Item> {
183        let mut reader = expr.get_operators_reader();
184        let op = reader.read()?;
185
186        let result = match op {
187            Operator::I32Const { value } => num(format_df_number_i64(value as i64)),
188            Operator::I64Const { value } => num(format_df_number_i64(value)),
189            Operator::F32Const { value: _ } => todo!("float constant expressions"),
190            Operator::F64Const { value: _ } => todo!("float constant expressions"),
191            Operator::RefNull { hty: _ } => num(-1),
192            Operator::RefFunc { function_index } => num(function_index),
193            Operator::GlobalGet { global_index: _ } => todo!("global get within a const expr"),
194            _ => return Err(DFWasmError::UnsupportedConstExpr),
195        };
196
197        if reader.eof() {
198            Ok(result)
199        } else {
200            match reader.read()? {
201                Operator::End => Ok(result),
202                // todo: should this be a separate error?
203                _ => Err(DFWasmError::UnsupportedConstExpr),
204            }
205        }
206    }
207
208    /// Evaluates a constant expression as an offset.
209    /// This is used when a compile-time known offset is needed, such as memory data section offsets,
210    /// or element section offsets.
211    pub(crate) fn eval_const_expr_as_offset(&self, expr: &ConstExpr) -> DFWasmResult<usize> {
212        let mut reader = expr.get_operators_reader();
213        let op = reader.read()?;
214
215        let result = match op {
216            Operator::I32Const { value } => usize::from_le_bytes((value as i64).to_le_bytes()),
217            Operator::I64Const { value } => usize::from_le_bytes(value.to_le_bytes()),
218            Operator::GlobalGet { global_index: _ } => todo!("global get within a const expr"),
219            _ => return Err(DFWasmError::UnsupportedConstExpr),
220        };
221
222        if reader.eof() {
223            Ok(result)
224        } else {
225            match reader.read()? {
226                Operator::End => Ok(result),
227                // todo: should this be a separate error?
228                _ => Err(DFWasmError::UnsupportedConstExpr),
229            }
230        }
231    }
232
233    /// Get the argument count of a function type.
234    pub(crate) fn arg_count_of_type(signature: &RecGroup) -> Option<usize> {
235        Some(signature.types().next()?.unwrap_func().params().len())
236    }
237
238    /// Get the result count of a function type.
239    /// In WASM, a function can have multiple results.
240    pub(crate) fn result_count_of_type(signature: &RecGroup) -> Option<usize> {
241        Some(signature.types().next()?.unwrap_func().results().len())
242    }
243}
dfwasm_compiler/compiler.rs

dfwasm_compiler/
compiler.rs
