hugr_llvm/extension/
prelude.rs

1use anyhow::{Ok, Result, anyhow, bail, ensure};
2use hugr_core::Node;
3use hugr_core::extension::prelude::generic::LoadNat;
4use hugr_core::extension::prelude::{
5    self, ConstError, ConstExternalSymbol, ConstString, ConstUsize, MakeTuple, TupleOpDef,
6    UnpackTuple, error_type, generic,
7};
8use hugr_core::extension::prelude::{ERROR_TYPE_NAME, STRING_TYPE_NAME};
9use hugr_core::ops::ExtensionOp;
10use hugr_core::types::TypeArg;
11use hugr_core::{
12    HugrView, extension::simple_op::MakeExtensionOp as _, ops::constant::CustomConst,
13    types::SumType,
14};
15use inkwell::{
16    AddressSpace,
17    types::{BasicType, IntType, PointerType},
18    values::{BasicValue as _, BasicValueEnum, StructValue},
19};
20use itertools::Itertools;
21
22use crate::emit::EmitOpArgs;
23use crate::{
24    custom::{CodegenExtension, CodegenExtsBuilder},
25    emit::{
26        func::EmitFuncContext,
27        libc::{emit_libc_abort, emit_libc_printf},
28    },
29    sum::LLVMSumValue,
30    types::TypingSession,
31};
32
33/// A helper trait for customising the lowering [`hugr_core::extension::prelude`]
34/// types, [`CustomConst`]s, and ops.
35///
36/// All methods have sensible defaults provided, and [`DefaultPreludeCodegen`] is
37/// a trivial implementation of this trait which delegates everything to those
38/// default implementations.
39pub trait PreludeCodegen: Clone {
40    /// Return the llvm type of [`hugr_core::extension::prelude::usize_t`]. That type
41    /// must be an [`IntType`].
42    fn usize_type<'c>(&self, session: &TypingSession<'c, '_>) -> IntType<'c> {
43        session.iw_context().i64_type()
44    }
45
46    /// Return the llvm type of [`hugr_core::extension::prelude::qb_t`].
47    fn qubit_type<'c>(&self, session: &TypingSession<'c, '_>) -> impl BasicType<'c> {
48        session.iw_context().i16_type()
49    }
50
51    /// Return the llvm type of [`hugr_core::extension::prelude::error_type()`].
52    ///
53    /// The returned type must always match the type of the returned value of
54    /// [`Self::emit_const_error`], and the `err` argument of [`Self::emit_panic`].
55    ///
56    /// The default implementation is a struct type with an i32 field and an i8*
57    /// field for the code and message.
58    fn error_type<'c>(&self, session: &TypingSession<'c, '_>) -> Result<impl BasicType<'c>> {
59        let ctx = session.iw_context();
60        Ok(session.iw_context().struct_type(
61            &[
62                ctx.i32_type().into(),
63                ctx.i8_type().ptr_type(AddressSpace::default()).into(),
64            ],
65            false,
66        ))
67    }
68
69    /// Return the llvm type of [`hugr_core::extension::prelude::string_type()`].
70    ///
71    /// The returned type must always match the type of the returned value of
72    /// [`Self::emit_const_string`], and the `text` argument of [`Self::emit_print`].
73    ///
74    /// The default implementation is i8*.
75    fn string_type<'c>(&self, session: &TypingSession<'c, '_>) -> Result<impl BasicType<'c>> {
76        Ok(session
77            .iw_context()
78            .i8_type()
79            .ptr_type(AddressSpace::default()))
80    }
81
82    /// Emit a [`hugr_core::extension::prelude::PRINT_OP_ID`] node.
83    fn emit_print<H: HugrView<Node = Node>>(
84        &self,
85        ctx: &mut EmitFuncContext<H>,
86        text: BasicValueEnum,
87    ) -> Result<()> {
88        let format_str = ctx
89            .builder()
90            .build_global_string_ptr("%s\n", "prelude.print_template")?
91            .as_basic_value_enum();
92        emit_libc_printf(ctx, &[format_str.into(), text.into()])
93    }
94
95    /// Emit instructions to materialise an LLVM value representing `err`.
96    ///
97    /// The type of the returned value must match [`Self::error_type`].
98    ///
99    /// The default implementation materialises an LLVM struct with the
100    /// [`ConstError::signal`] and [`ConstError::message`] of `err`.
101    fn emit_const_error<'c, H: HugrView<Node = Node>>(
102        &self,
103        ctx: &mut EmitFuncContext<'c, '_, H>,
104        err: &ConstError,
105    ) -> Result<BasicValueEnum<'c>> {
106        let builder = ctx.builder();
107        let err_ty = ctx.llvm_type(&error_type())?.into_struct_type();
108        let signal = err_ty
109            .get_field_type_at_index(0)
110            .unwrap()
111            .into_int_type()
112            .const_int(u64::from(err.signal), false);
113        let message = builder
114            .build_global_string_ptr(&err.message, "")?
115            .as_basic_value_enum();
116        let err = err_ty.const_named_struct(&[signal.into(), message]);
117        Ok(err.into())
118    }
119
120    /// Emit instructions to construct an error value from a signal and message.
121    ///
122    /// The type of the returned value must match [`Self::error_type`].
123    ///
124    /// The default implementation constructs a struct with the given signal and message.
125    fn emit_make_error<'c, H: HugrView<Node = Node>>(
126        &self,
127        ctx: &mut EmitFuncContext<'c, '_, H>,
128        signal: BasicValueEnum<'c>,
129        message: BasicValueEnum<'c>,
130    ) -> Result<BasicValueEnum<'c>> {
131        let builder = ctx.builder();
132
133        // The usize signal is an i64 but error struct stores an i32.
134        let i32_type = ctx.typing_session().iw_context().i32_type();
135        let signal_int = signal.into_int_value();
136        let signal_truncated = builder.build_int_truncate(signal_int, i32_type, "")?;
137
138        // Construct the error struct as runtime value.
139        let err_ty = ctx.llvm_type(&error_type())?.into_struct_type();
140        let undef = err_ty.get_undef();
141        let err_with_sig = builder
142            .build_insert_value(undef, signal_truncated, 0, "")?
143            .into_struct_value();
144        let err_complete = builder
145            .build_insert_value(err_with_sig, message, 1, "")?
146            .into_struct_value();
147
148        Ok(err_complete.into())
149    }
150
151    /// Emit instructions to halt execution with the error `err`.
152    ///
153    /// The type of `err` must match that returned from [`Self::error_type`].
154    ///
155    /// The default implementation emits calls to libc's `printf` and `abort`.
156    ///
157    /// Note that implementations of `emit_panic` must not emit `unreachable`
158    /// terminators, that, if appropriate, is the responsibility of the caller.
159    fn emit_panic<H: HugrView<Node = Node>>(
160        &self,
161        ctx: &mut EmitFuncContext<H>,
162        err: BasicValueEnum,
163    ) -> Result<()> {
164        let format_str = ctx
165            .builder()
166            .build_global_string_ptr(
167                "Program panicked (signal %i): %s\n",
168                "prelude.panic_template",
169            )?
170            .as_basic_value_enum();
171        let Some(err) = StructValue::try_from(err).ok() else {
172            bail!("emit_panic: Expected err value to be a struct type")
173        };
174        ensure!(err.get_type().count_fields() == 2);
175        let signal = ctx.builder().build_extract_value(err, 0, "")?;
176        ensure!(signal.get_type() == ctx.iw_context().i32_type().as_basic_type_enum());
177        let msg = ctx.builder().build_extract_value(err, 1, "")?;
178        ensure!(PointerType::try_from(msg.get_type()).is_ok());
179        emit_libc_printf(ctx, &[format_str.into(), signal.into(), msg.into()])?;
180        emit_libc_abort(ctx)
181    }
182
183    /// Emit instructions to halt execution with the error `err`.
184    ///
185    /// The type of `err` must match that returned from [`Self::error_type`].
186    ///
187    /// The default implementation emits calls to libc's `printf` and `abort`,
188    /// matching the default implementation of [`Self::emit_panic`].
189    ///
190    /// Note that implementations of `emit_panic` must not emit `unreachable`
191    /// terminators, that, if appropriate, is the responsibility of the caller.
192    fn emit_exit<H: HugrView<Node = Node>>(
193        &self,
194        ctx: &mut EmitFuncContext<H>,
195        err: BasicValueEnum,
196    ) -> Result<()> {
197        self.emit_panic(ctx, err)
198    }
199
200    /// Emit instructions to materialise an LLVM value representing `str`.
201    ///
202    /// The type of the returned value must match [`Self::string_type`].
203    ///
204    /// The default implementation creates a global C string.
205    fn emit_const_string<'c, H: HugrView<Node = Node>>(
206        &self,
207        ctx: &mut EmitFuncContext<'c, '_, H>,
208        str: &ConstString,
209    ) -> Result<BasicValueEnum<'c>> {
210        let default_str_type = ctx
211            .iw_context()
212            .i8_type()
213            .ptr_type(AddressSpace::default())
214            .as_basic_type_enum();
215        let str_type = ctx.llvm_type(&str.get_type())?.as_basic_type_enum();
216        ensure!(
217            str_type == default_str_type,
218            "The default implementation of PreludeCodegen::string_type was overridden, but the default implementation of emit_const_string was not. String type is: {str_type}"
219        );
220        let s = ctx.builder().build_global_string_ptr(str.value(), "")?;
221        Ok(s.as_basic_value_enum())
222    }
223
224    fn emit_barrier<'c, H: HugrView<Node = Node>>(
225        &self,
226        ctx: &mut EmitFuncContext<'c, '_, H>,
227        args: EmitOpArgs<'c, '_, ExtensionOp, H>,
228    ) -> Result<()> {
229        // By default, treat barriers as no-ops.
230        args.outputs.finish(ctx.builder(), args.inputs)
231    }
232}
233
234/// A trivial implementation of [`PreludeCodegen`] which passes all methods
235/// through to their default implementations.
236#[derive(Default, Clone)]
237pub struct DefaultPreludeCodegen;
238
239impl PreludeCodegen for DefaultPreludeCodegen {}
240
241#[derive(Clone, Debug, Default)]
242pub struct PreludeCodegenExtension<PCG>(PCG);
243
244impl<PCG: PreludeCodegen> PreludeCodegenExtension<PCG> {
245    pub fn new(pcg: PCG) -> Self {
246        Self(pcg)
247    }
248}
249
250impl<PCG: PreludeCodegen> From<PCG> for PreludeCodegenExtension<PCG> {
251    fn from(pcg: PCG) -> Self {
252        Self::new(pcg)
253    }
254}
255
256impl<PCG: PreludeCodegen> CodegenExtension for PreludeCodegenExtension<PCG> {
257    fn add_extension<'a, H: HugrView<Node = Node> + 'a>(
258        self,
259        builder: CodegenExtsBuilder<'a, H>,
260    ) -> CodegenExtsBuilder<'a, H>
261    where
262        Self: 'a,
263    {
264        add_prelude_extensions(builder, self.0)
265    }
266}
267
268impl<'a, H: HugrView<Node = Node> + 'a> CodegenExtsBuilder<'a, H> {
269    /// Add a [`PreludeCodegenExtension`] to the given [`CodegenExtsBuilder`] using `pcg`
270    /// as the implementation.
271    #[must_use]
272    pub fn add_default_prelude_extensions(self) -> Self {
273        self.add_prelude_extensions(DefaultPreludeCodegen)
274    }
275
276    /// Add a [`PreludeCodegenExtension`] to the given [`CodegenExtsBuilder`] using
277    /// [`DefaultPreludeCodegen`] as the implementation.
278    pub fn add_prelude_extensions(self, pcg: impl PreludeCodegen + 'a) -> Self {
279        self.add_extension(PreludeCodegenExtension::from(pcg))
280    }
281}
282
283/// Add a [`PreludeCodegenExtension`] to the given [`CodegenExtsBuilder`] using `pcg`
284/// as the implementation.
285pub fn add_prelude_extensions<'a, H: HugrView<Node = Node> + 'a>(
286    cem: CodegenExtsBuilder<'a, H>,
287    pcg: impl PreludeCodegen + 'a,
288) -> CodegenExtsBuilder<'a, H> {
289    cem.custom_type((prelude::PRELUDE_ID, "qubit".into()), {
290        let pcg = pcg.clone();
291        move |ts, _| Ok(pcg.qubit_type(&ts).as_basic_type_enum())
292    })
293    .custom_type((prelude::PRELUDE_ID, "usize".into()), {
294        let pcg = pcg.clone();
295        move |ts, _| Ok(pcg.usize_type(&ts).as_basic_type_enum())
296    })
297    .custom_type((prelude::PRELUDE_ID, ERROR_TYPE_NAME.clone()), {
298        let pcg = pcg.clone();
299        move |ts, _| Ok(pcg.error_type(&ts)?.as_basic_type_enum())
300    })
301    .custom_type((prelude::PRELUDE_ID, STRING_TYPE_NAME.clone()), {
302        let pcg = pcg.clone();
303        move |ts, _| Ok(pcg.string_type(&ts)?.as_basic_type_enum())
304    })
305    .custom_const::<ConstUsize>(|context, k| {
306        let ty: IntType = context
307            .llvm_type(&k.get_type())?
308            .try_into()
309            .map_err(|()| anyhow!("Failed to get ConstUsize as IntType"))?;
310        Ok(ty.const_int(k.value(), false).into())
311    })
312    .custom_const::<ConstExternalSymbol>(|context, k| {
313        // TODO we should namespace these symbols
314        // https://github.com/CQCL/hugr-llvm/issues/120
315        let llvm_type = context.llvm_type(&k.get_type())?;
316        let global = context.get_global(&k.symbol, llvm_type, k.constant)?;
317        Ok(context
318            .builder()
319            .build_load(global.as_pointer_value(), &k.symbol)?)
320    })
321    .custom_const::<ConstString>({
322        let pcg = pcg.clone();
323        move |context, k| {
324            let err = pcg.emit_const_string(context, k)?;
325            ensure!(
326                err.get_type()
327                    == pcg
328                        .string_type(&context.typing_session())?
329                        .as_basic_type_enum()
330            );
331            Ok(err)
332        }
333    })
334    .custom_const::<ConstError>({
335        let pcg = pcg.clone();
336        move |context, k| {
337            let err = pcg.emit_const_error(context, k)?;
338            ensure!(
339                err.get_type()
340                    == pcg
341                        .error_type(&context.typing_session())?
342                        .as_basic_type_enum()
343            );
344            Ok(err)
345        }
346    })
347    .extension_op(prelude::PRELUDE_ID, prelude::NOOP_OP_ID, |context, args| {
348        args.outputs.finish(context.builder(), args.inputs)
349    })
350    .simple_extension_op::<TupleOpDef>(|context, args, op| match op {
351        TupleOpDef::UnpackTuple => {
352            let unpack_tuple = UnpackTuple::from_extension_op(args.node().as_ref())?;
353            let llvm_sum_type = context.llvm_sum_type(SumType::new([unpack_tuple.0]))?;
354            let llvm_sum_value = args
355                .inputs
356                .into_iter()
357                .exactly_one()
358                .map_err(|_| anyhow!("UnpackTuple does not have exactly one input"))
359                .and_then(|v| LLVMSumValue::try_new(v, llvm_sum_type))?;
360            let rs = llvm_sum_value.build_untag(context.builder(), 0)?;
361            args.outputs.finish(context.builder(), rs)
362        }
363        TupleOpDef::MakeTuple => {
364            let make_tuple = MakeTuple::from_extension_op(args.node().as_ref())?;
365            let llvm_sum_type = context.llvm_sum_type(SumType::new([make_tuple.0]))?;
366            let r = llvm_sum_type.build_tag(context.builder(), 0, args.inputs)?;
367            args.outputs.finish(context.builder(), [r.into()])
368        }
369        _ => Err(anyhow!("Unsupported TupleOpDef")),
370    })
371    .extension_op(prelude::PRELUDE_ID, prelude::PRINT_OP_ID, {
372        let pcg = pcg.clone();
373        move |context, args| {
374            let text = args.inputs[0];
375            pcg.emit_print(context, text)?;
376            args.outputs.finish(context.builder(), [])
377        }
378    })
379    .extension_op(prelude::PRELUDE_ID, prelude::MAKE_ERROR_OP_ID, {
380        let pcg = pcg.clone();
381        move |context, args| {
382            let signal = args.inputs[0];
383            let message = args.inputs[1];
384            ensure!(
385                message.get_type()
386                    == pcg
387                        .string_type(&context.typing_session())?
388                        .as_basic_type_enum(),
389                signal.get_type() == pcg.usize_type(&context.typing_session()).into()
390            );
391            let err = pcg.emit_make_error(context, signal, message)?;
392            args.outputs.finish(context.builder(), [err])
393        }
394    })
395    .extension_op(prelude::PRELUDE_ID, prelude::PANIC_OP_ID, {
396        let pcg = pcg.clone();
397        move |context, args| {
398            let err = args.inputs[0];
399            ensure!(
400                err.get_type()
401                    == pcg
402                        .error_type(&context.typing_session())?
403                        .as_basic_type_enum()
404            );
405            pcg.emit_panic(context, err)?;
406            let returns = args
407                .outputs
408                .get_types()
409                .map(inkwell::types::BasicTypeEnum::const_zero)
410                .collect_vec();
411            args.outputs.finish(context.builder(), returns)
412        }
413    })
414    .extension_op(prelude::PRELUDE_ID, prelude::EXIT_OP_ID, {
415        // by default treat an exit like a panic
416        let pcg = pcg.clone();
417        move |context, args| {
418            let err = args.inputs[0];
419            ensure!(
420                err.get_type()
421                    == pcg
422                        .error_type(&context.typing_session())?
423                        .as_basic_type_enum()
424            );
425            pcg.emit_exit(context, err)?;
426            let returns = args
427                .outputs
428                .get_types()
429                .map(inkwell::types::BasicTypeEnum::const_zero)
430                .collect_vec();
431            args.outputs.finish(context.builder(), returns)
432        }
433    })
434    .extension_op(prelude::PRELUDE_ID, generic::LOAD_NAT_OP_ID.clone(), {
435        let pcg = pcg.clone();
436        move |context, args| {
437            let load_nat = LoadNat::from_extension_op(args.node().as_ref())?;
438            let v = match load_nat.get_nat() {
439                TypeArg::BoundedNat(n) => pcg
440                    .usize_type(&context.typing_session())
441                    .const_int(n, false),
442                arg => bail!("Unexpected type arg for LoadNat: {}", arg),
443            };
444            args.outputs.finish(context.builder(), vec![v.into()])
445        }
446    })
447    .extension_op(prelude::PRELUDE_ID, prelude::BARRIER_OP_ID, {
448        let pcg = pcg.clone();
449        move |context, args| pcg.emit_barrier(context, args)
450    })
451}
452
453#[cfg(test)]
454mod test {
455    use hugr_core::builder::{Dataflow, DataflowHugr};
456    use hugr_core::extension::PRELUDE;
457    use hugr_core::extension::prelude::{EXIT_OP_ID, MAKE_ERROR_OP_ID, Noop};
458    use hugr_core::types::{Term, Type};
459    use hugr_core::{Hugr, type_row};
460    use prelude::{PANIC_OP_ID, PRINT_OP_ID, bool_t, qb_t, usize_t};
461    use rstest::{fixture, rstest};
462
463    use crate::check_emission;
464    use crate::custom::CodegenExtsBuilder;
465    use crate::emit::test::SimpleHugrConfig;
466    use crate::test::{TestContext, exec_ctx, llvm_ctx};
467    use crate::types::HugrType;
468
469    use super::*;
470
471    #[derive(Clone)]
472    struct TestPreludeCodegen;
473    impl PreludeCodegen for TestPreludeCodegen {
474        fn usize_type<'c>(&self, session: &TypingSession<'c, '_>) -> IntType<'c> {
475            session.iw_context().i32_type()
476        }
477
478        fn qubit_type<'c>(&self, session: &TypingSession<'c, '_>) -> impl BasicType<'c> {
479            session.iw_context().f64_type()
480        }
481    }
482
483    #[rstest]
484    fn prelude_extension_types(llvm_ctx: TestContext) {
485        let iw_context = llvm_ctx.iw_context();
486        let type_converter = CodegenExtsBuilder::<Hugr>::default()
487            .add_prelude_extensions(TestPreludeCodegen)
488            .finish()
489            .type_converter;
490        let session = type_converter.session(iw_context);
491
492        assert_eq!(
493            iw_context.i32_type().as_basic_type_enum(),
494            session.llvm_type(&usize_t()).unwrap()
495        );
496        assert_eq!(
497            iw_context.f64_type().as_basic_type_enum(),
498            session.llvm_type(&qb_t()).unwrap()
499        );
500    }
501
502    #[rstest]
503    fn prelude_extension_types_in_test_context(mut llvm_ctx: TestContext) {
504        llvm_ctx.add_extensions(|x| x.add_prelude_extensions(TestPreludeCodegen));
505        let tc = llvm_ctx.get_typing_session();
506        assert_eq!(
507            llvm_ctx.iw_context().i32_type().as_basic_type_enum(),
508            tc.llvm_type(&usize_t()).unwrap()
509        );
510        assert_eq!(
511            llvm_ctx.iw_context().f64_type().as_basic_type_enum(),
512            tc.llvm_type(&qb_t()).unwrap()
513        );
514    }
515
516    #[rstest::fixture]
517    fn prelude_llvm_ctx(mut llvm_ctx: TestContext) -> TestContext {
518        llvm_ctx.add_extensions(CodegenExtsBuilder::add_default_prelude_extensions);
519        llvm_ctx
520    }
521
522    #[rstest]
523    fn prelude_const_usize(prelude_llvm_ctx: TestContext) {
524        let hugr = SimpleHugrConfig::new()
525            .with_outs(usize_t())
526            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
527            .finish(|mut builder| {
528                let k = builder.add_load_value(ConstUsize::new(17));
529                builder.finish_hugr_with_outputs([k]).unwrap()
530            });
531        check_emission!(hugr, prelude_llvm_ctx);
532    }
533
534    #[rstest]
535    fn prelude_const_external_symbol(prelude_llvm_ctx: TestContext) {
536        let konst1 = ConstExternalSymbol::new("sym1", usize_t(), true);
537        let konst2 = ConstExternalSymbol::new(
538            "sym2",
539            HugrType::new_sum([
540                vec![usize_t(), HugrType::new_unit_sum(3)].into(),
541                type_row![],
542            ]),
543            false,
544        );
545
546        let hugr = SimpleHugrConfig::new()
547            .with_outs(vec![konst1.get_type(), konst2.get_type()])
548            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
549            .finish(|mut builder| {
550                let k1 = builder.add_load_value(konst1);
551                let k2 = builder.add_load_value(konst2);
552                builder.finish_hugr_with_outputs([k1, k2]).unwrap()
553            });
554        check_emission!(hugr, prelude_llvm_ctx);
555    }
556
557    #[rstest]
558    fn prelude_noop(prelude_llvm_ctx: TestContext) {
559        let hugr = SimpleHugrConfig::new()
560            .with_ins(usize_t())
561            .with_outs(usize_t())
562            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
563            .finish(|mut builder| {
564                let in_wires = builder.input_wires();
565                let r = builder
566                    .add_dataflow_op(Noop::new(usize_t()), in_wires)
567                    .unwrap()
568                    .outputs();
569                builder.finish_hugr_with_outputs(r).unwrap()
570            });
571        check_emission!(hugr, prelude_llvm_ctx);
572    }
573
574    #[rstest]
575    fn prelude_make_tuple(prelude_llvm_ctx: TestContext) {
576        let hugr = SimpleHugrConfig::new()
577            .with_ins(vec![bool_t(), bool_t()])
578            .with_outs(Type::new_tuple(vec![bool_t(), bool_t()]))
579            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
580            .finish(|mut builder| {
581                let in_wires = builder.input_wires();
582                let r = builder.make_tuple(in_wires).unwrap();
583                builder.finish_hugr_with_outputs([r]).unwrap()
584            });
585        check_emission!(hugr, prelude_llvm_ctx);
586    }
587
588    #[rstest]
589    fn prelude_unpack_tuple(prelude_llvm_ctx: TestContext) {
590        let hugr = SimpleHugrConfig::new()
591            .with_ins(Type::new_tuple(vec![bool_t(), bool_t()]))
592            .with_outs(vec![bool_t(), bool_t()])
593            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
594            .finish(|mut builder| {
595                let unpack = builder
596                    .add_dataflow_op(
597                        UnpackTuple::new(vec![bool_t(), bool_t()].into()),
598                        builder.input_wires(),
599                    )
600                    .unwrap();
601                builder.finish_hugr_with_outputs(unpack.outputs()).unwrap()
602            });
603        check_emission!(hugr, prelude_llvm_ctx);
604    }
605
606    #[rstest]
607    fn prelude_panic(prelude_llvm_ctx: TestContext) {
608        let error_val = ConstError::new(42, "PANIC");
609        let type_arg_q: Term = qb_t().into();
610        let type_arg_2q = Term::new_list([type_arg_q.clone(), type_arg_q]);
611        let panic_op = PRELUDE
612            .instantiate_extension_op(&PANIC_OP_ID, [type_arg_2q.clone(), type_arg_2q.clone()])
613            .unwrap();
614
615        let hugr = SimpleHugrConfig::new()
616            .with_ins(vec![qb_t(), qb_t()])
617            .with_outs(vec![qb_t(), qb_t()])
618            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
619            .finish(|mut builder| {
620                let [q0, q1] = builder.input_wires_arr();
621                let err = builder.add_load_value(error_val);
622                let [q0, q1] = builder
623                    .add_dataflow_op(panic_op, [err, q0, q1])
624                    .unwrap()
625                    .outputs_arr();
626                builder.finish_hugr_with_outputs([q0, q1]).unwrap()
627            });
628
629        check_emission!(hugr, prelude_llvm_ctx);
630    }
631
632    #[rstest]
633    fn prelude_exit(prelude_llvm_ctx: TestContext) {
634        let error_val = ConstError::new(42, "EXIT");
635        let type_arg_q: Term = qb_t().into();
636        let type_arg_2q = Term::new_list([type_arg_q.clone(), type_arg_q]);
637        let exit_op = PRELUDE
638            .instantiate_extension_op(&EXIT_OP_ID, [type_arg_2q.clone(), type_arg_2q.clone()])
639            .unwrap();
640
641        let hugr = SimpleHugrConfig::new()
642            .with_ins(vec![qb_t(), qb_t()])
643            .with_outs(vec![qb_t(), qb_t()])
644            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
645            .finish(|mut builder| {
646                let [q0, q1] = builder.input_wires_arr();
647                let err = builder.add_load_value(error_val);
648                let [q0, q1] = builder
649                    .add_dataflow_op(exit_op, [err, q0, q1])
650                    .unwrap()
651                    .outputs_arr();
652                builder.finish_hugr_with_outputs([q0, q1]).unwrap()
653            });
654
655        check_emission!(hugr, prelude_llvm_ctx);
656    }
657
658    #[rstest]
659    fn prelude_print(prelude_llvm_ctx: TestContext) {
660        let greeting: ConstString = ConstString::new("Hello, world!".into());
661        let print_op = PRELUDE.instantiate_extension_op(&PRINT_OP_ID, []).unwrap();
662
663        let hugr = SimpleHugrConfig::new()
664            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
665            .finish(|mut builder| {
666                let greeting_out = builder.add_load_value(greeting);
667                builder.add_dataflow_op(print_op, [greeting_out]).unwrap();
668                builder.finish_hugr_with_outputs([]).unwrap()
669            });
670
671        check_emission!(hugr, prelude_llvm_ctx);
672    }
673
674    #[rstest]
675    fn prelude_make_error(prelude_llvm_ctx: TestContext) {
676        let sig: ConstUsize = ConstUsize::new(100);
677        let msg: ConstString = ConstString::new("Error!".into());
678
679        let make_error_op = PRELUDE
680            .instantiate_extension_op(&MAKE_ERROR_OP_ID, [])
681            .unwrap();
682
683        let hugr = SimpleHugrConfig::new()
684            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
685            .with_outs(error_type())
686            .finish(|mut builder| {
687                let sig_out = builder.add_load_value(sig);
688                let msg_out = builder.add_load_value(msg);
689                let [err] = builder
690                    .add_dataflow_op(make_error_op, [sig_out, msg_out])
691                    .unwrap()
692                    .outputs_arr();
693                builder.finish_hugr_with_outputs([err]).unwrap()
694            });
695
696        check_emission!(hugr, prelude_llvm_ctx);
697    }
698
699    #[rstest]
700    fn prelude_make_error_and_panic(prelude_llvm_ctx: TestContext) {
701        let sig: ConstUsize = ConstUsize::new(100);
702        let msg: ConstString = ConstString::new("Error!".into());
703
704        let make_error_op = PRELUDE
705            .instantiate_extension_op(&MAKE_ERROR_OP_ID, [])
706            .unwrap();
707
708        let panic_op = PRELUDE
709            .instantiate_extension_op(&PANIC_OP_ID, [Term::new_list([]), Term::new_list([])])
710            .unwrap();
711
712        let hugr = SimpleHugrConfig::new()
713            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
714            .finish(|mut builder| {
715                let sig_out = builder.add_load_value(sig);
716                let msg_out = builder.add_load_value(msg);
717                let [err] = builder
718                    .add_dataflow_op(make_error_op, [sig_out, msg_out])
719                    .unwrap()
720                    .outputs_arr();
721                builder.add_dataflow_op(panic_op, [err]).unwrap();
722                builder.finish_hugr_with_outputs([]).unwrap()
723            });
724
725        check_emission!(hugr, prelude_llvm_ctx);
726    }
727
728    #[rstest]
729    fn prelude_load_nat(prelude_llvm_ctx: TestContext) {
730        let hugr = SimpleHugrConfig::new()
731            .with_outs(usize_t())
732            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
733            .finish(|mut builder| {
734                let v = builder
735                    .add_dataflow_op(LoadNat::new(42u64.into()), vec![])
736                    .unwrap()
737                    .out_wire(0);
738                builder.finish_hugr_with_outputs([v]).unwrap()
739            });
740        check_emission!(hugr, prelude_llvm_ctx);
741    }
742
743    #[fixture]
744    fn barrier_hugr() -> Hugr {
745        SimpleHugrConfig::new()
746            .with_outs(vec![usize_t()])
747            .with_extensions(prelude::PRELUDE_REGISTRY.to_owned())
748            .finish(|mut builder| {
749                let i = builder.add_load_value(ConstUsize::new(42));
750                let [w1, _w2] = builder.add_barrier([i, i]).unwrap().outputs_arr();
751                builder.finish_hugr_with_outputs([w1]).unwrap()
752            })
753    }
754
755    #[rstest]
756    fn prelude_barrier(prelude_llvm_ctx: TestContext, barrier_hugr: Hugr) {
757        check_emission!(barrier_hugr, prelude_llvm_ctx);
758    }
759    #[rstest]
760    fn prelude_barrier_exec(mut exec_ctx: TestContext, barrier_hugr: Hugr) {
761        exec_ctx.add_extensions(|cem| add_prelude_extensions(cem, TestPreludeCodegen));
762        assert_eq!(exec_ctx.exec_hugr_u64(barrier_hugr, "main"), 42);
763    }
764}
hugr_llvm/extension/prelude.rs

hugr_llvm/extension/
prelude.rs
