Skip to main content

gaia_assembler/backends/wasi/
mod.rs

1//! WASI (WebAssembly System Interface) backend compiler
2
3use super::{Backend, GeneratedFiles};
4use crate::{
5    config::GaiaConfig,
6    instruction::{CmpCondition, CoreInstruction, GaiaInstruction, ManagedInstruction},
7    program::{GaiaConstant, GaiaFunction, GaiaModule},
8    types::GaiaType,
9};
10use gaia_types::{
11    helpers::{AbiCompatible, ApiCompatible, Architecture, CompilationTarget},
12    Result,
13};
14use std::collections::HashMap;
15use wasi_assembler::{WasiFunctionType, WasiInstruction, WasiProgram, WasmValueType};
16
17/// WASI Backend implementation
18#[derive(Default)]
19pub struct WasiBackend {}
20
21impl Backend for WasiBackend {
22    fn name(&self) -> &'static str {
23        "WASI"
24    }
25
26    fn primary_target(&self) -> CompilationTarget {
27        CompilationTarget {
28            build: Architecture::WASM32,
29            host: AbiCompatible::WebAssemblyTextFormat,
30            target: ApiCompatible::WASI,
31        }
32    }
33
34    fn match_score(&self, target: &CompilationTarget) -> f32 {
35        match target.host {
36            AbiCompatible::WebAssemblyTextFormat => 10.0,
37            AbiCompatible::Unknown => match target.build {
38                // wat output, 5% support
39                Architecture::WASM32 => 5.0,
40                Architecture::WASM64 => 0.0,
41                _ => -100.0,
42            },
43            _ => -100.0,
44        }
45    }
46
47    fn generate(&self, program: &GaiaModule, _config: &GaiaConfig) -> Result<GeneratedFiles> {
48        let mut context = create_wasi_context()?;
49        compile_program(&mut context, program)?;
50
51        let wasi_program = context.program;
52
53        let mut files = HashMap::new();
54
55        // 1. 生成 .wasm
56        files.insert("main.wasm".to_string(), wasi_program.to_wasm()?);
57
58        // 2. 生成 .wat
59        // if let Ok(wat_ast) = wasi_program.to_wat() {
60        // use oak_core::source::ToSource;
61        // files.insert("main.wat".to_string(), wat_ast.to_source_string().into_bytes());
62        // }
63
64        Ok(GeneratedFiles { files, diagnostics: vec![] })
65    }
66}
67
68impl WasiBackend {
69    /// Generate WASI WebAssembly bytecode from Gaia program
70    pub fn generate(program: &GaiaModule) -> Result<Vec<u8>> {
71        let mut context = create_wasi_context()?;
72        compile_program(&mut context, program)?;
73        context.program.to_wasm()
74    }
75}
76
77/// Compile Gaia program to WASI WebAssembly
78pub fn compile(program: &GaiaModule) -> Result<Vec<u8>> {
79    WasiBackend::generate(program)
80}
81
82/// Create WASI assembler context
83fn create_wasi_context() -> Result<WasiContext> {
84    Ok(WasiContext::new())
85}
86
87/// Compile entire program
88fn compile_program(context: &mut WasiContext, program: &GaiaModule) -> Result<()> {
89    // 1. 添加必要的导入
90    // console.log 映射到 WASI fd_write 或类似的
91    // fd_write(fd: i32, iovs: i32, iovs_len: i32, nwritten: i32) -> i32
92    context.add_import(
93        "wasi_snapshot_preview1",
94        "fd_write",
95        vec![WasmValueType::I32, WasmValueType::I32, WasmValueType::I32, WasmValueType::I32],
96        vec![WasmValueType::I32],
97    );
98
99    // 2. 编译所有函数
100    for function in &program.functions {
101        compile_function(context, function)?;
102    }
103
104    // 3. 导出 main 函数
105    if !context.program.functions.is_empty() {
106        let main_index = (context.program.imports.len() + context.program.functions.len() - 1) as u32;
107        context.program.add_export(wasi_assembler::WasiExport {
108            name: "main".to_string(),
109            export_type: wasi_assembler::WasmExportType::Function { function_index: main_index },
110        });
111    }
112
113    Ok(())
114}
115
116/// Compile single function
117fn compile_function(context: &mut WasiContext, function: &GaiaFunction) -> Result<()> {
118    let mut instrs = Vec::new();
119
120    // 编译代码块
121    for block in &function.blocks {
122        // TODO: 处理标签
123
124        for instruction in &block.instructions {
125            compile_instruction(context, &mut instrs, instruction)?;
126        }
127
128        // 编译终结符
129        match &block.terminator {
130            crate::program::GaiaTerminator::Jump(_label) => instrs.push(WasiInstruction::Br { label_index: 0 }), // FIXME
131            crate::program::GaiaTerminator::Branch { .. } => {
132                instrs.push(WasiInstruction::BrIf { label_index: 0 }); // FIXME
133            }
134            crate::program::GaiaTerminator::Return => instrs.push(WasiInstruction::Return),
135            crate::program::GaiaTerminator::Call { .. } => {
136                // TODO: 查找函数索引
137                instrs.push(WasiInstruction::Call { function_index: 0 });
138            }
139            crate::program::GaiaTerminator::Halt => {
140                // WASI: exit(0)
141                instrs.push(WasiInstruction::I32Const { value: 0 });
142                // TODO: 查找 proc_exit 索引
143                instrs.push(WasiInstruction::Call { function_index: 0 });
144            }
145        }
146    }
147
148    instrs.push(WasiInstruction::End);
149
150    // 映射参数和返回值类型
151    let params = function.signature.params.iter().map(|p| map_type(p)).collect();
152    let results = match function.signature.return_type {
153        GaiaType::Void => vec![],
154        _ => vec![map_type(&function.signature.return_type)],
155    };
156
157    // 添加函数到程序
158    context.add_function(instrs, params, results);
159
160    Ok(())
161}
162
163fn map_type(ty: &GaiaType) -> WasmValueType {
164    match ty {
165        GaiaType::I32 => WasmValueType::I32,
166        GaiaType::I64 => WasmValueType::I64,
167        GaiaType::F32 => WasmValueType::F32,
168        GaiaType::F64 => WasmValueType::F64,
169        _ => WasmValueType::I32, // Fallback
170    }
171}
172
173/// Compile single instruction
174fn compile_instruction(
175    _context: &mut WasiContext,
176    instrs: &mut Vec<WasiInstruction>,
177    instruction: &GaiaInstruction,
178) -> Result<()> {
179    match instruction {
180        GaiaInstruction::Core(core) => match core {
181            CoreInstruction::PushConstant(constant) => compile_load_constant(instrs, constant),
182            CoreInstruction::Load(gaia_type) => compile_load_indirect(instrs, gaia_type),
183            CoreInstruction::Store(gaia_type) => compile_store_indirect(instrs, gaia_type),
184            CoreInstruction::Add(_) => {
185                instrs.push(WasiInstruction::I32Add);
186                Ok(())
187            }
188            CoreInstruction::Sub(_) => {
189                instrs.push(WasiInstruction::I32Sub);
190                Ok(())
191            }
192            CoreInstruction::Mul(_) => {
193                instrs.push(WasiInstruction::I32Mul);
194                Ok(())
195            }
196            CoreInstruction::Div(_) => {
197                instrs.push(WasiInstruction::I32DivS);
198                Ok(())
199            }
200            CoreInstruction::Rem(_) => {
201                instrs.push(WasiInstruction::I32RemS);
202                Ok(())
203            }
204            CoreInstruction::And(_) => {
205                instrs.push(WasiInstruction::I32And);
206                Ok(())
207            }
208            CoreInstruction::Or(_) => {
209                instrs.push(WasiInstruction::I32Or);
210                Ok(())
211            }
212            CoreInstruction::Xor(_) => {
213                instrs.push(WasiInstruction::I32Xor);
214                Ok(())
215            }
216            CoreInstruction::Shl(_) => {
217                instrs.push(WasiInstruction::I32Shl);
218                Ok(())
219            }
220            CoreInstruction::Shr(_) => {
221                instrs.push(WasiInstruction::I32ShrS);
222                Ok(())
223            }
224            CoreInstruction::Pop => {
225                instrs.push(WasiInstruction::Drop);
226                Ok(())
227            }
228            CoreInstruction::LoadLocal(index, _) => {
229                instrs.push(WasiInstruction::LocalGet { local_index: *index });
230                Ok(())
231            }
232            CoreInstruction::StoreLocal(index, _) => {
233                instrs.push(WasiInstruction::LocalSet { local_index: *index });
234                Ok(())
235            }
236            CoreInstruction::LoadArg(index, _) => {
237                instrs.push(WasiInstruction::LocalGet { local_index: *index });
238                Ok(())
239            }
240            CoreInstruction::StoreArg(index, _) => {
241                instrs.push(WasiInstruction::LocalSet { local_index: *index });
242                Ok(())
243            }
244            CoreInstruction::Ret => {
245                instrs.push(WasiInstruction::Return);
246                Ok(())
247            }
248            CoreInstruction::Call(_, _) => {
249                // TODO: 查找函数索引
250                instrs.push(WasiInstruction::Call { function_index: 0 });
251                Ok(())
252            }
253            CoreInstruction::Cmp(cond, _) => compile_compare(instrs, cond),
254            CoreInstruction::StructNew(_) => {
255                // TODO: 查找类型索引
256                instrs.push(WasiInstruction::StructNew { type_index: 0 });
257                Ok(())
258            }
259            CoreInstruction::StructGet { field_index, .. } => {
260                // TODO: 查找类型索引
261                instrs.push(WasiInstruction::StructGet { type_index: 0, field_index: *field_index });
262                Ok(())
263            }
264            CoreInstruction::StructSet { field_index, .. } => {
265                // TODO: 查找类型索引
266                instrs.push(WasiInstruction::StructSet { type_index: 0, field_index: *field_index });
267                Ok(())
268            }
269            _ => Ok(()),
270        },
271        GaiaInstruction::Managed(managed) => match managed {
272            ManagedInstruction::CallStatic { .. } => {
273                // AOT 模式下静态方法调用映射为普通 call
274                instrs.push(WasiInstruction::Call { function_index: 0 });
275                Ok(())
276            }
277            _ => Ok(()),
278        },
279        _ => Ok(()),
280    }
281}
282
283fn compile_compare(instrs: &mut Vec<WasiInstruction>, cond: &CmpCondition) -> Result<()> {
284    match cond {
285        CmpCondition::Eq => instrs.push(WasiInstruction::I32Eq),
286        CmpCondition::Ne => instrs.push(WasiInstruction::I32Ne),
287        CmpCondition::Lt => instrs.push(WasiInstruction::I32LtS),
288        CmpCondition::Le => instrs.push(WasiInstruction::I32LeS),
289        CmpCondition::Gt => instrs.push(WasiInstruction::I32GtS),
290        CmpCondition::Ge => instrs.push(WasiInstruction::I32GeS),
291    }
292    Ok(())
293}
294
295fn compile_load_indirect(instrs: &mut Vec<WasiInstruction>, gaia_type: &GaiaType) -> Result<()> {
296    // 假设内存偏移已经在栈顶
297    match gaia_type {
298        GaiaType::I32 => instrs.push(WasiInstruction::I32Load { offset: 0, align: 0 }),
299        GaiaType::I64 => instrs.push(WasiInstruction::I64Load { offset: 0, align: 0 }),
300        GaiaType::F32 => instrs.push(WasiInstruction::F32Load { offset: 0, align: 0 }),
301        GaiaType::F64 => instrs.push(WasiInstruction::F64Load { offset: 0, align: 0 }),
302        _ => {}
303    }
304    Ok(())
305}
306
307fn compile_store_indirect(instrs: &mut Vec<WasiInstruction>, gaia_type: &GaiaType) -> Result<()> {
308    // 假设 [ptr, value] 已经在栈上
309    match gaia_type {
310        GaiaType::I32 => instrs.push(WasiInstruction::I32Store { offset: 0, align: 0 }),
311        GaiaType::I64 => instrs.push(WasiInstruction::I64Store { offset: 0, align: 0 }),
312        GaiaType::F32 => instrs.push(WasiInstruction::F32Store { offset: 0, align: 0 }),
313        GaiaType::F64 => instrs.push(WasiInstruction::F64Store { offset: 0, align: 0 }),
314        _ => {}
315    }
316    Ok(())
317}
318
319fn compile_load_constant(instrs: &mut Vec<WasiInstruction>, constant: &GaiaConstant) -> Result<()> {
320    match constant {
321        GaiaConstant::I32(value) => instrs.push(WasiInstruction::I32Const { value: *value }),
322        GaiaConstant::I64(value) => instrs.push(WasiInstruction::I64Const { value: *value }),
323        GaiaConstant::F32(value) => instrs.push(WasiInstruction::F32Const { value: *value }),
324        GaiaConstant::F64(value) => instrs.push(WasiInstruction::F64Const { value: *value }),
325        _ => {}
326    }
327    Ok(())
328}
329
330/// WASI assembler context
331struct WasiContext {
332    program: WasiProgram,
333    /// 导入映射
334    import_map: HashMap<String, u32>,
335}
336
337impl WasiContext {
338    fn new() -> Self {
339        WasiContext { program: WasiProgram::new_core_module(), import_map: HashMap::new() }
340    }
341
342    fn add_import(&mut self, module: &str, field: &str, params: Vec<WasmValueType>, results: Vec<WasmValueType>) -> u32 {
343        let key = format!("{}.{}", module, field);
344        if let Some(&index) = self.import_map.get(&key) {
345            index
346        }
347        else {
348            let index = self.import_map.len() as u32;
349            let func_type = WasiFunctionType { params, results };
350            let type_index = self.program.add_function_type(func_type);
351
352            self.program.add_import(wasi_assembler::WasiImport {
353                module: module.to_string(),
354                field: field.to_string(),
355                import_type: wasi_assembler::WasmImportType::Function { type_index },
356            });
357
358            self.import_map.insert(key, index);
359            index
360        }
361    }
362
363    fn add_function(&mut self, body: Vec<WasiInstruction>, params: Vec<WasmValueType>, results: Vec<WasmValueType>) {
364        let func_type = WasiFunctionType { params, results };
365        let type_index = self.program.add_function_type(func_type);
366        let func = wasi_assembler::WasiFunction { type_index, locals: vec![], body };
367        self.program.add_function(func);
368    }
369}