Skip to main content

thalir_core/codegen/
module.rs

1use cranelift::prelude::EntityRef;
2use cranelift_codegen::ir::Function;
3use cranelift_codegen::isa;
4use cranelift_codegen::settings;
5use cranelift_codegen::Context;
6use cranelift_frontend::Variable;
7use cranelift_module::Module;
8use cranelift_object::{ObjectBuilder, ObjectModule};
9use std::collections::HashMap;
10
11use crate::{
12    codegen::context::CodegenContext,
13    codegen::lowering::{lower_instruction, lower_terminator},
14    contract::Contract,
15    values::VarId,
16    IrError, Result,
17};
18
19pub struct ModuleBuilder {
20    module: ObjectModule,
21}
22
23impl ModuleBuilder {
24    pub fn new() -> Result<Self> {
25        let flags_builder = settings::builder();
26        let isa_builder = isa::lookup_by_name("x86_64-unknown-unknown-elf")
27            .map_err(|e| IrError::CraneliftError(format!("Failed to lookup ISA: {}", e)))?;
28
29        let flags = settings::Flags::new(flags_builder);
30        let isa = isa_builder
31            .finish(flags)
32            .map_err(|e| IrError::CraneliftError(format!("Failed to create ISA: {}", e)))?;
33
34        let object_builder = ObjectBuilder::new(
35            isa.clone(),
36            "contract",
37            cranelift_module::default_libcall_names(),
38        )
39        .unwrap();
40        let module = ObjectModule::new(object_builder);
41
42        Ok(Self { module })
43    }
44
45    pub fn compile_contract(mut self, contract: &Contract) -> Result<Vec<u8>> {
46        let mut func_ids = HashMap::new();
47
48        for (name, function) in &contract.functions {
49            let mut sig = self.module.make_signature();
50            for param in &function.signature.params {
51                sig.params.push(cranelift_codegen::ir::AbiParam::new(
52                    param.param_type.to_cranelift().unwrap(),
53                ));
54            }
55            for ret in &function.signature.returns {
56                sig.returns.push(cranelift_codegen::ir::AbiParam::new(
57                    ret.to_cranelift().unwrap(),
58                ));
59            }
60
61            let func_id = self
62                .module
63                .declare_function(name, cranelift_module::Linkage::Export, &sig)
64                .map_err(|e| {
65                    IrError::CraneliftError(format!("Failed to declare function: {}", e))
66                })?;
67            func_ids.insert(name.clone(), func_id);
68        }
69
70        for (name, function) in &contract.functions {
71            let func_id = func_ids.get(name).unwrap();
72            let mut clif_func = Function::new();
73            clif_func.signature = self
74                .module
75                .declarations()
76                .get_function_decl(*func_id)
77                .signature
78                .clone();
79
80            let mut ctx = CodegenContext::new(&mut clif_func);
81            let mut func_builder = ctx.func_builder();
82
83            let mut block_map = HashMap::new();
84            let mut variables = HashMap::new();
85            let mut ssa_values = HashMap::new();
86
87            for (block_id, _) in &function.body.blocks {
88                let clif_block = func_builder.create_block();
89                block_map.insert(*block_id, clif_block);
90            }
91
92            let entry_clif_block = block_map.get(&function.body.entry_block).unwrap();
93            func_builder.append_block_params_for_function_params(*entry_clif_block);
94            func_builder.switch_to_block(*entry_clif_block);
95
96            for (i, param) in function.signature.params.iter().enumerate() {
97                let var = Variable::new(i);
98                func_builder.declare_var(var, param.param_type.to_cranelift().unwrap());
99                let val = func_builder.block_params(*entry_clif_block)[i];
100                func_builder.def_var(var, val);
101                variables.insert(VarId(i as u32), var);
102                ssa_values.insert(
103                    crate::values::Value::Param(crate::values::ParamId(i as u32)),
104                    val,
105                );
106            }
107
108            for (block_id, block) in &function.body.blocks {
109                let clif_block = *block_map.get(block_id).unwrap();
110                func_builder.switch_to_block(clif_block);
111
112                for inst in &block.instructions {
113                    lower_instruction(inst, &variables, &mut ssa_values, &mut func_builder)?;
114                }
115
116                if !matches!(block.terminator, crate::block::Terminator::Invalid) {
117                    lower_terminator(
118                        &block.terminator,
119                        &ssa_values,
120                        &mut func_builder,
121                        &block_map,
122                    )?;
123                }
124            }
125
126            func_builder.seal_all_blocks();
127            func_builder.finalize();
128
129            let mut context = Context::for_function(clif_func);
130            self.module
131                .define_function(*func_id, &mut context)
132                .map_err(|e| {
133                    IrError::CraneliftError(format!("Failed to define function: {}", e))
134                })?;
135        }
136
137        let product = self.module.finish();
138        let obj_bytes = product
139            .emit()
140            .map_err(|e| IrError::CraneliftError(format!("Failed to emit object: {}", e)))?;
141
142        Ok(obj_bytes)
143    }
144}