Skip to main content

pliron_llvm/llvm_sys/
execution_engine.rs

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