hugr_core/extension/
const_fold.rs

1use std::fmt::Formatter;
2
3use std::fmt::Debug;
4
5use crate::ops::Value;
6use crate::types::TypeArg;
7
8use crate::IncomingPort;
9use crate::OutgoingPort;
10
11use crate::ops;
12
13/// Output of constant folding an operation, None indicates folding was either
14/// not possible or unsuccessful. An empty vector indicates folding was
15/// successful and no values are output.
16pub type ConstFoldResult = Option<Vec<(OutgoingPort, ops::Value)>>;
17
18/// Tag some output constants with [`OutgoingPort`] inferred from the ordering.
19pub fn fold_out_row(consts: impl IntoIterator<Item = Value>) -> ConstFoldResult {
20    let vec = consts
21        .into_iter()
22        .enumerate()
23        .map(|(i, c)| (i.into(), c))
24        .collect();
25    Some(vec)
26}
27
28/// Trait implemented by extension operations that can perform constant folding.
29pub trait ConstFold: Send + Sync {
30    /// Given type arguments `type_args` and
31    /// [`crate::ops::Const`] values for inputs at [`crate::IncomingPort`]s,
32    /// try to evaluate the operation.
33    fn fold(
34        &self,
35        type_args: &[TypeArg],
36        consts: &[(crate::IncomingPort, crate::ops::Value)],
37    ) -> ConstFoldResult;
38}
39
40impl Debug for Box<dyn ConstFold> {
41    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
42        write!(f, "<custom constant folding>")
43    }
44}
45
46/// Blanket implementation for functions that only require the constants to
47/// evaluate - type arguments are not relevant.
48impl<T> ConstFold for T
49where
50    T: Fn(&[(crate::IncomingPort, crate::ops::Value)]) -> ConstFoldResult + Send + Sync,
51{
52    fn fold(
53        &self,
54        _type_args: &[TypeArg],
55        consts: &[(crate::IncomingPort, crate::ops::Value)],
56    ) -> ConstFoldResult {
57        self(consts)
58    }
59}
60
61type FoldFn = dyn Fn(&[TypeArg], &[(IncomingPort, Value)]) -> ConstFoldResult + Send + Sync;
62
63/// Type holding a boxed const-folding function.
64pub struct Folder {
65    /// Const-folding function.
66    pub folder: Box<FoldFn>,
67}
68
69impl ConstFold for Folder {
70    fn fold(&self, type_args: &[TypeArg], consts: &[(IncomingPort, Value)]) -> ConstFoldResult {
71        (self.folder)(type_args, consts)
72    }
73}
hugr_core/extension/const_fold.rs

hugr_core/extension/
const_fold.rs
