petr_codegen/
lib.rs

1#![allow(dead_code, unused_variables)]
2use std::process::Command;
3
4use cranelift::{
5    codegen::settings::{self, Configurable},
6    frontend::{FunctionBuilder, FunctionBuilderContext},
7};
8use cranelift_module::{DataId, Linkage, Module};
9use cranelift_native::builder as isa_builder;
10use cranelift_object::{object::write::Object, ObjectBuilder, ObjectModule};
11use error::{IrError, IrErrorKind};
12
13mod error {
14    use miette::Diagnostic;
15    use thiserror::Error;
16
17    #[derive(Debug, Error, Diagnostic)]
18    pub struct IrError {
19        kind: IrErrorKind,
20    }
21
22    impl std::fmt::Display for IrError {
23        fn fmt(
24            &self,
25            f: &mut std::fmt::Formatter<'_>,
26        ) -> std::fmt::Result {
27            write!(f, "{:?}", self.kind)
28        }
29    }
30
31    #[derive(Debug, Error, Diagnostic)]
32    pub enum IrErrorKind {
33        #[error("petr-IR error: {0}")]
34        GenericIRError(String),
35        #[error(transparent)]
36        ModuleError(#[from] cranelift_module::ModuleError),
37    }
38
39    impl<T: Into<IrErrorKind>> From<T> for IrError {
40        fn from(value: T) -> Self {
41            Self { kind: value.into() }
42        }
43    }
44}
45pub struct IrContext {
46    module: ObjectModule,
47}
48
49impl IrContext {
50    // TODO: this is probably aarch64/mac specific.
51    // Need to make this platform independent. but it might already be.
52    pub fn new(file_name: &str) -> Result<Self, IrError> {
53        // Set up the ISA for the current machine
54        let mut flag_builder = settings::builder();
55        flag_builder.enable("is_pic").unwrap();
56
57        let isa_builder = isa_builder().map_err(ToString::to_string).map_err(IrErrorKind::GenericIRError)?;
58        let isa = isa_builder.finish(settings::Flags::new(flag_builder)).unwrap();
59
60        let builder = ObjectBuilder::new(isa, file_name, cranelift_module::default_libcall_names()).expect("TODO");
61
62        let module = ObjectModule::new(builder);
63        Ok(Self { module })
64    }
65
66    pub fn insert_data(
67        &mut self,
68        data: Box<[u8]>,
69        data_name: &str,
70        is_writable: bool,
71    ) -> Result<DataId, IrError> {
72        // Define the data section with "Hello, world!\n"
73        let mut data_ctx = cranelift_module::DataDescription::new();
74        data_ctx.define(data);
75        let data_id = self.module.declare_data(data_name, Linkage::Local, is_writable, false)?;
76        self.module.define_data(data_id, &data_ctx).expect("TODO");
77        Ok(data_id)
78    }
79
80    pub fn add_function(
81        &mut self,
82        func_name: &str,
83        function: petr_ir::Function,
84    ) -> Result<(), IrError> {
85        let sig = self.module.make_signature();
86        let func_id = self
87            .module
88            .declare_function(func_name, Linkage::Local, &sig)
89            .map_err(|e| IrErrorKind::GenericIRError(e.to_string()))?;
90        // func_sig
91        //     .returns
92        //     .push(AbiParam::new(cranelift::codegen::ir::types::I32));
93        let mut ctx = self.module.make_context();
94        // TODO: diff between ctx.func.signature and builder.func.signature?
95        ctx.func.signature = sig.clone();
96        let mut builder_ctx = FunctionBuilderContext::new();
97        let mut builder = FunctionBuilder::new(&mut ctx.func, &mut builder_ctx);
98        builder.func.signature = sig;
99
100        self.lower_function_body(&function, &mut builder).expect("TODO err");
101        self.module
102            .define_function(func_id, &mut ctx)
103            .map_err(|e| IrErrorKind::GenericIRError(e.to_string()))?;
104        self.module.clear_context(&mut ctx);
105
106        // module.define_function(func_id, &mut ctx
107
108        Ok(())
109    }
110
111    pub fn into_object_file(self) -> Result<Object<'static>, IrError> {
112        let product = self.module.finish();
113        let obj = product.object;
114        Ok(obj)
115    }
116
117    pub fn add_main_function(&mut self) -> Result<(), IrError> {
118        // let func_id = module.declare_function("main", Linkage::Export, &func_sig)?;
119        todo!()
120    }
121
122    fn lower_function_body(
123        &self,
124        _function: &petr_ir::Function,
125        _builder: &mut FunctionBuilder,
126    ) -> Result<(), IrError> {
127        todo!()
128        // let entry_block = builder.create_block();
129        // builder.switch_to_block(entry_block);
130        // builder.seal_block(entry_block);
131        // let mut puts_sig = module.make_signature();
132        // puts_sig.params.push(AbiParam::new(types::I64)); // pointer to char as parameter
133        // puts_sig.returns.push(AbiParam::new(types::I32)); // return type of `puts`
134        // let puts_func = module.declare_function("puts", Linkage::Import, &puts_sig)?;
135        // // get func ref for func id
136        // let puts_func_ref = module.declare_func_in_func(puts_func, builder.func);
137
138        // // Reference the data section
139        // let data_ref = module.declare_data_in_func(data_id, builder.func);
140        // let base_addr = builder.ins().global_value(types::I64, data_ref);
141
142        // // Call `puts` with the string address
143        // let call = builder.ins().call(puts_func_ref, &[base_addr]);
144
145        // // Return from the function
146        // let ret_val = builder.inst_results(call)[0];
147        // builder.ins().return_(&[ret_val]);
148
149        // builder.finalize();
150    }
151}
152
153fn write_obj_file(
154    file_name: &str,
155    obj: Object,
156) -> Result<(), IrError> {
157    let mut file = std::fs::File::create(file_name).expect("TODO errs");
158    use std::io::Write;
159
160    file.write_all(&obj.write().unwrap()).expect("TODO errs");
161    Ok(())
162}
163
164fn link_for_mac(
165    obj_file_name: &str,
166    output_file_name: &str,
167) -> Result<(), IrError> {
168    // Link the object file using clang
169    Command::new("clang")
170        .arg("output.o")
171        .arg("-o")
172        .arg("output")
173        .arg("-Wl")
174        .arg("-ld_classic")
175        .arg("-v")
176        .status()
177        .expect("TODO errs");
178    use std::{fs, os::unix::fs::PermissionsExt};
179
180    // Set the output file to be executable
181    let mut perms = fs::metadata("output").expect("TODO errs").permissions();
182    perms.set_mode(0o755);
183    fs::set_permissions("output", perms).expect("TODO errs");
184    Ok(())
185}
186
187fn main() -> Result<(), Box<dyn std::error::Error>> {
188    use cranelift::prelude::*;
189    use cranelift_module::{Linkage, Module};
190    use cranelift_native::builder as isa_builder;
191    use cranelift_object::{ObjectBuilder, ObjectModule};
192
193    // Set up the ISA for the current machine
194    let mut flag_builder = settings::builder();
195    flag_builder.enable("is_pic").unwrap();
196
197    let isa_builder = isa_builder()?;
198    let isa = isa_builder.finish(settings::Flags::new(flag_builder)).unwrap();
199
200    // Set up the object module
201    let builder = ObjectBuilder::new(isa, "my_program.o", cranelift_module::default_libcall_names())?;
202    let mut module = ObjectModule::new(builder);
203
204    // Define the data section with "Hello, world!\n"
205    let mut data_ctx = cranelift_module::DataDescription::new();
206    let data = "Hello, world!\n\0".as_bytes();
207    data_ctx.define(data.to_vec().into_boxed_slice());
208    let data_id = module.declare_data("hello_world", Linkage::Local, false, false)?;
209    module.define_data(data_id, &data_ctx)?;
210
211    // Create a function signature
212    let mut func_sig = module.make_signature();
213    func_sig.returns.push(AbiParam::new(cranelift::codegen::ir::types::I32));
214    let func_id = module.declare_function("main", Linkage::Export, &func_sig)?;
215
216    // Define the function
217    let mut context = module.make_context();
218    context.func.signature = func_sig;
219
220    {
221        let mut builder_ctx = FunctionBuilderContext::new();
222        let mut builder = FunctionBuilder::new(&mut context.func, &mut builder_ctx);
223        //        todo!("Alignment issues?");
224
225        let entry_block = builder.create_block();
226        builder.switch_to_block(entry_block);
227        builder.seal_block(entry_block);
228        let mut puts_sig = module.make_signature();
229        puts_sig.params.push(AbiParam::new(types::I64)); // pointer to char as parameter
230        puts_sig.returns.push(AbiParam::new(types::I32)); // return type of `puts`
231        let puts_func = module.declare_function("puts", Linkage::Import, &puts_sig)?;
232        // get func ref for func id
233        let puts_func_ref = module.declare_func_in_func(puts_func, builder.func);
234
235        // Reference the data section
236        let data_ref = module.declare_data_in_func(data_id, builder.func);
237        let base_addr = builder.ins().global_value(types::I64, data_ref);
238
239        // Call `puts` with the string address
240        let call = builder.ins().call(puts_func_ref, &[base_addr]);
241
242        // Return from the function
243        let ret_val = builder.inst_results(call)[0];
244        builder.ins().return_(&[ret_val]);
245
246        builder.finalize();
247    }
248
249    // Define and compile the function
250    module.define_function(func_id, &mut context)?;
251    module.clear_context(&mut context);
252
253    // Write the object file
254    let product = module.finish();
255    let obj = product.object;
256
257    let mut file = std::fs::File::create("output.o")?;
258    use std::io::Write;
259
260    file.write_all(&obj.write().unwrap())?;
261
262    // make the file executable
263    use std::process::Command;
264
265    // Link the object file using clang
266    Command::new("clang")
267        .arg("output.o")
268        .arg("-o")
269        .arg("output")
270        .arg("-Wl")
271        .arg("-ld_classic")
272        .arg("-v")
273        .status()?;
274
275    use std::{fs, os::unix::fs::PermissionsExt};
276
277    // Set the output file to be executable
278    let mut perms = fs::metadata("output")?.permissions();
279    perms.set_mode(0o755);
280    fs::set_permissions("output", perms)?;
281    // Make the output file executable
282
283    Ok(())
284}