lower_ir_utils/
builder.rs

1//! High-level helper to define a Cranelift function: declares the function,
2//! creates the `FunctionBuilderContext` + entry block, runs the user-supplied
3//! body closure to populate the IR, emits the return, and finalizes.
4//!
5//! The body closure receives `&mut FunctionBuilder`, `&mut Module`, and the
6//! entry-block parameters. It returns the values to be returned, via the
7//! [`IntoReturns`] trait — so `()`, a single `Value`, `[Value; N]`, or
8//! `Vec<Value>` all work without ceremony.
9
10use cranelift_codegen::ir::{InstBuilder, Signature, UserFuncName, Value};
11use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
12use cranelift_module::{FuncId, Linkage, Module, ModuleResult};
13use smallvec::SmallVec;
14
15/// Conversion from common Rust shapes to a return-Value list.
16///
17/// Used by [`define_function`] to let the body closure return whatever is most
18/// natural: `()` for void, a bare `Value` for one return, `[Value; N]` /
19/// `Vec<Value>` / `SmallVec` for multi-return.
20pub trait IntoReturns {
21    fn into_returns(self) -> SmallVec<[Value; 4]>;
22}
23
24impl IntoReturns for () {
25    fn into_returns(self) -> SmallVec<[Value; 4]> {
26        SmallVec::new()
27    }
28}
29
30impl IntoReturns for Value {
31    fn into_returns(self) -> SmallVec<[Value; 4]> {
32        let mut s = SmallVec::new();
33        s.push(self);
34        s
35    }
36}
37
38impl<const N: usize> IntoReturns for [Value; N] {
39    fn into_returns(self) -> SmallVec<[Value; 4]> {
40        SmallVec::from_iter(self)
41    }
42}
43
44impl IntoReturns for Vec<Value> {
45    fn into_returns(self) -> SmallVec<[Value; 4]> {
46        SmallVec::from_vec(self)
47    }
48}
49
50impl IntoReturns for SmallVec<[Value; 4]> {
51    fn into_returns(self) -> SmallVec<[Value; 4]> {
52        self
53    }
54}
55
56/// Declare and define a Cranelift function in one call.
57///
58/// `body` is invoked with the live `FunctionBuilder`, the module (re-borrowed
59/// so the body can declare imports / call other JITed functions), and the
60/// entry block's parameters. Whatever it returns is wrapped via [`IntoReturns`]
61/// and emitted as the function's `return_` instruction.
62///
63/// # Example
64///
65/// ```ignore
66/// let id = define_function(
67///     &mut module,
68///     "wrap",
69///     Linkage::Export,
70///     jit_signature!(&module; fn(i64) -> i64),
71///     |bcx, module, params| {
72///         double_i64_jit::call(bcx, module, ext_id, params[0])
73///     },
74/// )?;
75/// ```
76#[allow(clippy::result_large_err)] // ModuleError is cranelift's type; not our concern.
77pub fn define_function<M, F, R>(
78    module: &mut M,
79    name: &str,
80    linkage: Linkage,
81    signature: Signature,
82    body: F,
83) -> ModuleResult<FuncId>
84where
85    M: Module,
86    F: FnOnce(&mut FunctionBuilder<'_>, &mut M, &[Value]) -> R,
87    R: IntoReturns,
88{
89    let func_id = module.declare_function(name, linkage, &signature)?;
90
91    let mut ctx = module.make_context();
92    ctx.func.signature = signature;
93    ctx.func.name = UserFuncName::user(0, func_id.as_u32());
94
95    let mut bcx_ctx = FunctionBuilderContext::new();
96    {
97        let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut bcx_ctx);
98        let entry = bcx.create_block();
99        bcx.append_block_params_for_function_params(entry);
100        bcx.switch_to_block(entry);
101        bcx.seal_block(entry);
102        let params: SmallVec<[Value; 8]> = bcx.block_params(entry).iter().copied().collect();
103        let returns = body(&mut bcx, module, &params).into_returns();
104        bcx.ins().return_(&returns);
105        bcx.finalize();
106    }
107
108    module.define_function(func_id, &mut ctx)?;
109    module.clear_context(&mut ctx);
110    Ok(func_id)
111}
lower_ir_utils/builder.rs

lower_ir_utils/
builder.rs
