hugr_llvm/custom/
load_constant.rs

1//! Provides the implementation for a collection of [`CustomConst`] callbacks.
2use std::{any::TypeId, collections::BTreeMap};
3
4use hugr_core::{HugrView, Node, ops::constant::CustomConst};
5use inkwell::values::BasicValueEnum;
6
7use anyhow::{Result, anyhow, bail, ensure};
8
9use crate::emit::EmitFuncContext;
10
11/// A helper trait for describing the callback used for emitting [`CustomConst`]s,
12/// and for hanging documentation. We have the appropriate `Fn` as a supertrait,
13/// and there is a blanket impl for that `Fn`. We do not intend users to impl
14/// this trait.
15///
16/// `LoadConstantFn` callbacks for `CC`, which must impl [`CustomConst`], should
17/// materialise an appropriate [`BasicValueEnum`]. The type of this value must
18/// match the result of [`EmitFuncContext::llvm_type`] on [`CustomConst::get_type`].
19///
20/// Callbacks may hold references with lifetimes older than `'a`.
21pub trait LoadConstantFn<'a, H: ?Sized, CC: CustomConst + ?Sized>:
22    for<'c> Fn(&mut EmitFuncContext<'c, 'a, H>, &CC) -> Result<BasicValueEnum<'c>> + 'a
23{
24}
25
26impl<
27    'a,
28    H: ?Sized,
29    CC: ?Sized + CustomConst,
30    F: 'a + ?Sized + for<'c> Fn(&mut EmitFuncContext<'c, 'a, H>, &CC) -> Result<BasicValueEnum<'c>>,
31> LoadConstantFn<'a, H, CC> for F
32{
33}
34
35/// A collection of [`LoadConstantFn`] callbacks registered for various impls of [`CustomConst`].
36/// The callbacks are keyed by the [`TypeId`] of those impls.
37#[derive(Default)]
38pub struct LoadConstantsMap<'a, H>(
39    BTreeMap<TypeId, Box<dyn LoadConstantFn<'a, H, dyn CustomConst>>>,
40);
41
42impl<'a, H: HugrView<Node = Node>> LoadConstantsMap<'a, H> {
43    /// Register a callback to emit a `CC` value.
44    ///
45    /// If a callback is already registered for that type, we will replace it.
46    pub fn custom_const<CC: CustomConst>(&mut self, handler: impl LoadConstantFn<'a, H, CC>) {
47        self.0.insert(
48            TypeId::of::<CC>(),
49            Box::new(move |context, konst: &dyn CustomConst| {
50                let cc = konst.downcast_ref::<CC>().ok_or(anyhow!(
51                    "impossible! Failed to downcast in LoadConstantsMap::custom_const"
52                ))?;
53                handler(context, cc)
54            }),
55        );
56    }
57
58    /// Emit instructions to materialise `konst` by delegating to the
59    /// appropriate inner callbacks.
60    pub fn emit_load_constant<'c>(
61        &self,
62        context: &mut EmitFuncContext<'c, 'a, H>,
63        konst: &dyn CustomConst,
64    ) -> Result<BasicValueEnum<'c>> {
65        let type_id = konst.type_id();
66        let Some(handler) = self.0.get(&type_id) else {
67            bail!(
68                "No extension could load constant name: {} type_id: {type_id:?}",
69                konst.name()
70            )
71        };
72        let r = handler(context, konst)?;
73        let r_type = r.get_type();
74        let konst_type = context.llvm_type(&konst.get_type())?;
75        ensure!(
76            r_type == konst_type,
77            "CustomConst handler returned a value of the wrong type. Expected: {konst_type} Actual: {r_type}"
78        );
79        Ok(r)
80    }
81}
hugr_llvm/custom/load_constant.rs

hugr_llvm/custom/
load_constant.rs
