hugr_llvm/custom/
load_constant.rs1use 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
11pub 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#[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 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 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}