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

hugr_llvm/custom/
load_constant.rs
