hugr_llvm/
custom.rs

1//! Provides an interface for extending `hugr-llvm` to emit [CustomType]s,
2//! [CustomConst]s, and [ExtensionOp]s.
3//!
4//! [CustomType]: hugr_core::types::CustomType
5//! [CustomConst]: hugr_core::ops::constant::CustomConst
6//! [ExtensionOp]: hugr_core::ops::ExtensionOp
7use std::rc::Rc;
8
9use self::extension_op::{ExtensionOpFn, ExtensionOpMap};
10use hugr_core::{
11    extension::{simple_op::MakeOpDef, ExtensionId},
12    ops::{constant::CustomConst, ExtensionOp, OpName},
13    HugrView, Node,
14};
15
16use strum::IntoEnumIterator;
17use types::CustomTypeKey;
18
19use self::load_constant::{LoadConstantFn, LoadConstantsMap};
20use self::types::LLVMCustomTypeFn;
21use anyhow::Result;
22
23use crate::{
24    emit::{func::EmitFuncContext, EmitOpArgs},
25    types::TypeConverter,
26};
27
28pub mod extension_op;
29pub mod load_constant;
30pub mod types;
31
32/// A helper to register codegen extensions.
33///
34/// Types that implement this trait can be registered with a [CodegenExtsBuilder]
35/// via [CodegenExtsBuilder::add_extension].
36///
37/// See [crate::extension::PreludeCodegenExtension] for an example.
38pub trait CodegenExtension {
39    /// Implementers should add each of their handlers to `builder` and return the
40    /// resulting [CodegenExtsBuilder].
41    fn add_extension<'a, H: HugrView<Node = Node> + 'a>(
42        self,
43        builder: CodegenExtsBuilder<'a, H>,
44    ) -> CodegenExtsBuilder<'a, H>
45    where
46        Self: 'a;
47}
48
49/// A container for a collection of codegen callbacks as they are being
50/// assembled.
51///
52/// The callbacks are registered against several keys:
53///  - [CustomType]s, with [CodegenExtsBuilder::custom_type]
54///  - [CustomConst]s, with [CodegenExtsBuilder::custom_const]
55///  - [ExtensionOp]s, with [CodegenExtsBuilder::extension_op]
56///
57/// Each callback may hold references older than `'a`.
58///
59/// Registering any callback silently replaces any other callback registered for
60/// that same key.
61///
62/// [CustomType]: hugr_core::types::CustomType
63#[derive(Default)]
64pub struct CodegenExtsBuilder<'a, H> {
65    load_constant_handlers: LoadConstantsMap<'a, H>,
66    extension_op_handlers: ExtensionOpMap<'a, H>,
67    type_converter: TypeConverter<'a>,
68}
69
70impl<'a, H: HugrView<Node = Node> + 'a> CodegenExtsBuilder<'a, H> {
71    /// Forwards to [CodegenExtension::add_extension].
72    ///
73    /// ```
74    /// use hugr_llvm::{extension::{PreludeCodegenExtension, DefaultPreludeCodegen}, CodegenExtsBuilder};
75    /// let ext = PreludeCodegenExtension::from(DefaultPreludeCodegen);
76    /// CodegenExtsBuilder::<hugr_core::Hugr>::default().add_extension(ext);
77    /// ```
78    pub fn add_extension(self, ext: impl CodegenExtension + 'a) -> Self {
79        ext.add_extension(self)
80    }
81
82    /// Register a callback to map a [CustomType] to a [BasicTypeEnum].
83    ///
84    /// [CustomType]: hugr_core::types::CustomType
85    /// [BasicTypeEnum]: inkwell::types::BasicTypeEnum
86    pub fn custom_type(
87        mut self,
88        custom_type: CustomTypeKey,
89        handler: impl LLVMCustomTypeFn<'a>,
90    ) -> Self {
91        self.type_converter.custom_type(custom_type, handler);
92        self
93    }
94
95    /// Register a callback to emit a [ExtensionOp], keyed by fully
96    /// qualified [OpName].
97    pub fn extension_op(
98        mut self,
99        extension: ExtensionId,
100        op: OpName,
101        handler: impl ExtensionOpFn<'a, H>,
102    ) -> Self {
103        self.extension_op_handlers
104            .extension_op(extension, op, handler);
105        self
106    }
107
108    /// Register callbacks to emit [ExtensionOp]s that match the
109    /// definitions generated by `Op`s impl of [strum::IntoEnumIterator]>
110    pub fn simple_extension_op<Op: MakeOpDef + IntoEnumIterator>(
111        mut self,
112        handler: impl 'a
113            + for<'c> Fn(
114                &mut EmitFuncContext<'c, 'a, H>,
115                EmitOpArgs<'c, '_, ExtensionOp, H>,
116                Op,
117            ) -> Result<()>,
118    ) -> Self {
119        self.extension_op_handlers
120            .simple_extension_op::<Op>(handler);
121        self
122    }
123
124    /// Register a callback to materialise a constant implemented by `CC`.
125    pub fn custom_const<CC: CustomConst>(
126        mut self,
127        handler: impl LoadConstantFn<'a, H, CC>,
128    ) -> Self {
129        self.load_constant_handlers.custom_const(handler);
130        self
131    }
132
133    /// Consume `self` to return collections of callbacks for each of the
134    /// supported keys.`
135    pub fn finish(self) -> CodegenExtsMap<'a, H> {
136        CodegenExtsMap {
137            load_constant_handlers: Rc::new(self.load_constant_handlers),
138            extension_op_handlers: Rc::new(self.extension_op_handlers),
139            type_converter: Rc::new(self.type_converter),
140        }
141    }
142}
143
144/// The result of [CodegenExtsBuilder::finish]. Users are expected to
145/// deconstruct this type, and consume the fields independently.
146/// We expect to add further collections at a later date, and so this type is
147/// marked `non_exhaustive`
148#[derive(Default)]
149#[non_exhaustive]
150pub struct CodegenExtsMap<'a, H> {
151    pub load_constant_handlers: Rc<LoadConstantsMap<'a, H>>,
152    pub extension_op_handlers: Rc<ExtensionOpMap<'a, H>>,
153    pub type_converter: Rc<TypeConverter<'a>>,
154}
155
156#[cfg(test)]
157mod test {
158    use hugr_core::{
159        extension::prelude::{string_type, ConstString, PRELUDE_ID, PRINT_OP_ID, STRING_TYPE_NAME},
160        Hugr,
161    };
162    use inkwell::{
163        context::Context,
164        types::BasicType,
165        values::{BasicMetadataValueEnum, BasicValue},
166    };
167    use itertools::Itertools as _;
168
169    use crate::{emit::libc::emit_libc_printf, CodegenExtsBuilder};
170
171    #[test]
172    fn types_with_lifetimes() {
173        let n = "name_with_lifetime".to_string();
174
175        let cem = CodegenExtsBuilder::<Hugr>::default()
176            .custom_type((PRELUDE_ID, STRING_TYPE_NAME), |session, _| {
177                let ctx = session.iw_context();
178                Ok(ctx
179                    .get_struct_type(n.as_ref())
180                    .unwrap_or_else(|| ctx.opaque_struct_type(n.as_ref()))
181                    .as_basic_type_enum())
182            })
183            .finish();
184
185        let ctx = Context::create();
186
187        let ty = cem
188            .type_converter
189            .session(&ctx)
190            .llvm_type(&string_type())
191            .unwrap()
192            .into_struct_type();
193        let ty_n = ty.get_name().unwrap().to_str().unwrap();
194        assert_eq!(ty_n, n);
195    }
196
197    #[test]
198    fn custom_const_lifetime_of_context() {
199        let ctx = Context::create();
200
201        let _ = CodegenExtsBuilder::<Hugr>::default()
202            .custom_const::<ConstString>(|_, konst| {
203                Ok(ctx
204                    .const_string(konst.value().as_bytes(), true)
205                    .as_basic_value_enum())
206            })
207            .finish();
208    }
209
210    #[test]
211    fn extension_op_lifetime() {
212        let ctx = Context::create();
213
214        let _ = CodegenExtsBuilder::<Hugr>::default()
215            .extension_op(PRELUDE_ID, PRINT_OP_ID, |context, args| {
216                let mut print_args: Vec<BasicMetadataValueEnum> =
217                    vec![ctx.const_string("%s".as_bytes(), true).into()];
218                print_args.extend(args.inputs.into_iter().map_into::<BasicMetadataValueEnum>());
219                emit_libc_printf(context, &print_args)?;
220                args.outputs.finish(context.builder(), [])
221            })
222            .finish();
223    }
224}
hugr_llvm/custom.rs

hugr_llvm/
custom.rs
