Skip to main content

pliron_llvm/llvm_sys/
execution_engine.rs

1//! Safe(r) wrappers around llvm_sys::execution_engine
2
3use std::mem::{MaybeUninit, forget};
4
5use llvm_sys::{
6    core::LLVMDisposeMessage,
7    execution_engine::{
8        LLVMAddGlobalMapping, LLVMCreateExecutionEngineForModule, LLVMCreateGenericValueOfFloat,
9        LLVMCreateGenericValueOfInt, LLVMCreateGenericValueOfPointer,
10        LLVMCreateInterpreterForModule, LLVMCreateJITCompilerForModule,
11        LLVMCreateMCJITCompilerForModule, LLVMDisposeExecutionEngine, LLVMDisposeGenericValue,
12        LLVMExecutionEngineRef, LLVMFindFunction, LLVMGenericValueRef, LLVMGenericValueToFloat,
13        LLVMGenericValueToInt, LLVMGenericValueToPointer, LLVMInitializeMCJITCompilerOptions,
14        LLVMLinkInInterpreter, LLVMLinkInMCJIT, LLVMMCJITCompilerOptions, LLVMRunFunction,
15        LLVMRunFunctionAsMain, LLVMRunStaticConstructors, LLVMRunStaticDestructors,
16    },
17};
18
19use crate::llvm_sys::{
20    core::{LLVMModule, LLVMType, LLVMValue},
21    cstr_to_string, to_c_str,
22};
23
24/// Wrapper around [LLVMGenericValueRef].
25pub struct GenericValue(LLVMGenericValueRef);
26
27impl GenericValue {
28    /// Creates a [GenericValue] representing an integer.
29    pub fn from_u64(ty: LLVMType, n: u64, is_signed: bool) -> GenericValue {
30        let gv = unsafe { LLVMCreateGenericValueOfInt(ty.into(), n, is_signed.into()) };
31        GenericValue(gv)
32    }
33
34    /// Creates a [GenericValue] representing a pointer.
35    pub fn from_pointer<T>(ptr: *mut T) -> GenericValue {
36        let gv = unsafe { LLVMCreateGenericValueOfPointer(ptr as _) };
37        GenericValue(gv)
38    }
39
40    /// Creates a [GenericValue] representing a float.
41    pub fn from_f64(ty: LLVMType, n: f64) -> GenericValue {
42        let gv = unsafe { LLVMCreateGenericValueOfFloat(ty.into(), n) };
43        GenericValue(gv)
44    }
45
46    /// Returns [Self] as u64.
47    pub fn to_u64(&self) -> u64 {
48        unsafe { LLVMGenericValueToInt(self.0, 0) as u64 }
49    }
50
51    /// Returns [Self] as pointer.
52    pub fn to_pointer<T>(&self) -> *mut T {
53        unsafe { LLVMGenericValueToPointer(self.0) as *mut T }
54    }
55
56    /// Returns [Self] as f64.
57    pub fn to_f64(&self, ty: LLVMType) -> f64 {
58        unsafe { LLVMGenericValueToFloat(ty.into(), self.0) }
59    }
60}
61
62impl Drop for GenericValue {
63    fn drop(&mut self) {
64        unsafe {
65            LLVMDisposeGenericValue(self.0);
66        }
67    }
68}
69
70/// Rust wrapper around [LLVMMCJITCompilerOptions]
71#[derive(Clone, Debug, Copy)]
72pub struct MCJITCompilerOptions(pub LLVMMCJITCompilerOptions);
73
74impl Default for MCJITCompilerOptions {
75    fn default() -> Self {
76        let mut options = MaybeUninit::uninit();
77        unsafe {
78            LLVMInitializeMCJITCompilerOptions(
79                options.as_mut_ptr(),
80                std::mem::size_of::<LLVMMCJITCompilerOptions>(),
81            );
82            MCJITCompilerOptions(options.assume_init())
83        }
84    }
85}
86
87/// Rust wrapper around [LLVMExecutionEngineRef]
88pub struct ExecutionEngine(LLVMExecutionEngineRef);
89
90/// Code generation optimization level.
91/// This is a copy of LLVM's `CodeGenOptLevel` enum.
92#[derive(Clone, Copy, Debug, PartialEq, Eq)]
93pub enum CodeGenOptLevel {
94    /// -O0
95    None = 0,
96    /// -O1
97    Less = 1,
98    /// -O2, -Os
99    Default = 2,
100    /// -O3
101    Aggressive = 3,
102}
103
104/// Kind of execution engine to create.
105#[derive(Clone, Copy, Debug)]
106pub enum EngineKind {
107    JIT(CodeGenOptLevel),
108    Interpreter,
109    // Either of JIT or Interpreter.
110    Either,
111    MCJIT(MCJITCompilerOptions),
112}
113
114impl ExecutionEngine {
115    /// Creates a new [ExecutionEngine] for the given module.
116    ///
117    /// ### Example usage:
118    /// ```
119    /// use pliron_llvm::llvm_sys::{
120    ///     core::{LLVMMemoryBuffer, LLVMContext, LLVMModule},
121    ///     execution_engine::{ExecutionEngine, EngineKind, MCJITCompilerOptions},
122    ///     target::initialize_native,
123    /// };
124    /// fn main() -> Result<(), String> {
125    ///     let llvm_ctx = LLVMContext::default();
126    ///     ExecutionEngine::link_in_mcjit();
127    ///     initialize_native()?;
128    ///
129    ///     let llvm_ir = r#"
130    ///         define i32 @main() {
131    ///             ret i32 0
132    ///         }"#;
133    ///     let ir_mb = LLVMMemoryBuffer::from_str(llvm_ir, "test_buffer");
134    ///     let module = LLVMModule::from_ir_in_memory_buffer(&llvm_ctx, ir_mb)?;
135    ///
136    ///     let mcjit_options = MCJITCompilerOptions::default();
137    ///     let ee = ExecutionEngine::new_for_module(module, EngineKind::MCJIT(mcjit_options))?;
138    ///     let main = ee
139    ///         .find_function("main")
140    ///         .ok_or("Function 'main' not found")?;
141    ///     let ret_gv = unsafe { ee.run_function_as_main(main, &[]) };
142    ///     assert_eq!(ret_gv, 0);
143    ///     Ok(())
144    /// }
145    /// ```
146    pub fn new_for_module(module: LLVMModule, kind: EngineKind) -> Result<Self, String> {
147        let mut ee = MaybeUninit::uninit();
148        let mut error_string = MaybeUninit::uninit();
149        let result = unsafe {
150            match kind {
151                EngineKind::Either => LLVMCreateExecutionEngineForModule(
152                    ee.as_mut_ptr(),
153                    module.inner_ref(),
154                    error_string.as_mut_ptr(),
155                ),
156                EngineKind::Interpreter => LLVMCreateInterpreterForModule(
157                    ee.as_mut_ptr(),
158                    module.inner_ref(),
159                    error_string.as_mut_ptr(),
160                ),
161                EngineKind::JIT(opt_level) => LLVMCreateJITCompilerForModule(
162                    ee.as_mut_ptr(),
163                    module.inner_ref(),
164                    opt_level as _,
165                    error_string.as_mut_ptr(),
166                ),
167                EngineKind::MCJIT(options) => LLVMCreateMCJITCompilerForModule(
168                    ee.as_mut_ptr(),
169                    module.inner_ref(),
170                    // This is ok, `LLVMCreateMCJITCompilerForModule` creates a copy.
171                    &options.0 as *const LLVMMCJITCompilerOptions as *mut LLVMMCJITCompilerOptions,
172                    std::mem::size_of::<LLVMMCJITCompilerOptions>(),
173                    error_string.as_mut_ptr(),
174                ),
175            }
176        };
177        // We can forget the module, as the execution engine now owns it.
178        forget(module);
179        if result != 0 {
180            unsafe {
181                let err_str = error_string.assume_init();
182                let err_string = cstr_to_string(err_str).unwrap();
183                LLVMDisposeMessage(err_str);
184                Err(err_string)
185            }
186        } else {
187            let ee = unsafe { ee.assume_init() };
188            Ok(ExecutionEngine(ee))
189        }
190    }
191
192    /// Runs static constructors.
193    ///
194    /// ### Safety
195    /// This function executes arbitrary code.
196    pub unsafe fn run_static_constructors(&self) {
197        unsafe {
198            LLVMRunStaticConstructors(self.0);
199        }
200    }
201
202    /// Runs static destructors.
203    ///
204    /// ### Safety
205    /// This function executes arbitrary code.
206    pub unsafe fn run_static_destructors(&self) {
207        unsafe {
208            LLVMRunStaticDestructors(self.0);
209        }
210    }
211
212    /// Find a function by name.
213    pub fn find_function(&self, name: &str) -> Option<LLVMValue> {
214        let mut func = MaybeUninit::uninit();
215        let result =
216            unsafe { LLVMFindFunction(self.0, to_c_str(name).as_ptr(), func.as_mut_ptr()) };
217        if result == 0 {
218            let func = unsafe { func.assume_init() };
219            Some(func.into())
220        } else {
221            None
222        }
223    }
224
225    /// Map an external global into the execution engine.
226    ///
227    /// ### Example usage:
228    /// ```
229    /// use pliron_llvm::llvm_sys::{
230    ///     core::{LLVMMemoryBuffer, LLVMContext, LLVMModule, llvm_get_named_function},
231    ///     execution_engine::{ExecutionEngine, EngineKind, CodeGenOptLevel},
232    ///     target::initialize_native,
233    /// };
234    /// fn main() -> Result<(), String> {
235    ///     let llvm_ctx = LLVMContext::default();
236    ///     ExecutionEngine::link_in_mcjit();
237    ///     initialize_native()?;
238    ///
239    ///     let llvm_ir = r#"
240    ///         declare i32 @external_function()
241    ///
242    ///         define i32 @call_external() {
243    ///             %result = call i32 @external_function()
244    ///             ret i32 %result
245    ///         }"#;
246    ///     let ir_mb = LLVMMemoryBuffer::from_str(llvm_ir, "test_buffer");
247    ///     let module = LLVMModule::from_ir_in_memory_buffer(&llvm_ctx, ir_mb)?;
248    ///     let ext_fn_value = llvm_get_named_function(&module, "external_function")
249    ///         .ok_or("Function 'external_function' not found")?;
250    ///
251    ///     let ee =
252    ///         ExecutionEngine::new_for_module(module, EngineKind::JIT(CodeGenOptLevel::Default))?;
253    ///     let call_external_fn = ee
254    ///         .find_function("call_external")
255    ///         .ok_or("Function 'call_external' not found")?;
256    ///
257    ///     extern "C" fn external_function() -> i32 {
258    ///         1234
259    ///     }
260    ///
261    ///     ee.add_global_mapping(ext_fn_value, external_function as *mut fn() -> i32);
262    ///
263    ///     let ret_gv = unsafe { ee.run_function(call_external_fn, &[]) };
264    ///     let ret_val = ret_gv.to_u64();
265    ///     assert_eq!(ret_val, 1234);
266    ///     Ok(())
267    /// }
268    /// ```
269    pub fn add_global_mapping<T>(&self, function: LLVMValue, addr: *mut T) {
270        unsafe {
271            LLVMAddGlobalMapping(self.0, function.into(), addr as _);
272        }
273    }
274
275    /// Execute the function with the given arguments.
276    ///
277    /// ### Safety
278    /// This function executes arbitrary code.
279    ///
280    /// ### Example usage:
281    /// ```
282    /// use pliron_llvm::llvm_sys::{
283    ///     core::{LLVMMemoryBuffer, LLVMContext, LLVMModule, llvm_int_type_in_context},
284    ///     execution_engine::{ExecutionEngine, EngineKind, GenericValue},
285    ///     target::initialize_native,
286    /// };
287    /// fn main() -> Result<(), String> {
288    ///     let llvm_ctx = LLVMContext::default();
289    ///     ExecutionEngine::link_in_interpreter();
290    ///     initialize_native()?;
291    ///
292    ///     let llvm_ir = r#"
293    ///         define i32 @add(i32 %a, i32 %b) {
294    ///             %sum = add i32 %a, %b
295    ///             ret i32 %sum
296    ///         }"#;
297    ///     let ir_mb = LLVMMemoryBuffer::from_str(llvm_ir, "test_buffer");
298    ///     let module = LLVMModule::from_ir_in_memory_buffer(&llvm_ctx, ir_mb)?;
299    ///     let ee = ExecutionEngine::new_for_module(module, EngineKind::Interpreter)?;
300    ///     let add_fn = ee.find_function("add").ok_or("Function 'add' not found")?;
301    ///
302    ///     let i32_type = llvm_int_type_in_context(&llvm_ctx, 32);
303    ///     let arg1 = GenericValue::from_u64(i32_type, 10, false);
304    ///     let arg2 = GenericValue::from_u64(i32_type, 32, false);
305    ///     let ret_gv = unsafe { ee.run_function(add_fn, &[arg1, arg2]) };
306    ///     let ret_val = ret_gv.to_u64();
307    ///     assert_eq!(ret_val, 42);
308    ///     Ok(())
309    /// }
310    /// ```
311    pub unsafe fn run_function(&self, function: LLVMValue, args: &[GenericValue]) -> GenericValue {
312        let mut args = args
313            .iter()
314            .map(|gv| gv.0)
315            .collect::<Vec<LLVMGenericValueRef>>();
316        let gv = unsafe {
317            LLVMRunFunction(
318                self.0,
319                function.into(),
320                args.len().try_into().unwrap(),
321                args.as_mut_ptr(),
322            )
323        };
324        GenericValue(gv)
325    }
326
327    /// Execute a function as main.
328    ///
329    /// ### Safety
330    /// This function executes arbitrary code.
331    //
332    /// ### Example usage:
333    /// ```
334    /// use pliron_llvm::llvm_sys::{
335    ///     core::{LLVMMemoryBuffer, LLVMContext, LLVMModule},
336    ///     execution_engine::{ExecutionEngine, EngineKind},
337    ///     target::initialize_native,
338    /// };
339    /// fn main() -> Result<(), String> {
340    ///     let llvm_ctx = LLVMContext::default();
341    ///     ExecutionEngine::link_in_interpreter();
342    ///     initialize_native()?;
343    ///
344    ///     let llvm_ir = r#"
345    ///         define i32 @main() {
346    ///             ret i32 42
347    ///         }"#;
348    ///     let ir_mb = LLVMMemoryBuffer::from_str(llvm_ir, "test_buffer");
349    ///     let module = LLVMModule::from_ir_in_memory_buffer(&llvm_ctx, ir_mb)?;
350    ///     let ee = ExecutionEngine::new_for_module(module, EngineKind::Interpreter)?;
351    ///     let main = ee
352    ///         .find_function("main")
353    ///         .ok_or("Function 'main' not found")?;
354    ///     let ret_gv = unsafe { ee.run_function_as_main(main, &[]) };
355    ///     assert_eq!(ret_gv, 42);
356    ///     Ok(())
357    /// }
358    /// ```
359    pub unsafe fn run_function_as_main(&self, function: LLVMValue, args: &[&str]) -> i32 {
360        let c_args: Vec<_> = args.iter().map(|arg| to_c_str(arg)).collect();
361        let mut c_args_ptrs: Vec<_> = c_args.iter().map(|cstr| cstr.as_ptr()).collect();
362        unsafe {
363            LLVMRunFunctionAsMain(
364                self.0,
365                function.into(),
366                c_args_ptrs.len().try_into().unwrap(),
367                c_args_ptrs.as_mut_ptr(),
368                // TODO: Support env variables.
369                [std::ptr::null()].as_ptr(),
370            )
371        }
372    }
373
374    /// Link in Interpreter.
375    pub fn link_in_interpreter() {
376        unsafe {
377            LLVMLinkInInterpreter();
378        }
379    }
380
381    /// Link in MCJIT.
382    pub fn link_in_mcjit() {
383        unsafe {
384            LLVMLinkInMCJIT();
385        }
386    }
387}
388
389impl Drop for ExecutionEngine {
390    fn drop(&mut self) {
391        unsafe {
392            LLVMDisposeExecutionEngine(self.0);
393        }
394    }
395}